serve_file: add headers

Signed-off-by: Amneesh Singh <natto@weirdnatto.in>

Cargo.lock

@@ -295,6 +295,16 @@ version = "0.3.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
 
+[[package]]
+name = "mime_guess"
+version = "2.0.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e"
+dependencies = [
+ "mime",
+ "unicase",
+]
+
 [[package]]
 name = "mio"
 version = "1.1.0"
@@ -661,6 +671,12 @@ dependencies = [
  "once_cell",
 ]
 
+[[package]]
+name = "unicase"
+version = "2.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "75b844d17643ee918803943289730bec8aac480150456169e647ed0b576ba539"
+
 [[package]]
 name = "unicode-ident"
 version = "1.0.22"
@@ -789,6 +805,7 @@ version = "0.1.0"
 dependencies = [
  "axum",
  "futures-util",
+ "mime_guess",
  "rand",
  "tokio",
  "tokio-util",

Cargo.toml

@@ -8,6 +8,7 @@ edition = "2024"
 [dependencies]
 axum = { version = "0.8.6", features = ["multipart"] }
 futures-util = "0.3.31"
+mime_guess= "*"
+rand = "*"
 tokio = { version = "1.48.0", features = ["full"] }
 tokio-util = { version = "0.7.17", features = ["io"] }
-rand = "*"

src/main.rs

@@ -2,7 +2,7 @@ use axum::{
     Router,
     body::Body,
     extract::{DefaultBodyLimit, Multipart, Path},
-    http::StatusCode,
+    http::{HeaderMap, StatusCode, header},
     response::{Html, IntoResponse, Response},
     routing::get,
 };
@@ -120,6 +120,7 @@ enum YamafError {
     BadRequest(String),
     InternalError(String),
     FileTooBig(String),
+    FileNotFound,
 }
 
 impl IntoResponse for YamafError {
@@ -129,12 +130,12 @@ impl IntoResponse for YamafError {
             YamafError::InternalError(msg) => {
                 (StatusCode::INTERNAL_SERVER_ERROR, msg).into_response()
             }
-
             YamafError::FileTooBig(filename) => (
                 StatusCode::PAYLOAD_TOO_LARGE,
                 format!("File {} is too big!", filename),
             )
                 .into_response(),
+            YamafError::FileNotFound => (StatusCode::NOT_FOUND, "File Not Found!").into_response(),
         }
     }
 }
@@ -256,14 +257,34 @@ async fn upload(mut payload: Multipart) -> Result<impl IntoResponse, YamafError>
     .into_response())
 }
 
-async fn serve_file(Path(filename): Path<String>) -> impl IntoResponse {
+async fn serve_file(Path(filename): Path<String>) -> Result<impl IntoResponse, YamafError> {
     let path = PathBuf::from(&CONFIG.root_dir).join(&filename);
 
-    match File::open(&path).await {
+    let metadata = tokio::fs::metadata(&path)
-        Ok(file) => {
+        .await
-            let stream = ReaderStream::new(file);
+        .map_err(|_| YamafError::FileNotFound)?;
-            (StatusCode::OK, Body::from_stream(stream)).into_response()
+    let file = File::open(&path)
-        }
+        .await
-        Err(_) => (StatusCode::NOT_FOUND, "File not found!".to_string()).into_response(),
+        .map_err(|_| YamafError::FileNotFound)?;
-    }
+    let mime = mime_guess::from_path(&path).first_or_octet_stream();
+
+    let content_type = mime
+        .as_ref()
+        .parse()
+        .map_err(|_| YamafError::InternalError("Something went wrong".into()))?;
+    let content_length = metadata
+        .len()
+        .to_string()
+        .parse()
+        .map_err(|_| YamafError::InternalError("Something went wrong".into()))?;
+
+    let headers = HeaderMap::from_iter([
+        (header::CONTENT_TYPE, content_type),
+        (header::CONTENT_LENGTH, content_length),
+    ]);
+
+    let stream = ReaderStream::new(file);
+    let body = Body::from_stream(stream);
+
+    Ok((StatusCode::OK, headers, body).into_response())
 }