From e58b966319249e6899198e2a7ab87c8ca5369606 Mon Sep 17 00:00:00 2001
From: Zynh Ludwig <zynh0722@gmail.com>
Date: Thu, 29 Aug 2024 23:32:49 -0700
Subject: [PATCH] download router

---
 src/main.rs            | 57 ++++--------------------------------------
 src/router/download.rs | 52 ++++++++++++++++++++++++++++++++++++++
 src/router/mod.rs      |  9 +++++++
 3 files changed, 66 insertions(+), 52 deletions(-)
 create mode 100644 src/router/download.rs
 create mode 100644 src/router/mod.rs

diff --git a/src/main.rs b/src/main.rs
index 573ae17..de9208e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,11 +1,7 @@
-use tokio_util::io::ReaderStream;
-
 use axum::{
-    body::Body,
-    extract::{ConnectInfo, Request, State},
-    http::{HeaderMap, StatusCode},
+    extract::{ConnectInfo, Request},
     middleware::{self, Next},
-    response::{Html, IntoResponse, Redirect},
+    response::{Html, IntoResponse},
     routing::get,
     Router,
 };
@@ -14,25 +10,18 @@ use tower_http::{services::ServeDir, trace::TraceLayer};
 
 use std::{io, net::SocketAddr};
 
-mod router {
-    pub mod link;
-    pub mod records;
-    pub mod upload;
-}
 mod cache;
+mod router;
 mod state;
 mod util;
 mod views;
 
 use util::{headers::ForwardedFor, logging, ssr, sweeper};
 
+use router::*;
 use state::*;
 use views::*;
 
-use router::link::get_link_router;
-use router::records::get_records_router;
-use router::upload::get_upload_router;
-
 #[tokio::main]
 async fn main() -> io::Result<()> {
     logging::init_tracing();
@@ -47,7 +36,7 @@ async fn main() -> io::Result<()> {
     // Router Setup
     let app = Router::new()
         .route("/", get(welcome))
-        .route("/download/:id", get(download))
+        .nest("/download", get_download_router())
         .nest("/upload", get_upload_router())
         .nest("/records", get_records_router())
         .nest("/link", get_link_router())
@@ -79,42 +68,6 @@ async fn log_source(
     next.run(req).await
 }
 
-async fn download(
-    axum::extract::Path(id): axum::extract::Path<String>,
-    headers: HeaderMap,
-    State(state): State<AppState>,
-) -> Result<axum::response::Response, (StatusCode, String)> {
-    {
-        let mut records = state.records.lock().await;
-        if headers.get("hx-request").is_some() {
-            return Ok(axum::http::Response::builder()
-                .header("HX-Redirect", format!("/download/{id}"))
-                .status(204)
-                .body("".to_owned())
-                .unwrap()
-                .into_response());
-        }
-
-        if let Some(record) = records
-            .get_mut(&id)
-            .filter(|record| record.can_be_downloaded())
-        {
-            record.downloads += 1;
-
-            let file = tokio::fs::File::open(&record.file).await.unwrap();
-
-            return Ok(axum::response::Response::builder()
-                .header("Content-Type", "application/zip")
-                .body(Body::from_stream(ReaderStream::new(file)))
-                .unwrap());
-        } else {
-            records.remove_record(&id).await.unwrap()
-        }
-    }
-
-    Ok(Redirect::to("/404.html").into_response())
-}
-
 async fn serve(app: Router) {
     // // Server creation
     let addr = SocketAddr::from(([0, 0, 0, 0], 3000));
diff --git a/src/router/download.rs b/src/router/download.rs
new file mode 100644
index 0000000..d0817b3
--- /dev/null
+++ b/src/router/download.rs
@@ -0,0 +1,52 @@
+use axum::{
+    body::Body,
+    extract::State,
+    http::HeaderMap,
+    response::{IntoResponse, Redirect},
+    routing::get,
+    Router,
+};
+use reqwest::StatusCode;
+use tokio_util::io::ReaderStream;
+
+use crate::{AppState, AsyncRemoveRecord};
+
+pub fn get_download_router() -> Router<AppState> {
+    Router::new().route("/:id", get(download))
+}
+
+async fn download(
+    axum::extract::Path(id): axum::extract::Path<String>,
+    headers: HeaderMap,
+    State(state): State<AppState>,
+) -> Result<axum::response::Response, (StatusCode, String)> {
+    {
+        let mut records = state.records.lock().await;
+        if headers.get("hx-request").is_some() {
+            return Ok(axum::http::Response::builder()
+                .header("HX-Redirect", format!("/download/{id}"))
+                .status(204)
+                .body("".to_owned())
+                .unwrap()
+                .into_response());
+        }
+
+        if let Some(record) = records
+            .get_mut(&id)
+            .filter(|record| record.can_be_downloaded())
+        {
+            record.downloads += 1;
+
+            let file = tokio::fs::File::open(&record.file).await.unwrap();
+
+            return Ok(axum::response::Response::builder()
+                .header("Content-Type", "application/zip")
+                .body(Body::from_stream(ReaderStream::new(file)))
+                .unwrap());
+        } else {
+            records.remove_record(&id).await.unwrap()
+        }
+    }
+
+    Ok(Redirect::to("/404.html").into_response())
+}
diff --git a/src/router/mod.rs b/src/router/mod.rs
new file mode 100644
index 0000000..85d87ac
--- /dev/null
+++ b/src/router/mod.rs
@@ -0,0 +1,9 @@
+pub mod download;
+pub mod link;
+pub mod records;
+pub mod upload;
+
+pub use download::get_download_router;
+pub use link::get_link_router;
+pub use records::get_records_router;
+pub use upload::get_upload_router;