From ea7dee4dfd519d47d5ca2ba35a6706c532885462 Mon Sep 17 00:00:00 2001 From: Zynh Ludwig Date: Thu, 29 Aug 2024 23:16:14 -0700 Subject: [PATCH] extract upload router --- src/main.rs | 118 +++++++++---------------------------------- src/router/upload.rs | 101 ++++++++++++++++++++++++++++++++++++ 2 files changed, 126 insertions(+), 93 deletions(-) create mode 100644 src/router/upload.rs diff --git a/src/main.rs b/src/main.rs index f4ec960..2289e4e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,27 +1,22 @@ -use async_zip::{tokio::write::ZipFileWriter, Compression, ZipEntryBuilder}; - -use futures::TryStreamExt; -use tokio_util::{ - compat::FuturesAsyncWriteCompatExt, - io::{ReaderStream, StreamReader}, -}; +use tokio_util::io::ReaderStream; use axum::{ body::Body, - extract::{ConnectInfo, DefaultBodyLimit, Multipart, Request, State}, - http::{HeaderMap, Response, StatusCode}, + extract::{ConnectInfo, Request, State}, + http::{HeaderMap, StatusCode}, middleware::{self, Next}, response::{Html, IntoResponse, Redirect}, - routing::{get, post}, + routing::get, Json, Router, }; use axum_extra::TypedHeader; -use tower_http::{limit::RequestBodyLimitLayer, services::ServeDir, trace::TraceLayer}; +use tower_http::{services::ServeDir, trace::TraceLayer}; -use sanitize_filename_reader_friendly::sanitize; - -use std::{io, net::SocketAddr, path::Path}; +use std::{io, net::SocketAddr}; +mod router { + pub mod upload; +} mod cache; mod state; mod util; @@ -29,9 +24,12 @@ mod views; use util::{headers::ForwardedFor, logging, ssr, sweeper}; +use router::*; use state::*; use views::*; +use upload::get_upload_router; + #[tokio::main] async fn main() -> io::Result<()> { logging::init_tracing(); @@ -43,19 +41,23 @@ async fn main() -> io::Result<()> { sweeper::spawn(state.clone()); + // Records views + let record_router = Router::new() + .route("/", get(records)) + .route("/links", get(records_links)); + + // Link pages + let link_router = Router::new() + .route("/:id", get(link).delete(link_delete)) + .route("/:id/remaining", get(remaining)); + // Router Setup let app = Router::new() .route("/", get(welcome)) - .route("/upload", post(upload_to_zip)) - .route("/records", get(records)) - .route("/records/links", get(records_links)) .route("/download/:id", get(download)) - .route("/link/:id", get(link).delete(link_delete)) - .route("/link/:id/remaining", get(remaining)) - .layer(DefaultBodyLimit::disable()) - .layer(RequestBodyLimitLayer::new( - 10 * 1024 * 1024 * 1024, // 10GiB - )) + .nest("/upload", get_upload_router()) + .nest("/records", record_router) + .nest("/link", link_router) .with_state(state) .fallback_service(ServeDir::new("dist")) .layer(TraceLayer::new_for_http()) @@ -180,76 +182,6 @@ async fn log_source( next.run(req).await } -async fn upload_to_zip( - State(state): State, - mut body: Multipart, -) -> Result, (StatusCode, String)> { - tracing::debug!("{:?}", *state.records.lock().await); - - let cache_name = util::get_random_name(10); - - let archive_path = Path::new(".cache/serve").join(format!("{}.zip", &cache_name)); - - tracing::debug!("Zipping: {:?}", &archive_path); - - let mut archive = tokio::fs::File::create(&archive_path) - .await - .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?; - let mut writer = ZipFileWriter::with_tokio(&mut archive); - - while let Some(field) = body.next_field().await.unwrap() { - let file_name = match field.file_name() { - Some(file_name) => sanitize(file_name), - _ => continue, - }; - - tracing::debug!("Downloading to Zip: {file_name:?}"); - - let stream = field; - let body_with_io_error = stream.map_err(io::Error::other); - let mut body_reader = StreamReader::new(body_with_io_error); - - let builder = ZipEntryBuilder::new(file_name.into(), Compression::Deflate); - let mut entry_writer = writer - .write_entry_stream(builder) - .await - .unwrap() - .compat_write(); - - tokio::io::copy(&mut body_reader, &mut entry_writer) - .await - .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?; - - entry_writer - .into_inner() - .close() - .await - .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?; - } - - let mut records = state.records.lock().await; - let record = UploadRecord::new(archive_path); - records.insert(cache_name.clone(), record.clone()); - - cache::write_to_cache(&records) - .await - .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?; - - writer.close().await.unwrap(); - - let id = cache_name; - let response = Response::builder() - .status(200) - .header("Content-Type", "text/html") - .header("HX-Push-Url", format!("/link/{}", &id)) - .body(ssr::render(|| { - leptos::view! { } - })) - .unwrap(); - - Ok(response) -} - async fn download( axum::extract::Path(id): axum::extract::Path, headers: HeaderMap, diff --git a/src/router/upload.rs b/src/router/upload.rs new file mode 100644 index 0000000..750643b --- /dev/null +++ b/src/router/upload.rs @@ -0,0 +1,101 @@ +use std::path::Path; + +use async_zip::{base::write::ZipFileWriter, Compression, ZipEntryBuilder}; +use axum::{ + extract::{DefaultBodyLimit, Multipart, State}, + http::Response, + routing::post, + Router, +}; +use futures::TryStreamExt; +use reqwest::StatusCode; +use sanitize_filename_reader_friendly::sanitize; +use tokio::io; +use tokio_util::{compat::FuturesAsyncWriteCompatExt, io::StreamReader}; +use tower_http::limit::RequestBodyLimitLayer; + +use crate::{ + cache, + util::{self, ssr}, + AppState, LinkView, UploadRecord, +}; + +pub fn get_upload_router() -> Router { + // Upload needs a subrouter to increase the body limit + Router::new() + .route("/", post(upload_to_zip)) + .layer(DefaultBodyLimit::disable()) + .layer(RequestBodyLimitLayer::new( + 10 * 1024 * 1024 * 1024, // 10GiB + )) +} + +async fn upload_to_zip( + State(state): State, + mut body: Multipart, +) -> Result, (StatusCode, String)> { + tracing::debug!("{:?}", *state.records.lock().await); + + let cache_name = util::get_random_name(10); + + let archive_path = Path::new(".cache/serve").join(format!("{}.zip", &cache_name)); + + tracing::debug!("Zipping: {:?}", &archive_path); + + let mut archive = tokio::fs::File::create(&archive_path) + .await + .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?; + let mut writer = ZipFileWriter::with_tokio(&mut archive); + + while let Some(field) = body.next_field().await.unwrap() { + let file_name = match field.file_name() { + Some(file_name) => sanitize(file_name), + _ => continue, + }; + + tracing::debug!("Downloading to Zip: {file_name:?}"); + + let stream = field; + let body_with_io_error = stream.map_err(io::Error::other); + let mut body_reader = StreamReader::new(body_with_io_error); + + let builder = ZipEntryBuilder::new(file_name.into(), Compression::Deflate); + let mut entry_writer = writer + .write_entry_stream(builder) + .await + .unwrap() + .compat_write(); + + tokio::io::copy(&mut body_reader, &mut entry_writer) + .await + .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?; + + entry_writer + .into_inner() + .close() + .await + .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?; + } + + let mut records = state.records.lock().await; + let record = UploadRecord::new(archive_path); + records.insert(cache_name.clone(), record.clone()); + + cache::write_to_cache(&records) + .await + .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?; + + writer.close().await.unwrap(); + + let id = cache_name; + let response = Response::builder() + .status(200) + .header("Content-Type", "text/html") + .header("HX-Push-Url", format!("/link/{}", &id)) + .body(ssr::render(|| { + leptos::view! { } + })) + .unwrap(); + + Ok(response) +}