use axum::{ extract::{Multipart, Path, State}, http::StatusCode, response::{IntoResponse, Redirect, Response}, }; use askama::Template; use askama_web::WebTemplate; use sqlx::FromRow; use crate::auth_guard::{AuthUser, OptionalAuth}; use crate::state::AppState; const MAX_FILE_SIZE: usize = 10 * 1024 * 1024; const FREE_STORAGE_LIMIT: i64 = 50 * 1024 * 1024; const PRO_STORAGE_LIMIT: i64 = 1024 * 1024 * 1024; const THUMB_MAX_DIM: u32 = 300; const ALLOWED_TYPES: &[&str] = &["image/jpeg", "image/png", "image/gif", "image/webp"]; fn detect_content_type(data: &[u8]) -> Option<&'static str> { if data.len() < 12 { return None; } if data[0..3] == [0xFF, 0xD8, 0xFF] { Some("image/jpeg") } else if data[0..4] == [0x89, 0x50, 0x4E, 0x47] { Some("image/png") } else if data[0..3] == [0x47, 0x49, 0x46] { Some("image/gif") } else if data[0..4] == [0x52, 0x49, 0x46, 0x46] && data[8..12] == [0x57, 0x45, 0x42, 0x50] { Some("image/webp") } else { None } } fn generate_thumbnail(data: &[u8]) -> Result, String> { let img = ::image::load_from_memory(data).map_err(|e| format!("image decode failed: {e}"))?; let thumb = img.resize(THUMB_MAX_DIM, THUMB_MAX_DIM, ::image::imageops::FilterType::Lanczos3); let mut buf = std::io::Cursor::new(Vec::new()); thumb .write_to(&mut buf, ::image::ImageFormat::WebP) .map_err(|e| format!("thumbnail encode failed: {e}"))?; Ok(buf.into_inner()) } fn format_time(t: time::OffsetDateTime) -> String { format!( "{:04}-{:02}-{:02} {:02}:{:02} UTC", t.year(), t.month() as u8, t.day(), t.hour(), t.minute(), ) } fn format_size(bytes: i64) -> String { if bytes < 1024 { format!("{bytes} B") } else if bytes < 1024 * 1024 { format!("{:.1} KB", bytes as f64 / 1024.0) } else { format!("{:.1} MB", bytes as f64 / (1024.0 * 1024.0)) } } #[derive(FromRow)] struct ImageRow { id: String, user_id: String, filename: String, content_type: String, size_bytes: i64, width: Option, height: Option, created_at: time::OffsetDateTime, } #[derive(Template, WebTemplate)] #[template(path = "upload.html")] #[allow(dead_code)] pub struct UploadTemplate { pub email: Option, } pub async fn index(AuthUser(user): AuthUser) -> UploadTemplate { UploadTemplate { email: user.email, } } pub async fn upload( State(state): State, AuthUser(user): AuthUser, mut multipart: Multipart, ) -> Result { let is_pro = user.is_pro(); let storage_limit = if is_pro { PRO_STORAGE_LIMIT } else { FREE_STORAGE_LIMIT }; let used: (Option,) = sqlx::query_as("SELECT sum(size_bytes) FROM images WHERE user_id = $1") .bind(&user.sub) .fetch_one(&state.db) .await .map_err(|e| { tracing::error!("storage query failed: {e}"); (StatusCode::INTERNAL_SERVER_ERROR, "internal error").into_response() })?; let used_bytes = used.0.unwrap_or(0); let mut file_data: Option<(String, String, Vec)> = None; while let Ok(Some(field)) = multipart.next_field().await { if field.name() == Some("file") { let filename = field .file_name() .unwrap_or("upload") .to_string(); let content_type = field .content_type() .unwrap_or("application/octet-stream") .to_string(); let data = field.bytes().await.map_err(|e| { tracing::error!("multipart read failed: {e}"); (StatusCode::BAD_REQUEST, "failed to read upload").into_response() })?; file_data = Some((filename, content_type, data.to_vec())); break; } } let (filename, claimed_type, data) = file_data.ok_or_else(|| { (StatusCode::BAD_REQUEST, "no file uploaded").into_response() })?; if data.len() > MAX_FILE_SIZE { return Err((StatusCode::BAD_REQUEST, "file too large (10MB max)").into_response()); } let detected_type = detect_content_type(&data).ok_or_else(|| { (StatusCode::BAD_REQUEST, "unsupported image format").into_response() })?; if !ALLOWED_TYPES.contains(&claimed_type.as_str()) || detected_type != claimed_type { return Err( (StatusCode::BAD_REQUEST, "content type mismatch or unsupported format").into_response(), ); } let size_bytes = data.len() as i64; if used_bytes + size_bytes > storage_limit { return Err(( StatusCode::FORBIDDEN, "storage limit reached -- upgrade to pro for more space", ) .into_response()); } let dimensions = ::image::load_from_memory(&data) .ok() .map(|img| (img.width() as i32, img.height() as i32)); let (width, height) = dimensions.unzip(); let thumb_data = generate_thumbnail(&data).map_err(|e| { tracing::error!("thumbnail generation failed: {e}"); (StatusCode::INTERNAL_SERVER_ERROR, "thumbnail failed").into_response() })?; let id = nanoid::nanoid!(8); let original_key = format!("{id}/{filename}"); let thumb_key = format!("{id}/thumb.webp"); state .storage .put(&original_key, &data, &claimed_type) .await .map_err(|e| { tracing::error!("S3 upload failed: {e}"); (StatusCode::INTERNAL_SERVER_ERROR, "upload failed").into_response() })?; state .storage .put(&thumb_key, &thumb_data, "image/webp") .await .map_err(|e| { tracing::error!("S3 thumbnail upload failed: {e}"); (StatusCode::INTERNAL_SERVER_ERROR, "thumbnail upload failed").into_response() })?; sqlx::query( "INSERT INTO images (id, user_id, filename, content_type, size_bytes, width, height) VALUES ($1, $2, $3, $4, $5, $6, $7)", ) .bind(&id) .bind(&user.sub) .bind(&filename) .bind(&claimed_type) .bind(size_bytes) .bind(width) .bind(height) .execute(&state.db) .await .map_err(|e| { tracing::error!("image insert failed: {e}"); (StatusCode::INTERNAL_SERVER_ERROR, "internal error").into_response() })?; Ok(Redirect::see_other(&format!("/{id}"))) } #[derive(Template, WebTemplate)] #[template(path = "view.html")] pub struct ViewTemplate { pub id: String, pub filename: String, pub content_type: String, pub size: String, pub width: Option, pub height: Option, pub created_at: String, pub is_owner: bool, } pub async fn view( State(state): State, OptionalAuth(user): OptionalAuth, Path(id): Path, ) -> Result { let row = sqlx::query_as::<_, ImageRow>("SELECT * FROM images WHERE id = $1") .bind(&id) .fetch_optional(&state.db) .await .map_err(|e| { tracing::error!("image query failed: {e}"); StatusCode::INTERNAL_SERVER_ERROR })? .ok_or(StatusCode::NOT_FOUND)?; let is_owner = user.map(|u| u.sub == row.user_id).unwrap_or(false); Ok(ViewTemplate { id: row.id, filename: row.filename, content_type: row.content_type, size: format_size(row.size_bytes), width: row.width, height: row.height, created_at: format_time(row.created_at), is_owner, }) } pub async fn raw( State(state): State, Path(id): Path, ) -> Result { let row = sqlx::query_as::<_, ImageRow>("SELECT * FROM images WHERE id = $1") .bind(&id) .fetch_optional(&state.db) .await .map_err(|e| { tracing::error!("image query failed: {e}"); StatusCode::INTERNAL_SERVER_ERROR })? .ok_or(StatusCode::NOT_FOUND)?; let key = format!("{}/{}", row.id, row.filename); let (data, content_type) = state.storage.get(&key).await.map_err(|e| { tracing::error!("S3 get failed: {e}"); StatusCode::INTERNAL_SERVER_ERROR })?; Ok(( [(axum::http::header::CONTENT_TYPE, content_type)], data, ) .into_response()) } pub async fn thumb( State(state): State, Path(id): Path, ) -> Result { let exists: bool = sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM images WHERE id = $1)") .bind(&id) .fetch_one(&state.db) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; if !exists { return Err(StatusCode::NOT_FOUND); } let key = format!("{id}/thumb.webp"); let (data, _) = state.storage.get(&key).await.map_err(|e| { tracing::error!("S3 thumb get failed: {e}"); StatusCode::INTERNAL_SERVER_ERROR })?; Ok(( [(axum::http::header::CONTENT_TYPE, "image/webp".to_string())], data, ) .into_response()) } #[allow(dead_code)] pub struct ImageEntry { pub id: String, pub filename: String, pub size: String, pub created_at: String, } #[derive(Template, WebTemplate)] #[template(path = "my.html")] pub struct MyTemplate { pub images: Vec, } pub async fn my_images( State(state): State, AuthUser(user): AuthUser, ) -> Result { let rows = sqlx::query_as::<_, ImageRow>( "SELECT * FROM images WHERE user_id = $1 ORDER BY created_at DESC", ) .bind(&user.sub) .fetch_all(&state.db) .await .map_err(|e| { tracing::error!("image list query failed: {e}"); StatusCode::INTERNAL_SERVER_ERROR })?; let images = rows .into_iter() .map(|r| ImageEntry { id: r.id, filename: r.filename, size: format_size(r.size_bytes), created_at: format_time(r.created_at), }) .collect(); Ok(MyTemplate { images }) } pub async fn delete( State(state): State, AuthUser(user): AuthUser, Path(id): Path, ) -> Result { let row = sqlx::query_as::<_, ImageRow>("SELECT * FROM images WHERE id = $1 AND user_id = $2") .bind(&id) .bind(&user.sub) .fetch_optional(&state.db) .await .map_err(|e| { tracing::error!("image query failed: {e}"); StatusCode::INTERNAL_SERVER_ERROR })? .ok_or(StatusCode::NOT_FOUND)?; let original_key = format!("{}/{}", row.id, row.filename); let thumb_key = format!("{}/thumb.webp", row.id); let _ = state.storage.delete(&original_key).await; let _ = state.storage.delete(&thumb_key).await; sqlx::query("DELETE FROM images WHERE id = $1") .bind(&id) .execute(&state.db) .await .map_err(|e| { tracing::error!("image delete failed: {e}"); StatusCode::INTERNAL_SERVER_ERROR })?; Ok(Redirect::see_other("/my")) }