download router

This commit is contained in:
Zynh Ludwig 2024-08-29 23:32:49 -07:00
parent b7bcd9c5f1
commit e58b966319
3 changed files with 66 additions and 52 deletions

View file

@ -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));

52
src/router/download.rs Normal file
View file

@ -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())
}

9
src/router/mod.rs Normal file
View file

@ -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;