diff --git a/src/handlers/attachments.rs b/src/handlers/attachments.rs index 268c454..1102ae8 100644 --- a/src/handlers/attachments.rs +++ b/src/handlers/attachments.rs @@ -61,10 +61,9 @@ pub async fn upload_attachment( axum::Extension(_current_user): axum::Extension, mut multipart: Multipart, ) -> Result { - while let Some(field) = multipart.next_field().await.map_err(|e| AppError::from(anyhow::anyhow!(e)))? { + while let Some(mut field) = multipart.next_field().await.map_err(|e| AppError::from(anyhow::anyhow!(e)))? { let filename = field.file_name().unwrap_or("unknown").to_string(); let content_type = field.content_type().unwrap_or("application/octet-stream").to_string(); - let data = field.bytes().await.map_err(|e| AppError::from(anyhow::anyhow!(e)))?; let id: String = rand::rng() .sample_iter(&Alphanumeric) @@ -72,19 +71,44 @@ pub async fn upload_attachment( .map(char::from) .collect(); - let processed_data = if content_type.starts_with("image/") { - process_image(&data, &content_type).await.unwrap_or(data.to_vec()) - } else if content_type.starts_with("video/") { - process_video(&data).await.unwrap_or(data.to_vec()) - } else { - data.to_vec() - }; - - let size = processed_data.len() as i64; let file_path = PathBuf::from("attachments").join(&id); - fs::write(&file_path, &processed_data).await - .context("Failed to write attachment to disk")?; + // 1. Stream everything to disk directly to avoid OOM for 5GB+ files + use tokio::io::AsyncWriteExt; + let mut file = match tokio::fs::File::create(&file_path).await { + Ok(f) => f, + Err(e) => return Err(AppError::from(anyhow::anyhow!("Failed to create file: {}", e))), + }; + + let mut size = 0i64; + while let Some(chunk) = field.chunk().await.map_err(|e| AppError::from(anyhow::anyhow!(e)))? { + file.write_all(&chunk).await.map_err(|e| AppError::from(anyhow::anyhow!("Failed to write chunk: {}", e)))?; + size += chunk.len() as i64; + } + + file.flush().await.map_err(|e| AppError::from(anyhow::anyhow!("Failed to flush: {}", e)))?; + drop(file); + + // 2. Process image/video in place if needed + if content_type.starts_with("image/") { + let data = fs::read(&file_path).await.context("Failed to read image")?; + if let Ok(processed_data) = process_image(&data, &content_type).await { + fs::write(&file_path, &processed_data).await.context("Failed to overwrite image")?; + size = processed_data.len() as i64; + } + } else if content_type.starts_with("video/") { + let p_buf = file_path.clone(); + let _ = tokio::task::spawn_blocking(move || { + if let Ok(probed) = lofty::probe::Probe::open(&p_buf) { + if let Ok(mut video_file) = probed.read() { + video_file.clear(); // Clear all tags + let _ = video_file.save_to_path(&p_buf, lofty::config::WriteOptions::default()); + } + } + }).await; + + size = fs::metadata(&file_path).await.context("Failed to get size")?.len() as i64; + } sqlx::query("INSERT INTO attachments (id, filename, content_type, size, created_at) VALUES (?, ?, ?, ?, ?)") .bind(&id) @@ -166,19 +190,3 @@ async fn process_image(data: &[u8], content_type: &str) -> Result> { img.write_to(&mut out, format).context("Failed to re-encode image (stripping metadata)")?; Ok(out.into_inner()) } - -async fn process_video(data: &[u8]) -> Result> { - // Lofty can remove tags from various formats - let mut cursor = Cursor::new(data.to_vec()); - if let Ok(probed) = lofty::probe::Probe::new(&mut cursor).guess_file_type() { - if let Ok(mut file) = probed.read() { - file.clear(); - // We need to save it back. Lofty modify the file in place if it's a file but here we have a cursor. - let mut out = Cursor::new(Vec::new()); - if file.save_to(&mut out, lofty::config::WriteOptions::default()).is_ok() { - return Ok(out.into_inner()); - } - } - } - Ok(data.to_vec()) -} diff --git a/src/main.rs b/src/main.rs index eb035da..123ee2d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -120,6 +120,7 @@ async fn main() { .route("/__attachments/{id}", get(handlers::attachments::serve_attachment)) .nest("/__dungeon", handlers::router(&app_state)) // I'll create a single `handlers::router` .with_state(app_state.clone()) + .layer(axum::extract::DefaultBodyLimit::disable()) .layer(tower_http::trace::TraceLayer::new_for_http()); let listener = TcpListener::bind(&addr).await.unwrap();