More AI changes

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<dyn std::error::Error>> {
+    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<HashSet<String>>,
-    counters: Mutex<HashMap<String, u64>>,
+}
+
+struct UploadGuard {
+    state: Arc<ApiState>,
+    id: String,
+    tmp_path: Option<PathBuf>,
+    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::<u64>().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<ApiState>,
     id: &str,
-    body: Body,
+    mut reader: impl AsyncRead + Unpin,
 ) -> Result<u64, StatusCode> {
-    {
-        let mut busy = state.busy_ids.lock();
-        if !busy.insert(id.to_string()) {
+    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)
+        .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)?;
+    }
+
+    tmp_file
+        .sync_all()
         .await
         .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
 
-        write_counter(&state, id, next).await?;
+    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)
 }
-    .await;
 
-    state.busy_ids.lock().remove(id);
-    if result.is_err() {
-        let _ = tokio::fs::remove_file(&tmp_path).await;
+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<Arc<ApiState>>,
     Path(id): Path<String>,
+    headers: HeaderMap,
     body: Body,
-) -> Result<String, StatusCode> {
+) -> Result<String, Response> {
     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::<u64>().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<Arc<ApiState>>,
     Path(id): Path<String>,
     mut multipart: Multipart,
-) -> Result<String, StatusCode> {
+) -> Result<String, Response> {
     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();
+    std::fs::create_dir_all(&store_dir)?;
+    Ok(())
+}
 
-    if !store_dir.exists() {
-        std::fs::create_dir_all(&store_dir).unwrap();
-    }
-}
 #[inline]
 fn get_dir() -> PathBuf {
     PathBuf::from("/var/lib/learn_save")