diff --git a/src/main.rs b/src/main.rs index 27b830f..40f9080 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,27 +1,30 @@ use axum::body::Body; use axum::extract::{DefaultBodyLimit, Multipart, Path, State}; -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 axum::http::header::{CONTENT_LENGTH, CONTENT_TYPE, RETRY_AFTER}; +use axum::http::{HeaderMap, Method, StatusCode}; +use axum::response::{IntoResponse, Response}; +use axum::routing::{get, post}; +use axum::Router; use futures_util::TryStreamExt; use parking_lot::Mutex; use rand::distr::{Alphanumeric, SampleString}; use rand::rng; -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; use tokio::fs::File; -use tokio::io::AsyncWriteExt; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt}; use tokio_util::io::{ReaderStream, StreamReader}; use tower_http::cors::{Any, CorsLayer}; const MAX_UPLOAD_BYTES: u64 = 5 * 1024 * 1024; +const UPLOAD_CHUNK_SIZE: usize = 64 * 1024; +const DOWNLOAD_CHUNK_SIZE: usize = 64 * 1024; #[tokio::main] -async fn main() { - ensure_work_dir(); +async fn main() -> Result<(), Box> { + ensure_work_dir()?; + let cors = CorsLayer::new() .allow_origin(Any) .allow_methods(vec![Method::GET, Method::POST]) @@ -38,136 +41,226 @@ async fn main() { .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(); + let listener = tokio::net::TcpListener::bind("0.0.0.0:8089").await?; + axum::serve(listener, app).await?; + + Ok(()) } struct ApiState { work_directory: PathBuf, busy_ids: Mutex>, - counters: Mutex>, +} + +struct UploadGuard { + state: Arc, + id: String, + tmp_path: Option, + busy_inserted: bool, +} + +impl Drop for UploadGuard { + fn drop(&mut self) { + if self.busy_inserted { + self.state.busy_ids.lock().remove(&self.id); + } + + if let Some(tmp_path) = &self.tmp_path { + let _ = std::fs::remove_file(tmp_path); + } + } } async fn generate_id() -> String { - Alphanumeric.sample_string(&mut rng(), 24).to_uppercase() + Alphanumeric + .sample_string(&mut rng(), 24) + .to_ascii_uppercase() } 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) + + 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 + .unwrap_or(0) } 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()) + + let mut file = tokio::fs::OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .open(&tmp_meta) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + let data = value.to_string(); + + file.write_all(data.as_bytes()) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + file.sync_all() + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + drop(file); + tokio::fs::rename(&tmp_meta, &meta_path) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - state.counters.lock().insert(id.to_string(), value); + Ok(()) } async fn save_version( state: Arc, id: &str, - body: Body, + mut reader: impl AsyncRead + Unpin, ) -> Result { - { - let mut busy = state.busy_ids.lock(); - if !busy.insert(id.to_string()) { - return Err(StatusCode::CONFLICT); - } + let mut guard = UploadGuard { + state: state.clone(), + id: id.to_string(), + tmp_path: None, + busy_inserted: false, + }; + + if !state.busy_ids.lock().insert(guard.id.clone()) { + return Err(StatusCode::CONFLICT); } - let next = read_counter(&state, id).await + 1; + guard.busy_inserted = true; + + let current = read_counter(&state, id).await; + let next = current + .checked_add(1) + .ok_or(StatusCode::INTERNAL_SERVER_ERROR)?; let file_path = state.work_directory.join(id); - let tmp_path = state.work_directory.join(format!("{id}.tmp")); + let tmp_path = state.work_directory.join(format!( + "{id}.{}.tmp", + Alphanumeric.sample_string(&mut rng(), 8) + )); - 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 tmp_file = tokio::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&tmp_path) + .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) + guard.tmp_path = Some(tmp_path.clone()); + + let mut buf = vec![0u8; UPLOAD_CHUNK_SIZE]; + let mut written = 0u64; + + loop { + let n = reader + .read(&mut buf) .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + .map_err(|_| StatusCode::BAD_REQUEST)?; + + if n == 0 { + break; + } + + written += n as u64; if written > MAX_UPLOAD_BYTES { return Err(StatusCode::PAYLOAD_TOO_LARGE); } tmp_file - .flush() + .write_all(&buf[..n]) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - - 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; + tmp_file + .sync_all() + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + drop(tmp_file); + + tokio::fs::rename(&tmp_path, &file_path) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + guard.tmp_path = None; + + write_counter(&state, id, next).await?; + + Ok(next) +} + +fn error_response(code: StatusCode) -> Response { + if code == StatusCode::CONFLICT { + Response::builder() + .status(code) + .header(RETRY_AFTER, "1") + .body(Body::empty()) + .unwrap_or_else(|_| code.into_response()) + } else { + code.into_response() } - result } async fn upload_file_stream( State(state): State>, Path(id): Path, + headers: HeaderMap, body: Body, -) -> Result { +) -> Result { if !validate_id(&id) { - return Err(StatusCode::BAD_REQUEST); + return Err(StatusCode::BAD_REQUEST.into_response()); } - save_version(state, &id, body).await.map(|v| v.to_string()) + + if let Some(content_length) = headers + .get(CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()) + && content_length > MAX_UPLOAD_BYTES + { + return Err(StatusCode::PAYLOAD_TOO_LARGE.into_response()); + } + + let reader = StreamReader::new(body.into_data_stream().map_err(std::io::Error::other)); + + save_version(state, &id, reader) + .await + .map(|v| v.to_string()) + .map_err(error_response) } async fn upload_file( State(state): State>, Path(id): Path, mut multipart: Multipart, -) -> Result { +) -> Result { if !validate_id(&id) { - return Err(StatusCode::BAD_REQUEST); + return Err(StatusCode::BAD_REQUEST.into_response()); } + let field = multipart .next_field() .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? - .ok_or(StatusCode::BAD_REQUEST)?; - let data = field - .bytes() + .map_err(|_| StatusCode::BAD_REQUEST.into_response())? + .ok_or(StatusCode::BAD_REQUEST.into_response())?; + + let reader = StreamReader::new(field.map_err(|_| std::io::Error::other("multipart"))); + + save_version(state, &id, reader) .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - save_version(state, &id, Body::from(data)).await.map(|v| v.to_string()) + .map(|v| v.to_string()) + .map_err(error_response) } async fn download_file( @@ -177,14 +270,29 @@ async fn download_file( if !validate_id(&id) { return Err(StatusCode::BAD_REQUEST); } + let path = state.work_directory.join(id.as_str()); - let file = File::open(path) + + let file = File::open(&path) .await .map_err(|_| StatusCode::NOT_FOUND)?; - let body = Body::from_stream(ReaderStream::new(file)); + + let metadata = file + .metadata() + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + if !metadata.is_file() { + return Err(StatusCode::NOT_FOUND); + } + + let len = metadata.len(); + let body = Body::from_stream(ReaderStream::with_capacity(file, DOWNLOAD_CHUNK_SIZE)); + Response::builder() - .header(CONTENT_TYPE, "application/octet-stream") .status(StatusCode::OK) + .header(CONTENT_TYPE, "application/octet-stream") + .header(CONTENT_LENGTH, len.to_string()) .body(body) .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR) } @@ -196,22 +304,27 @@ async fn version( 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 path = state.work_directory.join(id.as_str()); + + match tokio::fs::metadata(&path).await { + Ok(metadata) if metadata.is_file() => {} + Ok(_) => return Err(StatusCode::NOT_FOUND), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => { + return Err(StatusCode::NOT_FOUND) + } + Err(_) => return Err(StatusCode::INTERNAL_SERVER_ERROR), } + Ok(read_counter(&state, &id).await.to_string()) } -fn ensure_work_dir() { +fn ensure_work_dir() -> std::io::Result<()> { let store_dir = get_dir(); - - if !store_dir.exists() { - std::fs::create_dir_all(&store_dir).unwrap(); - } + std::fs::create_dir_all(&store_dir)?; + Ok(()) } + #[inline] fn get_dir() -> PathBuf { PathBuf::from("/var/lib/learn_save") @@ -220,4 +333,4 @@ fn get_dir() -> PathBuf { #[inline] fn validate_id(id: &str) -> bool { id.len() == 24 && id.bytes().all(|b| matches!(b, b'A'..=b'Z' | b'0'..=b'9')) -} +} \ No newline at end of file