micialware/learn_save
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Ai improvements and new version system

Browse Source
This commit is contained in:
micialware
2026-08-12 18:04:08 +03:00
parent 8646fb7d3a
commit 15dc54aae3
1 changed files with 135 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/main.rs
+135 -79
View File
@@ -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<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(
State(state): State<Arc<RwLock<ApiState>>>,
Path(id): Path<String>,
body: Body,
) -> Result<impl IntoResponse, StatusCode> {
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::<u64>().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<String>,
mut multipart: Multipart,
) -> Result<impl IntoResponse, StatusCode> {
if !validate_id(id.as_str()) {
return Err(StatusCode::FORBIDDEN);
async fn save_version(
state: Arc<ApiState>,
id: &str,
body: Body,
) -> Result<u64, StatusCode> {
{
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<String>) -> Result<impl IntoResponse, StatusCode> {
if !validate_id(id.as_str()) {
async fn upload_file_stream(
State(state): State<Arc<ApiState>>,
Path(id): Path<String>,
body: Body,
) -> Result<String, StatusCode> {
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<String>) -> Result<impl IntoResponse, StatusCode> {
if !validate_id(id.as_str()) {
async fn upload_file(
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);
}
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() {