limit downloads! and clear the files!!

main
Zynh0722 2023-04-12 06:40:55 -07:00
parent af92f779e6
commit abf88d8656
3 changed files with 61 additions and 13 deletions

21
dist/404.html vendored Normal file
View File

@ -0,0 +1,21 @@
<!DOCTYPE html>
<html lang="en">
<head>
<title></title>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<link href="css/link.css" rel="stylesheet">
</head>
<body>
<h1>
Link not found 3:
</h1>
<a class="return-button" href="/">
Return to home
</a>
</body>
</html>

View File

@ -1,9 +1,11 @@
use async_zip::tokio::write::ZipFileWriter; use async_zip::tokio::write::ZipFileWriter;
use async_zip::{Compression, ZipEntryBuilder}; use async_zip::{Compression, ZipEntryBuilder};
use axum::body::StreamBody;
use axum::extract::State; use axum::extract::State;
use axum::http::StatusCode; use axum::http::StatusCode;
use axum::routing::post; use axum::response::IntoResponse;
use axum::routing::{get, post};
use axum::{ use axum::{
extract::{DefaultBodyLimit, Multipart}, extract::{DefaultBodyLimit, Multipart},
response::Redirect, response::Redirect,
@ -28,7 +30,7 @@ use std::io;
use std::net::SocketAddr; use std::net::SocketAddr;
use std::path::Path; use std::path::Path;
use tokio_util::io::StreamReader; use tokio_util::io::{ReaderStream, StreamReader};
use tower_http::{limit::RequestBodyLimitLayer, services::ServeDir, trace::TraceLayer}; use tower_http::{limit::RequestBodyLimitLayer, services::ServeDir, trace::TraceLayer};
@ -63,21 +65,15 @@ async fn main() -> io::Result<()> {
let state = fetch_cache().await; let state = fetch_cache().await;
// Router Setup // Router Setup
let with_big_body = Router::new() let app = Router::new()
.route("/upload", post(upload_to_zip)) .route("/upload", post(upload_to_zip))
.route("/download/:id", get(download))
.layer(DefaultBodyLimit::disable()) .layer(DefaultBodyLimit::disable())
.layer(RequestBodyLimitLayer::new( .layer(RequestBodyLimitLayer::new(
10 * 1024 * 1024 * 1024, // 10GiB 10 * 1024 * 1024 * 1024, // 10GiB
)) ))
.with_state(state); .with_state(state)
let base = Router::new()
.nest_service("/", ServeDir::new("dist")) .nest_service("/", ServeDir::new("dist"))
.nest_service("/download", ServeDir::new(".cache/serve"));
let app = Router::new()
.merge(with_big_body)
.merge(base)
.layer(TraceLayer::new_for_http()); .layer(TraceLayer::new_for_http());
// Server creation // Server creation
@ -147,7 +143,34 @@ async fn upload_to_zip(
writer.close().await.unwrap(); writer.close().await.unwrap();
Ok(Redirect::to(&format!("/link.html?link={}.zip", cache_name))) Ok(Redirect::to(&format!("/link.html?link={}", cache_name)))
}
async fn download(
axum::extract::Path(id): axum::extract::Path<String>,
State(state): State<AppState>,
) -> Result<axum::response::Response, (StatusCode, String)> {
let mut records = state.records.lock().await;
if let Some(record) = records.get_mut(&id) {
if record.can_be_downloaded() {
record.downloads += 1;
let file = tokio::fs::File::open(&record.file).await.unwrap();
return Ok(axum::http::Response::builder()
.header("Content-Type", "application/zip")
.body(StreamBody::new(ReaderStream::new(file)))
.unwrap()
.into_response());
} else {
let _ = tokio::fs::remove_file(&record.file);
records.remove(&id);
write_to_cache(&records).await.unwrap();
}
}
Ok(Redirect::to("/404.html").into_response())
} }
#[inline] #[inline]

View File

@ -24,6 +24,10 @@ impl UploadRecord {
..Default::default() ..Default::default()
} }
} }
pub fn can_be_downloaded(&self) -> bool {
self.downloads < self.max_downloads
}
} }
impl Default for UploadRecord { impl Default for UploadRecord {
@ -32,7 +36,7 @@ impl Default for UploadRecord {
uploaded: Utc::now(), uploaded: Utc::now(),
file: Path::new("").to_owned(), file: Path::new("").to_owned(),
downloads: 0, downloads: 0,
max_downloads: 5, max_downloads: 1,
} }
} }
} }