use axum::{ body::Body, extract::{Multipart, Path, State}, http::{header, HeaderValue, StatusCode}, response::{Html, IntoResponse, Response}, }; use std::path::PathBuf; use tokio::io::AsyncWriteExt; use tokio_util::io::ReaderStream; pub const MAX_TOTAL_STORAGE_BYTES: i64 = 10 * 1024 * 1024 * 1024; pub const MAX_FILE_SIZE_BYTES: i64 = 10 * 1024 * 1024 * 1024; use crate::web_app_state::{AppState, AppStateInner, AuthenticatedUserId}; pub struct UploadSuccess { pub(crate) id: i32, pub(crate) filename: String, pub(crate) size: u64, } pub fn render_upload_page(state: &AppStateInner, error: Option, success: Option) -> Html { let mut ctx = tera::Context::new(); if let Some(error) = error { ctx.insert("error", &error); } if let Some(success) = success { ctx.insert("success_id", &success.id); ctx.insert("success_filename", &success.filename); ctx.insert("success_size", &success.size); } let body = state.tera .render("upload.html", &ctx) .expect("render upload"); Html(body) } pub async fn upload_get( _: (), _state: AppState, _user: AuthenticatedUserId,) -> Html { render_upload_page(&_state, None, None) } pub async fn upload_post( (headers, mut multipart): (axum::http::HeaderMap, Multipart), _state: AppState, _user: AuthenticatedUserId, ) -> Html { let content_length: u64 = match headers .get(header::CONTENT_LENGTH) .and_then(|value| value.to_str().ok()) .and_then(|value| value.parse::().ok()) { Some(value) => value, _ => return render_upload_page(&_state, Some("Missing or invalid Content-Length".to_string()), None, ) }; if (content_length > MAX_FILE_SIZE_BYTES as u64){ return render_upload_page(&_state, Some("File size is too big".to_string()), None); } let old_total_size: i64 = match _state .db.query_one( "SELECT COALESCE(SUM(size), 0)::BIGINT FROM public.stored_file", &[], ).await { Ok(row) => row.get::<_, i64>(0), Err(err) => { eprintln!("Database error: {err}"); return render_upload_page(&_state, Some("Server error".to_string()), None, ); } }; assert!(0 <= old_total_size); assert!(old_total_size <= MAX_TOTAL_STORAGE_BYTES); let total_with_new: u64 = old_total_size as u64 + content_length; if total_with_new > MAX_TOTAL_STORAGE_BYTES as u64 { return render_upload_page(&_state, Some("Storage limit exceeded (10 GiB)".to_string()), None, ); } let mut field = loop { match multipart.next_field().await { Ok(Some(field)) => { if field.name() == Some("file") { break field; } } Ok(None) => {} Err(err) => return render_upload_page(&_state, Some(format!("Upload error: {err}")), None, ) } }; let filename = field .file_name() .map(|name| name.to_string()) .unwrap_or_else(|| "upload".to_string()); let row = match _state.db.query_one( "INSERT INTO public.stored_file (filename, size) VALUES ($1, 0) RETURNING id", &[&filename], ).await { Ok(row) => row, Err(err) => { return render_upload_page(&_state, Some(format!("Database error: {err}")), None, ); } }; let id: i32 = row.get::<_, i32>(0); let path = PathBuf::from(&_state.config.file_storage).join(id.to_string()); let mut file = match tokio::fs::File::create(&path).await { Ok(file) => file, Err(err) => { let _ = _state.db.execute( "DELETE FROM public.stored_file WHERE id = $1", &[&id] ).await; eprintln!("Cannot create file: {err}"); return render_upload_page(&_state, Some("Filesystem error".to_string()), None, ); } }; if let Err(err) = file.set_len(content_length).await { let _ = tokio::fs::remove_file(&path).await; let _ = _state.db.execute( "DELETE FROM public.stored_file WHERE id = $1", &[&id] ).await; eprintln!("Cannot preallocate file: {err}"); return render_upload_page(&_state, Some("Filesystem error".to_string()), None, ); } let mut written: u64 = 0; loop { match field.chunk().await { Ok(Some(chunk)) => { let chunk_len = chunk.len() as u64; if written + chunk_len > content_length { let _ = tokio::fs::remove_file(&path).await; let _ = _state.db.execute( "DELETE FROM public.stored_file WHERE id = $1", &[&id] ).await; return render_upload_page(&_state, Some("Your browser is bad".to_string()), None, ); } written += chunk_len; { let percent = 100. * (written as f32) / (content_length as f32); println!("DEBUG: uploaded {percent}, cur chunk is {chunk_len}") } if let Err(err) = file.write_all(&chunk).await { let _ = tokio::fs::remove_file(&path).await; let _ = _state.db.execute( "DELETE FROM public.stored_file WHERE id = $1", &[&id] ).await; eprintln!("Can't write chunk to file: {err}"); return render_upload_page(&_state, Some("Filesystem error".to_string()), None, ); } } Ok(None) => break, Err(err) => { let _ = tokio::fs::remove_file(&path).await; let _ = _state.db.execute( "DELETE FROM public.stored_file WHERE id = $1", &[&id] ).await; return render_upload_page(&_state, Some(format!("Upload error: {err}")), None, ); } } } if let Err(err) = _state.db.execute( "UPDATE public.stored_file SET size = $1 WHERE id = $2", &[&(written as i64), &id], ).await { let _ = tokio::fs::remove_file(&path).await; let _ = _state.db.execute( "DELETE FROM public.stored_file WHERE id = $1", &[&id] ).await; return render_upload_page(&_state, Some("Server error".to_string()), None, ); } render_upload_page(&_state, None, Some(UploadSuccess { id, filename, size: written, }) ) } pub async fn get_file( Path(id): Path, state: AppState, _user: AuthenticatedUserId, ) -> Response { let row = match state .db .query_opt( "SELECT filename, size FROM public.stored_file WHERE id = $1", &[&id], ) .await { Ok(Some(row)) => row, Ok(None) => return StatusCode::NOT_FOUND.into_response(), Err(err) => { eprintln!("file lookup error: {err}"); return StatusCode::INTERNAL_SERVER_ERROR.into_response(); } }; let filename: String = row.get(0); let size: i64 = row.get(1); let path = PathBuf::from(&state.config.file_storage).join(id.to_string()); let file = match tokio::fs::File::open(&path).await { Ok(file) => file, Err(err) => { eprintln!("file open error for {path:?}: {err}"); return StatusCode::NOT_FOUND.into_response(); } }; let stream = ReaderStream::new(file); let body = Body::from_stream(stream); let mut response = Response::new(body); if let Ok(value) = HeaderValue::from_str(&size.to_string()) { response.headers_mut().insert(header::CONTENT_LENGTH, value); } response .headers_mut() .insert(header::CONTENT_TYPE, HeaderValue::from_static("application/octet-stream")); response.headers_mut().insert( header::CONTENT_DISPOSITION, HeaderValue::from_str(&format!("attachment; filename=\"{}\"", filename)).unwrap_or_else( |_| HeaderValue::from_static("attachment"), ), ); response }