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..0e13e5b 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(), } } } @@ -256,14 +257,34 @@ async fn upload(mut payload: Multipart) -> Result .into_response()) } -async fn serve_file(Path(filename): Path) -> impl IntoResponse { +async fn serve_file(Path(filename): Path) -> Result { 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), + ]); + + let stream = ReaderStream::new(file); + let body = Body::from_stream(stream); + + Ok((StatusCode::OK, headers, body).into_response()) }