From 6df6b4e4f8e5c99625085e78ba7667d89cea8a7d Mon Sep 17 00:00:00 2001
From: Amneesh Singh <natto@weirdnatto.in>
Date: Sun, 9 Nov 2025 14:38:21 +0530
Subject: [PATCH] serve_file: add headers

Signed-off-by: Amneesh Singh <natto@weirdnatto.in>
---
 Cargo.lock  | 17 +++++++++++++++++
 Cargo.toml  |  3 ++-
 src/main.rs | 46 ++++++++++++++++++++++++++++++++++------------
 3 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 2197848..e3bd2f1 100644
--- a/Cargo.lock
+++ b/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",
diff --git a/Cargo.toml b/Cargo.toml
index 89f16ac..f29e80d 100644
--- a/Cargo.toml
+++ b/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 = "*"
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index 77e6bc8..013b391 100644
--- a/src/main.rs
+++ b/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(),
         }
     }
 }
@@ -233,11 +234,11 @@ async fn upload(mut payload: Multipart) -> Result<impl IntoResponse, YamafError>
                 }
 
                 responses.push(format!(
-                    r#"<a href="{proto}://{host}/{file}">{proto}://{host}/{file}</a> (size ~ {size}k)"#,
+                    r#"<a href="{proto}://{host}/{file}">{proto}://{host}/{file}</a> (size ~ {size:.2}k)"#,
                     proto = CONFIG.external_protocol,
                     host = CONFIG.external_host,
                     file = filename,
-                    size = written / 1024
+                    size = written as f64 / 1024 as f64
                 ));
             }
 
@@ -256,14 +257,35 @@ 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 {
-        Ok(file) => {
-            let stream = ReaderStream::new(file);
-            (StatusCode::OK, Body::from_stream(stream)).into_response()
-        }
-        Err(_) => (StatusCode::NOT_FOUND, "File not found!".to_string()).into_response(),
-    }
+    let metadata = tokio::fs::metadata(&path)
+        .await
+        .map_err(|_| YamafError::FileNotFound)?;
+    let file = File::open(&path)
+        .await
+        .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),
+        (header::ACCEPT_RANGES, "bytes".parse().unwrap()),
+    ]);
+
+    let stream = ReaderStream::new(file);
+    let body = Body::from_stream(stream);
+
+    Ok((StatusCode::OK, headers, body).into_response())
 }