Ai improvements and new version system

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::header::CONTENT_TYPE;
-use axum::http::StatusCode;
+use axum::http::{Method, StatusCode};
-use axum::response::{IntoResponse, Response};
+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<HashSet<String>>,
+    counters: Mutex<HashMap<String, u64>>,
 }
 
 async fn generate_id() -> String {
     Alphanumeric.sample_string(&mut rng(), 24).to_uppercase()
 }
 
-async fn upload_file_stream(
+async fn read_counter(state: &ApiState, id: &str) -> u64 {
-    State(state): State<Arc<RwLock<ApiState>>>,
+    if let Some(&counter) = state.counters.lock().get(id) {
-    Path(id): Path<String>,
+        return counter;
-    body: Body,
-) -> Result<impl IntoResponse, StatusCode> {
-    if !validate_id(id.as_str()) {
-        return Err(StatusCode::BAD_REQUEST);
     }
+    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::<u64>().ok())
+        .unwrap_or(0);
+    state.counters.lock().insert(id.to_string(), counter);
+    counter
+}
 
-    let temp_path;
+async fn write_counter(state: &ApiState, id: &str, value: u64) -> Result<(), StatusCode> {
-    let store_path;
+    let meta_path = state.work_directory.join(format!("{id}.meta"));
-    {
+    let tmp_meta = state.work_directory.join(format!("{id}.meta.tmp"));
-        let state = state.read();
+    tokio::fs::write(&tmp_meta, value.to_string())
-        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)
         .await
         .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
-
+    tokio::fs::rename(&tmp_meta, &meta_path)
-    let stream = TryStreamExt::map_err(body.into_data_stream(), std::io::Error::other);
+        .await
-    let mut reader = StreamReader::new(stream);
-    tokio::io::copy(&mut reader, &mut temp_file).await
         .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
-
+    state.counters.lock().insert(id.to_string(), value);
-    tokio::fs::rename(temp_path, store_path).await.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
-
     Ok(())
 }
 
-async fn upload_file(
+async fn save_version(
-    Path(id): Path<String>,
+    state: Arc<ApiState>,
-    mut multipart: Multipart,
+    id: &str,
-) -> Result<impl IntoResponse, StatusCode> {
+    body: Body,
-    if !validate_id(id.as_str()) {
+) -> Result<u64, StatusCode> {
-        return Err(StatusCode::FORBIDDEN);
+    {
+        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 file_path = state.work_directory.join(id);
-        let data = file
+    let tmp_path = state.work_directory.join(format!("{id}.tmp"));
-            .bytes()
+
+    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();
+        let stream = body.into_data_stream().map_err(std::io::Error::other);
-        path.push(id.clone());
+        let mut reader = StreamReader::new(stream);
-        if !path.exists() {
+        let written = tokio::io::copy(&mut reader, &mut tmp_file)
-            File::create(path.clone())
+            .await
-                .await
+            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
-                .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();
+        write_counter(&state, id, next).await?;
-        Ok(file.metadata().await.unwrap().mtime().to_string())
+        tokio::fs::rename(&tmp_path, &file_path)
-    } else {
+            .await
-        Err(StatusCode::BAD_REQUEST)
+            .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<String>) -> Result<impl IntoResponse, StatusCode> {
+async fn upload_file_stream(
-    if !validate_id(id.as_str()) {
+    State(state): State<Arc<ApiState>>,
+    Path(id): Path<String>,
+    body: Body,
+) -> Result<String, StatusCode> {
+    if !validate_id(&id) {
         return Err(StatusCode::BAD_REQUEST);
     }
-
+    save_version(state, &id, body).await.map(|v| v.to_string())
-    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)
 }
 
-async fn version(Path(id): Path<String>) -> Result<impl IntoResponse, StatusCode> {
+async fn upload_file(
-    if !validate_id(id.as_str()) {
+    State(state): State<Arc<ApiState>>,
+    Path(id): Path<String>,
+    mut multipart: Multipart,
+) -> Result<String, StatusCode> {
+    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<Arc<ApiState>>,
+    Path(id): Path<String>,
+) -> Result<Response, StatusCode> {
+    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<Arc<ApiState>>,
+    Path(id): Path<String>,
+) -> Result<String, StatusCode> {
+    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);
     }
-
+    Ok(read_counter(&state, &id).await.to_string())
-    let dir = get_dir().join(id);
-    let meta = tokio::fs::metadata(dir)
-        .await
-        .map_err(|_| StatusCode::BAD_REQUEST)?;
-
-    Ok(meta.mtime().to_string())
 }
 
 fn ensure_work_dir() {