fix: Fix attachment upload size
This commit is contained in:
@@ -61,10 +61,9 @@ pub async fn upload_attachment(
|
|||||||
axum::Extension(_current_user): axum::Extension<User>,
|
axum::Extension(_current_user): axum::Extension<User>,
|
||||||
mut multipart: Multipart,
|
mut multipart: Multipart,
|
||||||
) -> Result<Response, AppError> {
|
) -> Result<Response, AppError> {
|
||||||
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 filename = field.file_name().unwrap_or("unknown").to_string();
|
||||||
let content_type = field.content_type().unwrap_or("application/octet-stream").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()
|
let id: String = rand::rng()
|
||||||
.sample_iter(&Alphanumeric)
|
.sample_iter(&Alphanumeric)
|
||||||
@@ -72,19 +71,44 @@ pub async fn upload_attachment(
|
|||||||
.map(char::from)
|
.map(char::from)
|
||||||
.collect();
|
.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);
|
let file_path = PathBuf::from("attachments").join(&id);
|
||||||
|
|
||||||
fs::write(&file_path, &processed_data).await
|
// 1. Stream everything to disk directly to avoid OOM for 5GB+ files
|
||||||
.context("Failed to write attachment to disk")?;
|
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 (?, ?, ?, ?, ?)")
|
sqlx::query("INSERT INTO attachments (id, filename, content_type, size, created_at) VALUES (?, ?, ?, ?, ?)")
|
||||||
.bind(&id)
|
.bind(&id)
|
||||||
@@ -166,19 +190,3 @@ async fn process_image(data: &[u8], content_type: &str) -> Result<Vec<u8>> {
|
|||||||
img.write_to(&mut out, format).context("Failed to re-encode image (stripping metadata)")?;
|
img.write_to(&mut out, format).context("Failed to re-encode image (stripping metadata)")?;
|
||||||
Ok(out.into_inner())
|
Ok(out.into_inner())
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn process_video(data: &[u8]) -> Result<Vec<u8>> {
|
|
||||||
// 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())
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -120,6 +120,7 @@ async fn main() {
|
|||||||
.route("/__attachments/{id}", get(handlers::attachments::serve_attachment))
|
.route("/__attachments/{id}", get(handlers::attachments::serve_attachment))
|
||||||
.nest("/__dungeon", handlers::router(&app_state)) // I'll create a single `handlers::router`
|
.nest("/__dungeon", handlers::router(&app_state)) // I'll create a single `handlers::router`
|
||||||
.with_state(app_state.clone())
|
.with_state(app_state.clone())
|
||||||
|
.layer(axum::extract::DefaultBodyLimit::disable())
|
||||||
.layer(tower_http::trace::TraceLayer::new_for_http());
|
.layer(tower_http::trace::TraceLayer::new_for_http());
|
||||||
|
|
||||||
let listener = TcpListener::bind(&addr).await.unwrap();
|
let listener = TcpListener::bind(&addr).await.unwrap();
|
||||||
|
|||||||
Reference in New Issue
Block a user