From 3009f0a82cb401aa6982a803b080221797835f9c Mon Sep 17 00:00:00 2001 From: lilUnkelDemon Date: Thu, 5 Feb 2026 20:56:32 +0330 Subject: [PATCH] feat: add file upload, attachments, config, and cleanup logic --- Cargo.lock | 45 +++++++++++++ Cargo.toml | 3 +- README.md | 124 ++++++++++++++++++++++++++++++++++ src/config.rs | 57 ++++++++++++++++ src/main.rs | 168 +++++++++++++++++++++++++++++++++++++++++----- static/app.js | 64 ++++++++++++++++-- static/index.html | 68 +++++++++++++++---- 7 files changed, 491 insertions(+), 38 deletions(-) create mode 100644 src/config.rs diff --git a/Cargo.lock b/Cargo.lock index bd6e585..734bfae 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -43,6 +43,7 @@ dependencies = [ "matchit", "memchr", "mime", + "multer", "percent-encoding", "pin-project-lite", "serde_core", @@ -187,6 +188,15 @@ version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77c90badedccf4105eca100756a0b1289e191f6fcbdadd3cee1d2f614f97da8f" +[[package]] +name = "encoding_rs" +version = "0.8.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3" +dependencies = [ + "cfg-if", +] + [[package]] name = "errno" version = "0.3.14" @@ -511,6 +521,23 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "multer" +version = "3.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83e87776546dc87511aa5ee218730c92b666d7264ab6ed41f9d215af9cd5224b" +dependencies = [ + "bytes", + "encoding_rs", + "futures-util", + "http", + "httparse", + "memchr", + "mime", + "spin", + "version_check", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -771,6 +798,12 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "spin" +version = "0.9.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" + [[package]] name = "syn" version = "2.0.114" @@ -975,6 +1008,17 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" +[[package]] +name = "uuid" +version = "1.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee48d38b119b0cd71fe4141b30f5ba9c7c5d9f4e7a3a8b4a674e4b6ef789976f" +dependencies = [ + "getrandom", + "js-sys", + "wasm-bindgen", +] + [[package]] name = "version_check" version = "0.9.5" @@ -1201,6 +1245,7 @@ dependencies = [ "serde_json", "tokio", "tower-http", + "uuid", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 08c5f57..30d6f46 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,7 +4,7 @@ version = "0.1.0" edition = "2024" [dependencies] -axum = { version = "0.8", features = ["ws"] } +axum = { version = "0.8", features = ["ws","multipart"] } tokio = { version = "1", features = ["full"] } futures = "0.3" tower-http = { version = "0.5", features = ["fs"] } @@ -12,3 +12,4 @@ dotenv = "0.15.0" serde = { version = "1", features = ["derive"] } serde_json = "1" chrono = "0.4.42" +uuid = { version = "1", features = ["v4"] } diff --git a/README.md b/README.md index 781ea4d..df66bd1 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,72 @@ A **minimal, lightweight WebSocket chat server** written in **Rust** using **Axu --- + +## 📁 File Upload & Attachments + +### 🚀 Upload Endpoint +Files can now be uploaded via HTTP: + +POST /upload + +✨ Details: +- Uses `multipart/form-data` +- Files are stored on disk (default: `static/uploads`) +- Returns file metadata as JSON + +📦 Example response: +{ +"url": "/uploads/uuid-filename.png", +"filename": "photo.png", +"mime": "image/png" +} + +--- + +## 📎 Attachments in Chat Messages + +Chat messages now support: +- 📝 Text only +- 📎 File only +- 📝 + 📎 Text and file together + +⚠️ Empty messages are ignored **unless an attachment exists**. + +🧪 Example WebSocket payload: +{ +"text": "optional text", +"attachment": { +"url": "/uploads/file.png", +"filename": "file.png", +"mime": "image/png" +} +} + +--- + + +## 🧹 Automatic Cleanup of Uploaded Files + +✨ Behavior: +- Uploaded files are linked to chat messages +- When old messages are removed (due to retention limits): + - 🗑️ Their files are automatically deleted from disk +- Prevents unused files from accumulating in `static/uploads` + +--- + +## 🧠 Configurable Message Retention (Updated) + +⚙️ Configuration: +MESSAGE_LIMIT=50 + +📌 Rules: +- MESSAGE_LIMIT > 0 → keep only the last **N** messages +- MESSAGE_LIMIT = 0 → keep **all** messages (no deletion) + +📁 File cleanup follows the same rule. + + ## 📡 Architecture ```text @@ -95,6 +161,18 @@ export TOKEN=super-secret-token ``` +--- + + +## 🔒 MIME Type Whitelisting + +Only explicitly allowed MIME types can be uploaded. + +✅ Example: +ALLOWED_MIMES=image/jpeg,image/png,image/webp,image/gif,application/pdf,application/zip + +❌ Unsupported types are rejected automatically. + --- ## How become an admin @@ -200,3 +278,49 @@ In-memory History (Last 50 messages) * History is sent to **new clients** on connect * Backend handles `join`, `leave`, `message` events + + +## 🧾 MIME Types Reference + +### 🖼 Images +image/jpeg +image/png +image/webp +image/gif +image/bmp +image/heic +image/heif +image/svg+xml ⚠️ + +### 📄 Documents +application/pdf +text/plain +text/csv +text/markdown + +### 📦 Archives +application/zip +application/x-zip-compressed +application/x-7z-compressed +application/gzip +application/x-tar + +### 🎵 Audio +audio/mpeg +audio/wav +audio/ogg +audio/flac +audio/aac + +### 🎥 Video +video/mp4 +video/webm +video/ogg + +### ☠️ Not Recommended / Dangerous +application/x-msdownload +application/x-sh +application/x-bat +application/java-archive + +--- \ No newline at end of file diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..2c950a2 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,57 @@ +use std::{collections::HashSet, env}; + +#[derive(Clone)] +pub struct Config { + pub bind_addr: String, + pub token: String, + + pub upload_dir: String, + pub max_upload_bytes: usize, + pub body_limit_bytes: usize, + pub allowed_mimes: HashSet, + + /// 0 => unlimited (do not delete) + pub message_limit: usize, +} + +fn get_env(name: &str, default: &str) -> String { + env::var(name).unwrap_or_else(|_| default.to_string()) +} + +fn parse_usize_env(name: &str, default: usize) -> usize { + env::var(name) + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(default) +} + +pub fn load_config() -> Config { + let bind_addr = get_env("BIND_ADDR", "127.0.0.1:3000"); + let token = get_env("TOKEN", "token"); + + let upload_dir = get_env("UPLOAD_DIR", "static/uploads"); + let max_upload_bytes = parse_usize_env("MAX_UPLOAD_BYTES", 10 * 1024 * 1024); + let body_limit_bytes = parse_usize_env("BODY_LIMIT_BYTES", 20 * 1024 * 1024); + + let allowed_raw = get_env( + "ALLOWED_MIMES", + "image/jpeg,image/png,image/webp,image/gif,application/pdf,application/zip", + ); + let allowed_mimes: HashSet = allowed_raw + .split(',') + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect(); + + let message_limit = parse_usize_env("MESSAGE_LIMIT", 50); + + Config { + bind_addr, + token, + upload_dir, + max_upload_bytes, + body_limit_bytes, + allowed_mimes, + message_limit, + } +} diff --git a/src/main.rs b/src/main.rs index 2f598ba..718fbb3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,5 @@ +mod config; + use axum::{ extract::{ ws::{Message, WebSocket, WebSocketUpgrade}, @@ -12,17 +14,26 @@ use std::{net::SocketAddr, sync::Arc}; use tokio::sync::broadcast; use tower_http::services::ServeDir; use dotenv::dotenv; - use std::env; use tokio::sync::Mutex; use axum::extract::ConnectInfo; -use serde::{Serialize}; +use serde::{Serialize,Deserialize}; use chrono::Utc; use axum::http::HeaderMap; +use axum::extract::Multipart; +use axum::http::StatusCode; +use uuid::Uuid; +use std::path::Path; +use tokio::fs; +use axum::extract::DefaultBodyLimit; +use config::{Config, load_config}; + + #[derive(Clone)] struct AppState { tx: Arc>, messages: Arc>>, // Store last 50 messages + config: Arc, // Load config file } #[derive(Clone, Serialize)] @@ -37,6 +48,21 @@ struct Member{ ip: String, } + +#[derive(Clone, Serialize, Deserialize)] +struct Attachment { + url: String, + filename: String, + mime: String, +} + +#[derive(Clone, Deserialize)] +struct IncomingMessage { + text: Option, + attachment: Option, +} + + #[derive(Clone, Serialize)] struct ChatMessage { text: String, @@ -44,28 +70,35 @@ struct ChatMessage { username: String, role: Role, r#type: String, // "join" | "leave" | "message" + #[serde(skip_serializing_if = "Option::is_none")] + attachment: Option, } + + + #[tokio::main] async fn main() { dotenv().ok(); - + let config = Arc::new(load_config()); let (tx, _rx) = broadcast::channel::(100); let state = AppState { tx: Arc::new(tx), messages: Arc::new(Mutex::new(Vec::new())), + config: config.clone(), }; let static_files = ServeDir::new("static"); let app = Router::new() .route("/ws", get(ws_handler)) + .route("/upload", axum::routing::post(upload_handler)) .fallback_service(static_files) - .with_state(state); + .layer(DefaultBodyLimit::max(state.config.body_limit_bytes)) + .with_state(state.clone()); - let bind_addr = env::var("BIND_ADDR") - .unwrap_or_else(|_| "127.0.0.1:3000".to_string()); + let bind_addr = state.config.bind_addr.clone(); let addr: SocketAddr = bind_addr .parse() @@ -133,7 +166,7 @@ async fn handle_socket(socket: WebSocket, state: AppState, ip : String) { // Admin token validation if parts.len() == 2 && parts[1] - == env::var("TOKEN").unwrap_or_else(|_| "token".to_string()) + == state.config.token.as_str() { member.role = Role::Admin; member.name = parts[0].to_string(); @@ -155,20 +188,27 @@ async fn handle_socket(socket: WebSocket, state: AppState, ip : String) { username: "system".to_string(), role: member.role.clone(), r#type: "join".into(), + attachment: None, }, ) .await; // Receive user messages - while let Some(Ok(Message::Text(msg))) = receiver.next().await { + while let Some(Ok(Message::Text(raw))) = receiver.next().await { + let (text, attachment) = match serde_json::from_str::(&raw) { + Ok(v) => (v.text.unwrap_or_default(), v.attachment), + Err(_) => (raw.to_string(), None), + }; + broadcast_message( &state, ChatMessage { - text: msg.to_string(), + text, time: Utc::now().to_rfc3339(), username: member.name.clone(), role: member.role.clone(), r#type: "message".into(), + attachment, }, ) .await; @@ -183,6 +223,7 @@ async fn handle_socket(socket: WebSocket, state: AppState, ip : String) { username: "system".to_string(), role: member.role.clone(), r#type: "leave".into(), + attachment: None, }, ) .await; @@ -190,24 +231,119 @@ async fn handle_socket(socket: WebSocket, state: AppState, ip : String) { send_task.abort(); } +fn attachment_disk_path_from_url(url: &str) -> Option { + // Covers both /uploads/xxx and http(s)://.../uploads/xxx + let marker = "/uploads/"; + let idx = url.find(marker)?; + let rel = &url[idx + marker.len()..]; + if rel.is_empty() { + return None; + } + Some(format!("static/uploads/{}", rel)) +} + + // Broadcast messages as JSON ARRAY and keep only last 50 async fn broadcast_message(state: &AppState, msg: ChatMessage) { - - if msg.text.trim().is_empty() { + // dont send msg or file when are empty + if msg.text.trim().is_empty() && msg.attachment.is_none() { return; } - if msg.username != "system" { + let mut to_delete: Vec = vec![]; + + if msg.username != "system" { let mut messages = state.messages.lock().await; messages.push(msg.clone()); - if messages.len() > 50 { - messages.remove(0); + let limit = state.config.message_limit; + + + // Delete last 50 msg + if limit > 0 { + while messages.len() > limit { + let removed = messages.remove(0); + if let Some(att) = removed.attachment { + if let Some(p) = attachment_disk_path_from_url(&att.url) { + to_delete.push(p); + } + } + } + } + + } + + // delete files outs of Lock + for p in to_delete { + match fs::remove_file(&p).await { + Ok(_) => println!("Deleted upload file: {}", p), + Err(e) => eprintln!("Failed to delete upload file {}: {}", p, e), } - // Keep only last 50 messages } - // Always send messages as an array + let json = serde_json::to_string(&vec![msg]).unwrap(); let _ = state.tx.send(json); } + +async fn upload_handler( + State(state):State, + mut multipart: Multipart ) -> impl IntoResponse { + + let allowed = &state.config.allowed_mimes; + let max_bytes = state.config.max_upload_bytes; + let upload_dir = state.config.upload_dir.clone(); + + while let Some(field) = multipart.next_field().await.ok().flatten() { + if field.name() != Some("file") { + continue; + } + + let filename = field.file_name().unwrap_or("file").to_string(); + let mime = field.content_type().unwrap_or("application/octet-stream").to_string(); + + if !allowed.contains(&mime) { + return (StatusCode::UNSUPPORTED_MEDIA_TYPE, "MIME not allowed").into_response(); + } + + let mut data = Vec::new(); + let mut field_stream = field; + + while let Some(chunk) = field_stream.chunk().await.unwrap_or(None) { + data.extend_from_slice(&chunk); + + if data.len() > max_bytes { + return (StatusCode::PAYLOAD_TOO_LARGE, "File too large").into_response(); + } + } + + + let dir = Path::new(&upload_dir); + if fs::create_dir_all(dir).await.is_err() { + return (StatusCode::INTERNAL_SERVER_ERROR, "failed to create upload dir").into_response(); + } + + let ext = Path::new(&filename).extension().and_then(|e| e.to_str()).unwrap_or(""); + + let saved = if ext.is_empty() { + Uuid::new_v4().to_string() + } else { + format!("{}.{}", Uuid::new_v4(), ext) + }; + + let path = dir.join(&saved); + if fs::write(&path, data).await.is_err() { + return (StatusCode::INTERNAL_SERVER_ERROR, "failed to save file").into_response(); + } + + let body = serde_json::json!({ + "url": format!("/uploads/{}", saved), + "filename": filename, + "mime": mime + }); + + return axum::Json(body).into_response(); + } + + (StatusCode::BAD_REQUEST, "No file").into_response() +} diff --git a/static/app.js b/static/app.js index 1961745..33bfacc 100644 --- a/static/app.js +++ b/static/app.js @@ -36,7 +36,7 @@ function getCookie(name) { cookie = cookie.substring(1); } if (cookie.indexOf(cookieName) === 0) { - // استفاده از decodeURIComponent برای رمزگشایی مقدار + // Using decodeURIComponent to decode the value return decodeURIComponent(cookie.substring(cookieName.length, cookie.length)); } } @@ -59,11 +59,11 @@ var app = new Vue({ if (!this.canChat || this.ws == null) { this.handleWebSocket(this.name); } else { - if (this.msg.trim().length == 0) { - return; - } + const t = this.msg.trim(); + if (t.length === 0) return; + e.preventDefault(); - this.ws.send(this.msg.trim()); + this.ws.send(JSON.stringify({ text: t })); this.msg = ""; } }, @@ -73,10 +73,12 @@ var app = new Vue({ return false; } try { - this.ws = new WebSocket(`ws://${location.host}/ws`); + const proto = location.protocol === "https:" ? "wss" : "ws"; + + this.ws = new WebSocket(`${proto}://${location.host}/ws`); this.ws.onopen = () => { - this.ws.send(username); // اولین پیام = username + this.ws.send(username); // first msg = username }; this.ws.onmessage = (e) => { @@ -99,6 +101,54 @@ var app = new Vue({ } }, + async onFilePicked(e) { + const file = e.target.files && e.target.files[0]; + if (!file) return; + + // if not connect first must be connect + if (!this.canChat || !this.ws || this.ws.readyState !== 1) { + this.handleWebSocket(this.name); + try { + await this.waitForWsOpen(); + } catch (err) { + alert("WebSocket connection failed"); + return; + } + } + + try { + const fd = new FormData(); + fd.append("file", file); + + const res = await fetch("/upload", { method: "POST", body: fd }); + if (!res.ok) { + const t = await res.text(); + alert("upload failed: " + t); + return; + } + + const meta = await res.json(); + + // msg with file (msg without file is allowed) + this.ws.send(JSON.stringify({ text: "", attachment: meta })); + } catch (err) { + console.error(err); + alert("upload failed"); + } finally { + e.target.value = ""; + } + }, + waitForWsOpen(timeoutMs = 5000) { + return new Promise((resolve, reject) => { + const start = Date.now(); + const tick = () => { + if (this.ws && this.ws.readyState === 1) return resolve(); + if (Date.now() - start > timeoutMs) return reject(new Error("ws open timeout")); + setTimeout(tick, 50); + }; + tick(); + }); + }, fixTime: function (datetime) { let splited = datetime.split("T"); return splited[0] + " " + splited[1].split('.')[0] diff --git a/static/index.html b/static/index.html index 3ef3b3e..21f7763 100644 --- a/static/index.html +++ b/static/index.html @@ -12,31 +12,71 @@
- - - + + + + + + + + +
    -
  • - - {{msg.username}}: +
  • + + {{m.username}}: -

    - {{msg.text}} + +

    + {{m.text}}

    + + + +
- - +