improved image upload: no more hardcoded png extension

src/routes/users.rs

@@ -1,7 +1,7 @@
 use axum::{
     Extension, Json, Router,
     extract::{Path, Request},
-    http::{HeaderMap, StatusCode},
+    http::{HeaderMap, StatusCode, header},
     middleware::Next,
     response::Response,
     routing::{get, post, put},
@@ -12,7 +12,7 @@ use uuid::Uuid;
 use validator::ValidateEmail;
 
 use crate::{
-    AppConfig, MAX_USERNAME_LENGTH,
+    AppConfig, MAX_UPLOAD_SIZE, MAX_USERNAME_LENGTH,
     auth::{create_jwt, hash_password, validate_token, verify_jwt, verify_password},
     db::{user_id_from_uuid, username_from_uuid},
     errors::APIError,
@@ -257,18 +257,49 @@ async fn upload_avatar(
 ) -> Result<StatusCode, APIError> {
     let claims = verify_jwt(headers)?;
 
+    if body.len() > MAX_UPLOAD_SIZE {
+        // TODO: FileTooLarge error
+        return Err(APIError::WrongFileFormat);
+    }
+
+    let kind = infer::get(&body).ok_or(APIError::WrongFileFormat)?;
+
+    let ("image/png" | "image/jpeg" | "image/webp") = kind.mime_type() else {
+        return Err(APIError::WrongFileFormat);
+    };
+
     let user_id = user_id_from_uuid(&db, claims.sub).await?;
-    tracing::debug!("User ID {} is uploading {} bytes)", user_id, body.len());
+    tracing::debug!(
+        "User ID {} is uploading {} bytes ({})",
+        user_id,
+        body.len(),
+        kind.mime_type()
+    );
 
     let base_dir = &config.avatar_dir;
-    let file_extension = "png"; // TODO: detect MIME type
-    let filename = format!("{}.{}", claims.sub, file_extension);
-    let full_path = std::path::Path::new(&base_dir).join(&filename);
+
+    let supported_extensions = ["png", "jpg", "jpeg", "webp"];
 
     tokio::fs::create_dir_all(&base_dir)
         .await
         .map_err(|e| APIError::Internal(format!("Failed to create storage: {e}")))?;
 
+    // Delete all other files for this user first
+    for ext in supported_extensions {
+        let old_filename = format!("{}.{}", claims.sub, ext);
+        let old_path = std::path::Path::new(base_dir).join(&old_filename);
+
+        match tokio::fs::remove_file(old_path).await {
+            Ok(_) => tracing::debug!("Deleted old avatar: {}", old_filename),
+            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
+            Err(e) => tracing::warn!("Failed to delete old avatar {}: {}", old_filename, e),
+        }
+    }
+
+    let file_extension = kind.extension();
+    let filename = format!("{}.{}", claims.sub, file_extension);
+    let full_path = std::path::Path::new(&base_dir).join(&filename);
+
     tokio::fs::write(&full_path, body)
         .await
         .map_err(|e| APIError::Internal(format!("Failed to save file: {e}")))?;
@@ -288,19 +319,29 @@ async fn get_avatar(
     Extension(config): Extension<Arc<AppConfig>>,
 ) -> Result<Response, APIError> {
     let base_dir = &config.avatar_dir;
-    let filename = format!("{}.png", uuid);
-    let full_path = std::path::Path::new(&base_dir).join(filename);
 
-    if !full_path.exists() {
-        return Err(APIError::AvatarNotFound);
+    // Helper to try finding the file with allowed extensions
+    let mut file_path = None;
+    for ext in ["png", "jpg", "jpeg", "webp"] {
+        let path = std::path::Path::new(&base_dir).join(format!("{}.{}", uuid, ext));
+        if path.exists() {
+            file_path = Some(path);
+            break;
+        }
     }
 
+    let full_path = file_path.ok_or(APIError::AvatarNotFound)?;
+
     let file_contents = tokio::fs::read(&full_path)
         .await
         .map_err(|e| APIError::Internal(format!("Could not read avatar file: {e}")))?;
 
+    let mime_type = infer::get(&file_contents)
+        .map(|k| k.mime_type())
+        .unwrap_or("application/octet-stream"); // Fallback
+
     Ok(Response::builder()
-        .header("Content-Type", "image/png")
+        .header(header::CONTENT_TYPE, mime_type)
         .body(axum::body::Body::from(file_contents))
         .unwrap())
 }