diff --git a/src/engine.rs b/src/engine.rs index 39c82d2..e20d913 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -8,7 +8,6 @@ use std::{ use archived::Archive; use axum::extract::BodyStream; use bytes::{BufMut, Bytes, BytesMut}; -use hyper::StatusCode; use rand::Rng; use tokio::{ fs::File, @@ -21,7 +20,7 @@ use tokio::{ use tokio_stream::StreamExt; use walkdir::WalkDir; -use crate::view::ViewResponse; +use crate::view::{ViewSuccess, ViewError}; pub struct Engine { // state @@ -51,7 +50,7 @@ impl Engine { cache_lifetime, cache_mem_capacity, )), - upl_count: AtomicUsize::new(WalkDir::new(&save_path).into_iter().count()), // count the amount of files in the save path and initialise our cached count with it + upl_count: AtomicUsize::new(WalkDir::new(&save_path).min_depth(1).into_iter().count()), // count the amount of files in the save path and initialise our cached count with it base_url, save_path, @@ -208,7 +207,7 @@ impl Engine { Some(data) } - pub async fn get_upload(&self, original_path: &PathBuf) -> Result { + pub async fn get_upload(&self, original_path: &PathBuf) -> Result { // extract upload file name let name = original_path .file_name() @@ -222,7 +221,7 @@ impl Engine { // check if the upload exists, if not then 404 if !self.upload_exists(&path).await { - return Err(StatusCode::NOT_FOUND); + return Err(ViewError::NotFound); } // attempt to read upload from cache @@ -231,7 +230,7 @@ impl Engine { if let Some(data) = cached_data { info!(target: "get_upload", "got upload from cache!!"); - return Ok(ViewResponse::FromCache(data)); + return Ok(ViewSuccess::FromCache(data)); } else { let mut file = File::open(&path).await.unwrap(); @@ -258,8 +257,8 @@ impl Engine { } } Err(_) => { - return Err(StatusCode::INTERNAL_SERVER_ERROR); - }, + return Err(ViewError::InternalServerError); + } } } @@ -271,12 +270,12 @@ impl Engine { info!(target: "get_upload", "recached upload from disk!"); - return Ok(ViewResponse::FromCache(data)); + return Ok(ViewSuccess::FromCache(data)); } info!(target: "get_upload", "got upload from disk!"); - return Ok(ViewResponse::FromDisk(file)); + return Ok(ViewSuccess::FromDisk(file)); } } } diff --git a/src/view.rs b/src/view.rs index 8de516e..9828ae1 100644 --- a/src/view.rs +++ b/src/view.rs @@ -14,22 +14,27 @@ use hyper::StatusCode; use tokio::fs::File; use tokio_util::io::ReaderStream; -pub enum ViewResponse { +pub enum ViewSuccess { FromDisk(File), FromCache(Bytes), } -impl IntoResponse for ViewResponse { +pub enum ViewError { + NotFound, + InternalServerError, +} + +impl IntoResponse for ViewSuccess { fn into_response(self) -> Response { match self { - ViewResponse::FromDisk(file) => { + ViewSuccess::FromDisk(file) => { // create a streamed body response (we want to stream larger files) let reader = ReaderStream::new(file); let stream = StreamBody::new(reader); stream.into_response() } - ViewResponse::FromCache(data) => { + ViewSuccess::FromCache(data) => { // extract mutable headers from the response let mut res = data.into_response(); let headers = res.headers_mut(); @@ -43,11 +48,32 @@ impl IntoResponse for ViewResponse { } } +impl IntoResponse for ViewError { + fn into_response(self) -> Response { + match self { + ViewError::NotFound => { + // convert string into response, change status code + let mut res = "not found!".into_response(); + *res.status_mut() = StatusCode::NOT_FOUND; + + res + } + ViewError::InternalServerError => { + // convert string into response, change status code + let mut res = "internal server error!".into_response(); + *res.status_mut() = StatusCode::NOT_FOUND; + + res + } + } + } +} + #[axum::debug_handler] pub async fn view( State(engine): State>, Path(original_path): Path, -) -> Result { +) -> Result { // (hopefully) prevent path traversal, just check for any non-file components if original_path .components() @@ -55,7 +81,7 @@ pub async fn view( .any(|x| !matches!(x, Component::Normal(_))) { warn!(target: "view", "a request attempted path traversal"); - return Err(StatusCode::NOT_FOUND); + return Err(ViewError::NotFound); } engine.get_upload(&original_path).await