use std::sync::Arc; use axum::extract::Multipart; use axum::extract::Path as AxumPath; use axum::response::{IntoResponse, Response}; use axum::{Extension, Json}; use bytes::Bytes; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use crate::auth::session::SessionState; use crate::config::AppConfig; use crate::db; use crate::error::AppError; use crate::imap::client::{ImapClient, ImapCredentials}; use crate::storage::{self, BlobStorage, attachment_key}; #[derive(Debug, Serialize)] struct UploadResponse { id: String, filename: String, content_type: String, size: i64, } #[derive(Debug, Serialize)] struct DeleteResponse { status: String, } #[derive(Debug, Deserialize)] pub(crate) struct AttachMessageSourceRequest { folder: String, uid: u32, } pub async fn upload_attachment( Extension(session): Extension, Extension(config): Extension>, Extension(storage): Extension>, AxumPath(draft_id): AxumPath, mut multipart: Multipart, ) -> Result { db::drafts::validate_draft_id(&draft_id).map_err(AppError::BadRequest)?; let conn = db::pool::open_user_db(&config.data_dir, &session.user_hash) .map_err(|e| AppError::InternalError(format!("Failed to open database: {e}")))?; let existing = db::drafts::get_draft(&conn, &draft_id).map_err(AppError::InternalError)?; if existing.is_none() { db::drafts::upsert_draft( &conn, &draft_id, "", "", "", "", "", None, None, None, false, "normal", "html", "none", false, false, false, ) .map_err(AppError::InternalError)?; } let mut uploaded = Vec::new(); let max_attachment_size = config.max_attachment_size_bytes(); while let Some(mut field) = multipart .next_field() .await .map_err(|e| AppError::BadRequest(format!("Invalid multipart data: {e}")))? { let filename = field.file_name().unwrap_or("unnamed").to_string(); let content_type = field .content_type() .unwrap_or("application/octet-stream") .to_string(); let att_id = uuid::Uuid::new_v4().to_string(); let key = attachment_key(&session.user_hash, &draft_id, &att_id) .map_err(AppError::InternalError)?; let mut writer = storage .begin_upload(&key) .await .map_err(|e| AppError::ServiceUnavailable(format!("Storage unavailable: {e}")))?; let mut size: usize = 0; let mut checksum = Sha256::new(); loop { let chunk = match field.chunk().await { Ok(chunk) => chunk, Err(e) => { if writer.abort().await.is_err() && let Err(cleanup_error) = storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key) .await { tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue aborted attachment cleanup"); } return Err(AppError::BadRequest(format!( "Failed to read field data: {e}" ))); } }; let Some(chunk) = chunk else { break }; size = size.saturating_add(chunk.len()); if size > max_attachment_size { if writer.abort().await.is_err() && let Err(cleanup_error) = storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await { tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue oversized attachment cleanup"); } return Err(AppError::BadRequest(format!( "File '{}' exceeds maximum size of {} MiB", filename, config.max_attachment_size_mb ))); } checksum.update(&chunk); if let Err(e) = writer.write(chunk).await { if writer.abort().await.is_err() && let Err(cleanup_error) = storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await { tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue interrupted attachment cleanup"); } return Err(AppError::ServiceUnavailable(format!( "Failed to upload attachment: {e}" ))); } } let stat = match writer.finish().await { Ok(stat) => stat, Err(e) => { if let Err(cleanup_error) = storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await { tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue failed-finalize attachment cleanup"); } return Err(AppError::ServiceUnavailable(format!( "Failed to finalize attachment upload: {e}" ))); } }; if stat.size != size as u64 { if let Err(cleanup_error) = storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await { tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue invalid attachment cleanup"); } return Err(AppError::InternalError(format!( "Attachment size verification failed for '{filename}'" ))); } let checksum = format!("{:x}", checksum.finalize()); let size = size as i64; let attachment = db::drafts::NewDraftAttachment { id: &att_id, draft_id: &draft_id, filename: &filename, content_type: &content_type, size, storage_key: &key, checksum_sha256: Some(&checksum), }; if let Err(e) = db::drafts::add_draft_attachment(&conn, &attachment) { if let Err(cleanup_error) = storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await { tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue unattached blob cleanup"); } return Err(AppError::InternalError(e)); } uploaded.push(UploadResponse { id: att_id, filename, content_type, size, }); } if uploaded.is_empty() { return Err(AppError::BadRequest("No files uploaded".to_string())); } Ok(Json(serde_json::json!({ "attachments": uploaded, })) .into_response()) } pub async fn attach_message_source( Extension(session): Extension, Extension(config): Extension>, Extension(storage): Extension>, Extension(imap_client): Extension>, AxumPath(draft_id): AxumPath, Json(req): Json, ) -> Result { db::drafts::validate_draft_id(&draft_id).map_err(AppError::BadRequest)?; if req.folder.trim().is_empty() || req.uid == 0 { return Err(AppError::BadRequest( "Message folder and UID are required".to_string(), )); } let creds = ImapCredentials { host: session.imap_host.clone(), port: session.imap_port, tls: session.imap_tls, tls_verify: session.imap_tls_verify, email: session.email.clone(), password: session.password.clone(), }; let raw = imap_client .fetch_raw_message(&creds, &req.folder, req.uid) .await .map_err(|e| AppError::ServiceUnavailable(format!("IMAP error: {e}")))?; if raw.is_empty() { return Err(AppError::BadRequest("Message source is empty".to_string())); } if raw.len() > config.max_attachment_size_bytes() { return Err(AppError::BadRequest(format!( "Message source exceeds maximum attachment size of {} MiB", config.max_attachment_size_mb ))); } let conn = db::pool::open_user_db(&config.data_dir, &session.user_hash) .map_err(|e| AppError::InternalError(format!("Failed to open database: {e}")))?; if db::drafts::get_draft(&conn, &draft_id) .map_err(AppError::InternalError)? .is_none() { db::drafts::upsert_draft( &conn, &draft_id, "", "", "", "", "", None, None, None, false, "normal", "html", "none", false, false, false, ) .map_err(AppError::InternalError)?; } let filename = format!("message-{}.eml", req.uid); let content_type = "message/rfc822".to_string(); let att_id = uuid::Uuid::new_v4().to_string(); let key = attachment_key(&session.user_hash, &draft_id, &att_id).map_err(AppError::InternalError)?; let size = raw.len() as i64; let checksum = format!("{:x}", Sha256::digest(&raw)); let mut writer = storage .begin_upload(&key) .await .map_err(|e| AppError::ServiceUnavailable(format!("Storage unavailable: {e}")))?; if let Err(error) = writer.write(Bytes::from(raw)).await { if writer.abort().await.is_err() && let Err(cleanup_error) = storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await { tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue generated attachment cleanup"); } return Err(AppError::ServiceUnavailable(format!( "Failed to store message attachment: {error}" ))); } let stat = match writer.finish().await { Ok(stat) => stat, Err(error) => { if let Err(cleanup_error) = storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await { tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue failed generated attachment cleanup"); } return Err(AppError::ServiceUnavailable(format!( "Failed to finalize message attachment: {error}" ))); } }; if stat.size != size as u64 { if let Err(cleanup_error) = storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await { tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue invalid generated attachment cleanup"); } return Err(AppError::InternalError( "Message attachment size verification failed".to_string(), )); } let attachment = db::drafts::NewDraftAttachment { id: &att_id, draft_id: &draft_id, filename: &filename, content_type: &content_type, size, storage_key: &key, checksum_sha256: Some(&checksum), }; if let Err(error) = db::drafts::add_draft_attachment(&conn, &attachment) { if let Err(cleanup_error) = storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await { tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue unattached generated blob cleanup"); } return Err(AppError::InternalError(error)); } Ok(Json(serde_json::json!({ "attachments": [UploadResponse { id: att_id, filename, content_type, size, }], })) .into_response()) } pub async fn get_attachment_content( Extension(session): Extension, Extension(config): Extension>, Extension(storage): Extension>, AxumPath((draft_id, attachment_id)): AxumPath<(String, String)>, ) -> Result { db::drafts::validate_draft_id(&draft_id).map_err(AppError::BadRequest)?; let conn = db::pool::open_user_db(&config.data_dir, &session.user_hash) .map_err(|e| AppError::InternalError(format!("Failed to open database: {e}")))?; let attachments = db::drafts::get_draft_attachments(&conn, &draft_id).map_err(AppError::InternalError)?; let attachment = attachments .iter() .find(|a| a.id == attachment_id) .ok_or_else(|| AppError::NotFound("Attachment not found".to_string()))?; let stream = storage .get(&attachment.storage_key) .await .map_err(|e| AppError::ServiceUnavailable(format!("Storage read failed: {e}")))?; Ok(Response::builder() .header("content-type", &attachment.content_type) .header("content-length", attachment.size.max(0).to_string()) .header("cache-control", "private, max-age=3600") .body(axum::body::Body::from_stream(stream)) .unwrap()) } pub async fn delete_attachment( Extension(session): Extension, Extension(config): Extension>, Extension(storage): Extension>, AxumPath((draft_id, attachment_id)): AxumPath<(String, String)>, ) -> Result { db::drafts::validate_draft_id(&draft_id).map_err(AppError::BadRequest)?; let conn = db::pool::open_user_db(&config.data_dir, &session.user_hash) .map_err(|e| AppError::InternalError(format!("Failed to open database: {e}")))?; let transaction = conn.unchecked_transaction().map_err(|e| { AppError::InternalError(format!( "Failed to start attachment delete transaction: {e}" )) })?; let attachment = db::drafts::get_draft_attachment_for_update(&transaction, &draft_id, &attachment_id) .map_err(AppError::InternalError)? .ok_or_else(|| AppError::NotFound("Attachment not found".to_string()))?; let storage_key = attachment.storage_key.clone(); let cleanup_event = db::blob_delete_outbox::enqueue(&transaction, &storage_key) .map_err(AppError::InternalError)?; let deleted = db::drafts::delete_draft_attachment(&transaction, &attachment_id) .map_err(AppError::InternalError)?; if !deleted { return Err(AppError::NotFound("Attachment not found".to_string())); } transaction.commit().map_err(|e| { AppError::InternalError(format!("Failed to commit attachment deletion: {e}")) })?; if let Err(error) = storage::process_queued_delete(&conn, storage.as_ref(), &cleanup_event).await { tracing::warn!(error = %error, storage_key = %storage_key, "Failed to process queued attachment blob deletion"); } Ok(Json(DeleteResponse { status: "deleted".to_string(), }) .into_response()) }