diff --git a/src/main.rs b/src/main.rs index c7954c0..27b830f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,22 +1,24 @@ use axum::body::Body; use axum::extract::{DefaultBodyLimit, Multipart, Path, State}; -use axum::http::Method; -use axum::http::StatusCode; -use axum::response::{IntoResponse, Response}; +use axum::http::header::CONTENT_TYPE; +use axum::http::{Method, StatusCode}; +use axum::response::Response; use axum::routing::post; use axum::{routing::get, Router}; use futures_util::TryStreamExt; -use parking_lot::RwLock; +use parking_lot::Mutex; use rand::distr::{Alphanumeric, SampleString}; use rand::rng; -use std::os::unix::prelude::MetadataExt; +use std::collections::{HashMap, HashSet}; use std::path::PathBuf; use std::sync::Arc; -use axum::http::header::CONTENT_TYPE; use tokio::fs::File; +use tokio::io::AsyncWriteExt; use tokio_util::io::{ReaderStream, StreamReader}; use tower_http::cors::{Any, CorsLayer}; +const MAX_UPLOAD_BYTES: u64 = 5 * 1024 * 1024; + #[tokio::main] async fn main() { ensure_work_dir(); @@ -31,11 +33,13 @@ async fn main() { .route("/upload/{id}", post(upload_file)) .route("/download/{id}", get(download_file)) .route("/{id}/version", get(version)) - .layer(DefaultBodyLimit::max(5 * 1024 * 1024)) + .layer(DefaultBodyLimit::max(MAX_UPLOAD_BYTES as usize)) .layer(cors) - .with_state(Arc::new(RwLock::new(ApiState { + .with_state(Arc::new(ApiState { work_directory: get_dir(), - }))); + busy_ids: Mutex::new(HashSet::new()), + counters: Mutex::new(HashMap::new()), + })); let listener = tokio::net::TcpListener::bind("0.0.0.0:8089").await.unwrap(); axum::serve(listener, app).await.unwrap(); @@ -43,110 +47,162 @@ async fn main() { struct ApiState { work_directory: PathBuf, + busy_ids: Mutex>, + counters: Mutex>, } async fn generate_id() -> String { Alphanumeric.sample_string(&mut rng(), 24).to_uppercase() } -async fn upload_file_stream( - State(state): State>>, - Path(id): Path, - body: Body, -) -> Result { - if !validate_id(id.as_str()) { - return Err(StatusCode::BAD_REQUEST); +async fn read_counter(state: &ApiState, id: &str) -> u64 { + if let Some(&counter) = state.counters.lock().get(id) { + return counter; } + let meta_path = state.work_directory.join(format!("{id}.meta")); + let counter = tokio::fs::read_to_string(&meta_path) + .await + .ok() + .and_then(|s| s.trim().parse::().ok()) + .unwrap_or(0); + state.counters.lock().insert(id.to_string(), counter); + counter +} - let temp_path; - let store_path; - { - let state = state.read(); - store_path = state.work_directory.join(id.as_str()); - temp_path = state.work_directory.join(format!("{id}.tmp")); - } - - let mut temp_file = tokio::fs::OpenOptions::new() - .write(true) - .create_new(true) - .open(&temp_path) +async fn write_counter(state: &ApiState, id: &str, value: u64) -> Result<(), StatusCode> { + let meta_path = state.work_directory.join(format!("{id}.meta")); + let tmp_meta = state.work_directory.join(format!("{id}.meta.tmp")); + tokio::fs::write(&tmp_meta, value.to_string()) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - - let stream = TryStreamExt::map_err(body.into_data_stream(), std::io::Error::other); - let mut reader = StreamReader::new(stream); - tokio::io::copy(&mut reader, &mut temp_file).await + tokio::fs::rename(&tmp_meta, &meta_path) + .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - - tokio::fs::rename(temp_path, store_path).await.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - + state.counters.lock().insert(id.to_string(), value); Ok(()) } -async fn upload_file( - Path(id): Path, - mut multipart: Multipart, -) -> Result { - if !validate_id(id.as_str()) { - return Err(StatusCode::FORBIDDEN); +async fn save_version( + state: Arc, + id: &str, + body: Body, +) -> Result { + { + let mut busy = state.busy_ids.lock(); + if !busy.insert(id.to_string()) { + return Err(StatusCode::CONFLICT); + } } - let file = multipart.next_field().await; + let next = read_counter(&state, id).await + 1; - if let Ok(Some(file)) = file { - let data = file - .bytes() + let file_path = state.work_directory.join(id); + let tmp_path = state.work_directory.join(format!("{id}.tmp")); + + let result = async { + let mut tmp_file = tokio::fs::OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .open(&tmp_path) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - let mut path = get_dir(); - path.push(id.clone()); - if !path.exists() { - File::create(path.clone()) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + let stream = body.into_data_stream().map_err(std::io::Error::other); + let mut reader = StreamReader::new(stream); + let written = tokio::io::copy(&mut reader, &mut tmp_file) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + if written > MAX_UPLOAD_BYTES { + return Err(StatusCode::PAYLOAD_TOO_LARGE); } - tokio::fs::write(path.clone(), data) + tmp_file + .flush() .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - let file = File::open(path).await.unwrap(); - Ok(file.metadata().await.unwrap().mtime().to_string()) - } else { - Err(StatusCode::BAD_REQUEST) + write_counter(&state, id, next).await?; + tokio::fs::rename(&tmp_path, &file_path) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + Ok(next) } + .await; + + state.busy_ids.lock().remove(id); + if result.is_err() { + let _ = tokio::fs::remove_file(&tmp_path).await; + } + result } -async fn download_file(Path(id): Path) -> Result { - if !validate_id(id.as_str()) { +async fn upload_file_stream( + State(state): State>, + Path(id): Path, + body: Body, +) -> Result { + if !validate_id(&id) { return Err(StatusCode::BAD_REQUEST); } - - let mut path = get_dir(); - path.push(id.clone()); - - let file = tokio::fs::File::open(path).await.map_err(|_| StatusCode::NOT_FOUND)?; - let stream = ReaderStream::new(file); - let body = Body::from_stream(stream); - let response = Response::builder() - .header(CONTENT_TYPE, "application/octet-stream") // Or detect mime type - .status(StatusCode::OK).body(body) - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - Ok(response) + save_version(state, &id, body).await.map(|v| v.to_string()) } -async fn version(Path(id): Path) -> Result { - if !validate_id(id.as_str()) { +async fn upload_file( + State(state): State>, + Path(id): Path, + mut multipart: Multipart, +) -> Result { + if !validate_id(&id) { + return Err(StatusCode::BAD_REQUEST); + } + let field = multipart + .next_field() + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .ok_or(StatusCode::BAD_REQUEST)?; + let data = field + .bytes() + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + save_version(state, &id, Body::from(data)).await.map(|v| v.to_string()) +} + +async fn download_file( + State(state): State>, + Path(id): Path, +) -> Result { + if !validate_id(&id) { + return Err(StatusCode::BAD_REQUEST); + } + let path = state.work_directory.join(id.as_str()); + let file = File::open(path) + .await + .map_err(|_| StatusCode::NOT_FOUND)?; + let body = Body::from_stream(ReaderStream::new(file)); + Response::builder() + .header(CONTENT_TYPE, "application/octet-stream") + .status(StatusCode::OK) + .body(body) + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR) +} + +async fn version( + State(state): State>, + Path(id): Path, +) -> Result { + if !validate_id(&id) { + return Err(StatusCode::BAD_REQUEST); + } + if !tokio::fs::try_exists(state.work_directory.join(id.as_str())) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + { return Err(StatusCode::NOT_FOUND); } - - let dir = get_dir().join(id); - let meta = tokio::fs::metadata(dir) - .await - .map_err(|_| StatusCode::BAD_REQUEST)?; - - Ok(meta.mtime().to_string()) + Ok(read_counter(&state, &id).await.to_string()) } fn ensure_work_dir() {