forked from Zynh0722/nyazoom
download router
parent
b7bcd9c5f1
commit
e58b966319
57
src/main.rs
57
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));
|
||||
|
|
|
@ -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())
|
||||
}
|
|
@ -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;
|
Loading…
Reference in New Issue