//! Chunked Upload Service — TUS-like Protocol for Large File Uploads //! //! Enables parallel chunk uploads for files >10 MB with: //! - **Persistent, resumable uploads** — progress survives server restarts. //! A `session.json` (written once on create) + a `progress.bin` bitmask //! (updated atomically on each chunk) are stored alongside the chunk files. //! On boot the service scans `temp_base_dir` and recovers any active sessions. //! - Parallel chunk transfers (up to 6 concurrent) //! - Automatic reassembly with hash-on-write (BLAKE3) //! - Expiration cleanup (24 h) //! //! Protocol: //! 1. POST /api/uploads → Create upload session, get upload_id //! 2. PATCH /api/uploads/:id → Upload chunks (parallel OK) //! 3. HEAD /api/uploads/:id → Check progress //! 4. POST /api/uploads/:id/complete → Finalize and assemble use chrono::{DateTime, Utc}; use dashmap::DashMap; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; use tokio::fs::{self, File}; use tokio::io::AsyncWriteExt; use uuid::Uuid; use crate::application::ports::chunked_upload_ports::{ ChunkUploadResponseDto, ChunkedUploadPort, CreateUploadResponseDto, UploadStatusResponseDto, }; use crate::domain::errors::{DomainError, ErrorKind}; /// Minimum file size to use chunked upload (10 MB) pub const CHUNKED_UPLOAD_THRESHOLD: usize = 10 * 1024 * 1024; /// Default chunk size (5 MB) — optimised for parallel transfers pub const DEFAULT_CHUNK_SIZE: usize = 5 * 1024 * 1024; /// Maximum concurrent chunks per upload pub const MAX_PARALLEL_CHUNKS: usize = 6; /// Upload session expiration time (24 h) const SESSION_EXPIRATION: Duration = Duration::from_secs(24 * 60 * 60); /// Prefix every session directory name with this string so the cleanup /// loop can be safely co-located with unrelated writers (PUT spool /// tempfiles, the NC chunked subtree, anything else a sysadmin places /// under the same `OXICLOUD_CHUNK_DIR`). The orphan-cleanup scan /// filters by this prefix, so non-OxiCloud directories sharing the /// root are never touched. const SESSION_DIR_PREFIX: &str = "oxi-chunk-"; /// Sentinel file names inside each session directory const SESSION_META_FILE: &str = "session.json"; const PROGRESS_FILE: &str = "progress.bin"; /// Build a session directory name from an upload_id by attaching the /// well-known prefix. Symmetric with [`strip_session_prefix`]. fn session_dir_name(upload_id: &str) -> String { format!("{}{}", SESSION_DIR_PREFIX, upload_id) } /// Extract the upload_id from a session directory name. Returns /// `None` when the directory wasn't created by this service (no /// `oxi-chunk-` prefix) — the recovery and cleanup paths use this to /// skip foreign directories cohabiting under `OXICLOUD_CHUNK_DIR`. fn strip_session_prefix(dir_name: &str) -> Option<&str> { dir_name.strip_prefix(SESSION_DIR_PREFIX) } // ─── Serialisable types ────────────────────────────────────────────────────── /// Chunk status #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub enum ChunkStatus { Pending, Uploading, Complete, Failed(String), } /// Individual chunk metadata #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ChunkInfo { pub index: usize, pub offset: u64, pub size: usize, pub status: ChunkStatus, pub checksum: Option, } /// Upload session state — fully serialisable for crash recovery. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct UploadSession { pub id: String, pub user_id: String, pub filename: String, pub folder_id: Option, pub content_type: String, pub total_size: u64, pub chunk_size: usize, pub chunks: Vec, pub created_at: DateTime, pub last_activity: DateTime, pub temp_dir: PathBuf, pub bytes_received: u64, } impl UploadSession { /// Calculate number of chunks needed pub fn calculate_chunk_count(total_size: u64, chunk_size: usize) -> usize { (total_size as usize).div_ceil(chunk_size).max(1) } /// Get upload progress (0.0 – 1.0) pub fn progress(&self) -> f64 { if self.total_size == 0 { return 1.0; } self.bytes_received as f64 / self.total_size as f64 } /// Check if all chunks are complete pub fn is_complete(&self) -> bool { self.chunks .iter() .all(|c| c.status == ChunkStatus::Complete) } /// Get pending chunk indices pub fn pending_chunks(&self) -> Vec { self.chunks .iter() .enumerate() .filter(|(_, c)| c.status == ChunkStatus::Pending) .map(|(i, _)| i) .collect() } /// Check if session has expired pub fn is_expired(&self) -> bool { let elapsed = Utc::now() .signed_duration_since(self.last_activity) .to_std() .unwrap_or(Duration::ZERO); elapsed > SESSION_EXPIRATION } // ── Persistence helpers ────────────────────────────────────────────── /// Build the completed-chunks bitmask (1 bit per chunk). fn build_progress_bitmask(&self) -> Vec { let len = self.chunks.len().div_ceil(8); let mut bitmask = vec![0u8; len]; for chunk in &self.chunks { if chunk.status == ChunkStatus::Complete { bitmask[chunk.index / 8] |= 1 << (chunk.index % 8); } } bitmask } /// Apply a bitmask read from disk, marking matching chunks as `Complete` /// and recalculating `bytes_received`. fn apply_progress_bitmask(&mut self, bitmask: &[u8]) { self.bytes_received = 0; for chunk in &mut self.chunks { let byte_idx = chunk.index / 8; let bit_idx = chunk.index % 8; if byte_idx < bitmask.len() && (bitmask[byte_idx] & (1 << bit_idx)) != 0 { chunk.status = ChunkStatus::Complete; self.bytes_received += chunk.size as u64; } } } /// Persist the full session metadata once (on create). async fn persist_metadata(&self) -> Result<(), String> { let path = self.temp_dir.join(SESSION_META_FILE); let json = serde_json::to_vec(self).map_err(|e| format!("Failed to serialise session: {e}"))?; // Atomic write: write to .tmp then rename let tmp = self.temp_dir.join("session.json.tmp"); fs::write(&tmp, &json) .await .map_err(|e| format!("Failed to write session metadata: {e}"))?; fs::rename(&tmp, &path) .await .map_err(|e| format!("Failed to rename session metadata: {e}"))?; Ok(()) } /// Persist the lightweight progress bitmask (on each chunk upload). async fn persist_progress(&self) -> Result<(), String> { let bitmask = self.build_progress_bitmask(); let path = self.temp_dir.join(PROGRESS_FILE); // Bitmask is tiny (< 512 B for up to 4 096 chunks). // A single write() of < 512 B is atomic on POSIX. fs::write(&path, &bitmask) .await .map_err(|e| format!("Failed to write progress bitmask: {e}"))?; Ok(()) } } // ─── Service ───────────────────────────────────────────────────────────────── /// Chunked Upload Service /// /// Uses `DashMap` (sharded concurrent map) instead of a global `RwLock` /// so that operations on independent upload sessions never contend with each /// other. Disk I/O (temp-dir cleanup) is always performed **outside** any /// map lock to avoid blocking concurrent uploads. pub struct ChunkedUploadService { sessions: Arc>, temp_base_dir: PathBuf, } impl ChunkedUploadService { /// Create the service, recover any persisted sessions, and start the /// background cleanup task. pub async fn new(temp_base_dir: PathBuf) -> Self { // Ensure the base directory exists let _ = fs::create_dir_all(&temp_base_dir).await; // Recover sessions that survived a restart let recovered = Self::recover_sessions(&temp_base_dir).await; let recovered_count = recovered.len(); let service = Self { sessions: Arc::new(DashMap::from_iter(recovered)), temp_base_dir, }; if recovered_count > 0 { tracing::info!("♻️ Recovered {recovered_count} chunked-upload session(s) from disk"); } // Start cleanup task let sessions_clone = service.sessions.clone(); let temp_dir_clone = service.temp_base_dir.clone(); tokio::spawn(async move { Self::cleanup_loop(sessions_clone, temp_dir_clone).await; }); service } /// Lightweight constructor that skips recovery and cleanup. /// Used only by `AppState::default()` (stub wiring). pub fn new_stub(temp_base_dir: PathBuf) -> Self { Self { sessions: Arc::new(DashMap::new()), temp_base_dir, } } // ── Recovery ───────────────────────────────────────────────────────── /// Scan `temp_base_dir` for directories containing `session.json`, /// deserialise each session, apply the `progress.bin` bitmask, and /// verify the chunk files on disk actually exist for completed chunks. async fn recover_sessions(base: &Path) -> HashMap { let mut sessions = HashMap::new(); let mut entries = match fs::read_dir(base).await { Ok(e) => e, Err(_) => return sessions, }; while let Ok(Some(entry)) = entries.next_entry().await { let dir = entry.path(); if !dir.is_dir() { continue; } // Only consider directories WE created — anything without the // `oxi-chunk-` prefix belongs to a sibling writer (NC subtree, // PUT spool tempfiles, sysadmin-placed dirs) and must be left // strictly alone. See `SESSION_DIR_PREFIX`. let dir_name = match dir.file_name().and_then(|n| n.to_str()) { Some(n) => n, None => continue, }; if strip_session_prefix(dir_name).is_none() { continue; } let meta_path = dir.join(SESSION_META_FILE); let meta_bytes = match fs::read(&meta_path).await { Ok(b) => b, Err(_) => continue, // no session.json → orphaned dir, skip }; let mut session: UploadSession = match serde_json::from_slice(&meta_bytes) { Ok(s) => s, Err(e) => { tracing::warn!("Skipping corrupt session in {:?}: {e}", dir); continue; } }; // Apply progress bitmask if present if let Ok(bitmask) = fs::read(dir.join(PROGRESS_FILE)).await { session.apply_progress_bitmask(&bitmask); } else { // No progress file → all chunks still pending (freshly created) session.bytes_received = 0; for chunk in &mut session.chunks { chunk.status = ChunkStatus::Pending; } } // Verify chunk files on disk — downgrade to Pending if missing for chunk in &mut session.chunks { if chunk.status == ChunkStatus::Complete { let chunk_path = dir.join(format!("chunk_{:06}", chunk.index)); if !chunk_path.exists() { tracing::warn!( "Chunk {} missing on disk for session {}, marking pending", chunk.index, session.id ); chunk.status = ChunkStatus::Pending; session.bytes_received = session.bytes_received.saturating_sub(chunk.size as u64); } } } // Skip expired sessions if session.is_expired() { tracing::info!("Skipping expired recovered session: {}", session.id); let _ = fs::remove_dir_all(&dir).await; continue; } tracing::info!( "♻️ Recovered session {} — {}/{} chunks ({:.0}%)", session.id, session .chunks .iter() .filter(|c| c.status == ChunkStatus::Complete) .count(), session.chunks.len(), session.progress() * 100.0 ); sessions.insert(session.id.clone(), session); } sessions } // ── Cleanup ────────────────────────────────────────────────────────── /// Background task to clean expired sessions async fn cleanup_loop(sessions: Arc>, temp_base_dir: PathBuf) { let mut interval = tokio::time::interval(Duration::from_secs(3600)); // Every hour loop { interval.tick().await; // Collect expired session ids + temp dirs (lock-free iteration) let expired: Vec<(String, PathBuf)> = sessions .iter() .filter(|entry| entry.value().is_expired()) .map(|entry| (entry.key().clone(), entry.value().temp_dir.clone())) .collect(); // Remove from map (microseconds per entry) then clean disk OUTSIDE lock for (id, temp_dir) in expired { sessions.remove(&id); if let Err(e) = fs::remove_dir_all(&temp_dir).await { tracing::warn!("Failed to cleanup expired upload {}: {}", id, e); } else { tracing::info!("🧹 Cleaned expired upload session: {}", id); } } // Also clean orphaned temp directories (no session.json or very old). // Filter strictly on the `oxi-chunk-` prefix so we never touch // sibling directories sharing `OXICLOUD_CHUNK_DIR` (NC subtree // `nextcloud/`, PUT spool tempfiles which are files anyway, // operator-placed dirs). Without the prefix filter this loop // would silently delete anything older than 24 h sitting at the // root of the chunked-upload dir. if let Ok(mut entries) = fs::read_dir(&temp_base_dir).await { while let Ok(Some(entry)) = entries.next_entry().await { let path = entry.path(); if !path.is_dir() { continue; } let dir_name = path.file_name().and_then(|n| n.to_str()).unwrap_or(""); let upload_id = match strip_session_prefix(dir_name) { Some(id) => id, None => continue, // not ours — never touch }; if !sessions.contains_key(upload_id) && let Ok(metadata) = fs::metadata(&path).await && let Ok(modified) = metadata.modified() && modified.elapsed().unwrap_or_default() > SESSION_EXPIRATION { let _ = fs::remove_dir_all(&path).await; tracing::info!("🧹 Cleaned orphaned upload dir: {:?}", path); } } } } } // ── Core operations ────────────────────────────────────────────────── /// Verify that the given session belongs to the given user. /// Returns 404 (not 403) to avoid revealing the existence of other users' sessions. fn verify_session_owner(&self, upload_id: &str, user_id: &str) -> Result<(), String> { let session = self .sessions .get(upload_id) .ok_or_else(|| format!("Upload session not found: {}", upload_id))?; if session.user_id != user_id { return Err(format!("Upload session not found: {}", upload_id)); } Ok(()) } /// Create a new upload session (persists `session.json` + empty `progress.bin`) async fn create_session_inner( &self, user_id: String, filename: String, folder_id: Option, content_type: String, total_size: u64, chunk_size: Option, ) -> Result { let upload_id = Uuid::new_v4().to_string(); let chunk_size = chunk_size.unwrap_or(DEFAULT_CHUNK_SIZE); let chunk_count = UploadSession::calculate_chunk_count(total_size, chunk_size); // Create temp directory for chunks. The `oxi-chunk-` prefix // tags the directory as belonging to this service so the // shared-`OXICLOUD_CHUNK_DIR` story holds — see // `SESSION_DIR_PREFIX` for the full rationale. let temp_dir = self.temp_base_dir.join(session_dir_name(&upload_id)); fs::create_dir_all(&temp_dir) .await .map_err(|e| format!("Failed to create temp directory: {e}"))?; // Initialise chunk metadata let mut chunks = Vec::with_capacity(chunk_count); let mut offset: u64 = 0; for i in 0..chunk_count { let size = if i == chunk_count - 1 { (total_size - offset) as usize } else { chunk_size }; chunks.push(ChunkInfo { index: i, offset, size, status: ChunkStatus::Pending, checksum: None, }); offset += size as u64; } let now = Utc::now(); let session = UploadSession { id: upload_id.clone(), user_id, filename, folder_id, content_type, total_size, chunk_size, chunks, created_at: now, last_activity: now, temp_dir, bytes_received: 0, }; let expires_at = SESSION_EXPIRATION.as_secs(); // Persist metadata + empty progress to disk BEFORE inserting into RAM session.persist_metadata().await?; session.persist_progress().await?; self.sessions.insert(upload_id.clone(), session); tracing::info!( "📤 Created chunked upload session: {} ({} chunks, {} bytes each)", upload_id, chunk_count, chunk_size ); Ok(CreateUploadResponseDto { upload_id, chunk_size, total_chunks: chunk_count, expires_at, }) } /// Prepare a chunk write — validates session ownership and chunk /// index, returns the on-disk path the caller should stream the /// HTTP body to plus the expected byte count for that chunk. /// /// Used by the streaming REST PUT path: the handler calls /// `prepare_chunk` → streams body to disk via /// `interfaces::upload_spool::stream_body_to_path` → calls /// `commit_chunk` to finalise. This lets the body bypass the /// in-memory `Bytes` allocation entirely (peak heap ~one HTTP /// frame instead of "chunk size"). /// /// Returns `Err` if the session is unknown, owned by another user, /// the chunk index is out of range, or the chunk is already complete. pub async fn prepare_chunk( &self, upload_id: &str, user_id: Uuid, chunk_index: usize, ) -> Result<(PathBuf, usize), DomainError> { self.verify_session_owner(upload_id, &user_id.to_string()) .map_err(|e| DomainError::new(ErrorKind::NotFound, "ChunkedUpload", e))?; let session = self.sessions.get(upload_id).ok_or_else(|| { DomainError::new( ErrorKind::NotFound, "ChunkedUpload", format!("Upload session not found: {}", upload_id), ) })?; if chunk_index >= session.chunks.len() { return Err(DomainError::new( ErrorKind::InvalidInput, "ChunkedUpload", format!( "Invalid chunk index: {} (max: {})", chunk_index, session.chunks.len() - 1 ), )); } let chunk = &session.chunks[chunk_index]; if chunk.status == ChunkStatus::Complete { return Err(DomainError::new( ErrorKind::InvalidInput, "ChunkedUpload", format!("Chunk {} already uploaded", chunk_index), )); } Ok(( session.temp_dir.join(format!("chunk_{:06}", chunk_index)), chunk.size, )) } /// Finalise a chunk write — verifies the actually-written byte count /// matches the chunk's declared size, validates an optional /// algorithm-tagged checksum, and updates session state. The chunk /// file at `{session.temp_dir}/chunk_{index:06}` must already have /// been written by the caller (typically via /// `stream_body_to_path`). /// /// `actual_size` is the byte count the streaming write reported; /// `computed_checksum` is the hex digest computed during streaming /// (or `None` if the client didn't request a checksum). When /// `expected_checksum` is supplied the two are compared; a /// mismatch removes the partial file and returns `ValidationError` /// so a client retry against the same chunk index gets a clean /// slot. A size mismatch does the same. pub async fn commit_chunk( &self, upload_id: &str, user_id: Uuid, chunk_index: usize, actual_size: u64, computed_checksum: Option, expected_checksum: Option, ) -> Result { self.verify_session_owner(upload_id, &user_id.to_string()) .map_err(|e| DomainError::new(ErrorKind::NotFound, "ChunkedUpload", e))?; // Re-fetch chunk metadata under fresh lock — guards against the // (vanishingly unlikely) case of a session expiry / cancellation // racing with the write. let (chunk_path, expected_size, persist_path) = { let session = self.sessions.get(upload_id).ok_or_else(|| { DomainError::new( ErrorKind::NotFound, "ChunkedUpload", "Session disappeared".to_string(), ) })?; if chunk_index >= session.chunks.len() { return Err(DomainError::new( ErrorKind::InvalidInput, "ChunkedUpload", format!("Invalid chunk index: {}", chunk_index), )); } ( session.temp_dir.join(format!("chunk_{:06}", chunk_index)), session.chunks[chunk_index].size, session.temp_dir.join(PROGRESS_FILE), ) }; // Size check — the streaming body may have been truncated by // the client mid-flight or exceeded the chunk's declared // length. Either way we don't want a partial chunk to count // as complete; nuke it and ask the client to retry. if actual_size != expected_size as u64 { let _ = fs::remove_file(&chunk_path).await; return Err(DomainError::new( ErrorKind::InvalidInput, "ChunkedUpload", format!( "Invalid chunk size: expected {} bytes, got {} bytes", expected_size, actual_size ), )); } // Checksum check — case-insensitive compare so clients that // send uppercase hex still match. if let Some(expected) = expected_checksum.as_ref() && let Some(actual) = computed_checksum.as_ref() && !expected.eq_ignore_ascii_case(actual) { let _ = fs::remove_file(&chunk_path).await; return Err(DomainError::new( ErrorKind::InvalidInput, "ChunkedUpload", format!("Checksum mismatch: expected {}, got {}", expected, actual), )); } // Update session state — DashMap shard lock held only for the // RAM updates (~µs). The bitmask write happens AFTER the ref // is dropped so concurrent uploads to other sessions are never // blocked. Mirrors the legacy `upload_chunk_inner` semantics. let (bytes_received, progress, is_complete, persist_bitmask) = { let mut session = self.sessions.get_mut(upload_id).ok_or_else(|| { DomainError::new( ErrorKind::NotFound, "ChunkedUpload", "Session disappeared".to_string(), ) })?; session.chunks[chunk_index].status = ChunkStatus::Complete; session.chunks[chunk_index].checksum = expected_checksum; session.bytes_received += actual_size; session.last_activity = Utc::now(); let bitmask = session.build_progress_bitmask(); ( session.bytes_received, session.progress(), session.is_complete(), bitmask, ) }; if let Err(e) = fs::write(&persist_path, &persist_bitmask).await { tracing::warn!("Failed to persist progress for {upload_id}: {e}"); } tracing::debug!( "📦 Chunk {} committed for {} ({:.1}% complete)", chunk_index, upload_id, progress * 100.0 ); Ok(ChunkUploadResponseDto { chunk_index, bytes_received, progress, is_complete, }) } /// Upload a single chunk (persists `progress.bin` after success) async fn upload_chunk_inner( &self, upload_id: &str, user_id: &str, chunk_index: usize, data: bytes::Bytes, checksum: Option, ) -> Result { // Verify session exists AND belongs to the requesting user self.verify_session_owner(upload_id, user_id)?; // Validate chunk index is valid let (chunk_path, expected_size) = { let session = self .sessions .get(upload_id) .ok_or_else(|| format!("Upload session not found: {}", upload_id))?; if chunk_index >= session.chunks.len() { return Err(format!( "Invalid chunk index: {} (max: {})", chunk_index, session.chunks.len() - 1 )); } let chunk = &session.chunks[chunk_index]; if chunk.status == ChunkStatus::Complete { return Err(format!("Chunk {} already uploaded", chunk_index)); } ( session.temp_dir.join(format!("chunk_{:06}", chunk_index)), chunk.size, ) }; // Validate chunk size if data.len() != expected_size { return Err(format!( "Invalid chunk size: expected {} bytes, got {} bytes", expected_size, data.len() )); } // Verify checksum if provided — MD5 is CPU-bound (~1.2 ms per 5 MB), // so we offload it to the blocking thread-pool to keep the Tokio // worker free for other connections. if let Some(ref expected_checksum) = checksum { let data_clone = data.clone(); // Bytes::clone is O(1) — just an Arc increment let actual_checksum = tokio::task::spawn_blocking(move || { use md5::{Digest, Md5}; let hash = Md5::digest(&data_clone); hash.iter().map(|b| format!("{b:02x}")).collect::() }) .await .map_err(|e| format!("MD5 checksum task failed: {e}"))?; if actual_checksum != *expected_checksum { return Err(format!( "Checksum mismatch: expected {}, got {}", expected_checksum, actual_checksum )); } } // Write chunk to temp file let mut file = File::create(&chunk_path) .await .map_err(|e| format!("Failed to create chunk file: {e}"))?; file.write_all(&data) .await .map_err(|e| format!("Failed to write chunk: {e}"))?; // Update session state — DashMap shard lock held only for RAM updates (~µs). // Disk I/O (persist_progress) is done AFTER the ref is dropped so // concurrent uploads to other sessions are never blocked. let (bytes_received, progress, is_complete, persist_path, persist_bitmask) = { let mut session = self .sessions .get_mut(upload_id) .ok_or_else(|| "Session disappeared".to_string())?; session.chunks[chunk_index].status = ChunkStatus::Complete; session.chunks[chunk_index].checksum = checksum; session.bytes_received += data.len() as u64; session.last_activity = Utc::now(); // Build bitmask while under lock (CPU-only, ~microseconds) let bitmask = session.build_progress_bitmask(); let path = session.temp_dir.join(PROGRESS_FILE); ( session.bytes_received, session.progress(), session.is_complete(), path, bitmask, ) }; // DashMap shard ref dropped here — held only for RAM updates (~µs) // Persist bitmask to disk OUTSIDE the lock — no longer blocks other uploads if let Err(e) = fs::write(&persist_path, &persist_bitmask).await { tracing::warn!("Failed to persist progress for {upload_id}: {e}"); } tracing::debug!( "📦 Chunk {}/{} uploaded for {} ({:.1}% complete)", chunk_index + 1, expected_size, upload_id, progress * 100.0 ); Ok(ChunkUploadResponseDto { chunk_index, bytes_received, progress, is_complete, }) } /// Get upload status async fn get_status_inner( &self, upload_id: &str, user_id: &str, ) -> Result { self.verify_session_owner(upload_id, user_id)?; let session = self .sessions .get(upload_id) .ok_or_else(|| format!("Upload session not found: {}", upload_id))?; let completed_chunks = session .chunks .iter() .filter(|c| c.status == ChunkStatus::Complete) .count(); Ok(UploadStatusResponseDto { upload_id: session.id.clone(), filename: session.filename.clone(), total_size: session.total_size, bytes_received: session.bytes_received, progress: session.progress(), total_chunks: session.chunks.len(), completed_chunks, pending_chunks: session.pending_chunks(), is_complete: session.is_complete(), }) } /// Assemble chunks into final file and return the path + pre-computed BLAKE3 hash. /// /// **Hash-on-Write**: BLAKE3 is computed while copying chunks into the /// assembled file, eliminating the second sequential read that dedup_service /// would otherwise need. /// /// Returns `(assembled_file_path, filename, folder_id, content_type, total_size, blake3_hash)`. async fn complete_upload_inner( &self, upload_id: &str, user_id: &str, ) -> Result<(PathBuf, String, Option, String, u64, String), String> { // Verify ownership before assembly self.verify_session_owner(upload_id, user_id)?; // Get session and validate completion. // Clone the session data and drop the DashMap ref immediately // so the shard is not held during the expensive assembly step. let session = { let entry = self .sessions .get(upload_id) .ok_or_else(|| format!("Upload session not found: {}", upload_id))?; if !entry.is_complete() { let pending = entry.pending_chunks(); return Err(format!( "Upload not complete. Missing chunks: {:?}", pending )); } entry.clone() }; // Assemble file with hash-on-write. // // The entire loop is offloaded to spawn_blocking because BLAKE3 // hashing is CPU-bound and would otherwise block a Tokio worker, // starving all other connections. // Synchronous I/O is used inside the blocking thread — it avoids // the async reactor overhead and is actually faster for this // sequential workload. let assembled_path = session.temp_dir.join("assembled"); let chunks_meta: Vec<(usize, PathBuf)> = session .chunks .iter() .map(|c| { ( c.index, session.temp_dir.join(format!("chunk_{:06}", c.index)), ) }) .collect(); let total_size = session.total_size; let hash = tokio::task::spawn_blocking(move || -> Result { use std::io::{BufWriter as StdBufWriter, Read, Write}; let raw_output = std::fs::OpenOptions::new() .create(true) .write(true) .truncate(true) .open(&assembled_path) .map_err(|e| format!("Failed to create assembled file: {e}"))?; // Pre-allocate assembled file to reduce fragmentation let _ = raw_output.set_len(total_size); // 512 KB I/O buffers — 8× fewer syscalls than 64 KB let mut output = StdBufWriter::with_capacity(524_288, raw_output); let mut hasher = blake3::Hasher::new(); // For files >10 MB, use multithreaded BLAKE3 hashing (all cores) const RAYON_THRESHOLD: u64 = 10 * 1024 * 1024; let use_rayon = total_size > RAYON_THRESHOLD; // Single 512 KB read buffer reused across all chunks (avoids N allocations) let mut buf = vec![0u8; 524_288]; for (index, chunk_path) in &chunks_meta { let mut chunk_file = std::fs::File::open(chunk_path) .map_err(|e| format!("Failed to open chunk {index}: {e}"))?; loop { let n = chunk_file .read(&mut buf) .map_err(|e| format!("Failed to read chunk {index}: {e}"))?; if n == 0 { break; } if use_rayon { hasher.update_rayon(&buf[..n]); } else { hasher.update(&buf[..n]); } output.write_all(&buf[..n]).map_err(|e| { format!("Failed to write chunk {index} to assembled file: {e}") })?; } } output .flush() .map_err(|e| format!("Failed to flush assembled file: {e}"))?; // ── Durability boundary ──────────────────────────────────── // `flush` drains BufWriter's userspace buffer but leaves the // bytes in the kernel page cache. Without `sync_all`, a // power loss between this `complete_upload` returning 2xx // and the OS writeback timer firing (~5 s default) loses // the merged blob — and PG's metadata row references a hash // that no longer exists on disk. Reclaim the BufWriter's // inner File via `into_inner` so we can `sync_all` it; the // BufWriter would otherwise drop without flushing on the // inner handle. let raw_output = output .into_inner() .map_err(|e| format!("into_inner on BufWriter failed: {e}"))?; raw_output .sync_all() .map_err(|e| format!("Failed to fsync assembled file: {e}"))?; // Clean up chunk files (keep assembled) — already on a blocking thread for (_index, chunk_path) in &chunks_meta { let _ = std::fs::remove_file(chunk_path); } Ok(hasher.finalize().to_hex().to_string()) }) .await .map_err(|e| format!("Assembly task panicked: {e}"))??; let assembled_path = session.temp_dir.join("assembled"); tracing::info!( "✅ Assembled chunked upload: {} ({} bytes from {} chunks)", session.filename, session.total_size, session.chunks.len() ); Ok(( assembled_path, session.filename.clone(), session.folder_id.clone(), session.content_type.clone(), session.total_size, hash, )) } /// Finalize upload: remove session from RAM, then clean disk OUTSIDE lock. async fn finalize_upload_inner(&self, upload_id: &str, user_id: &str) -> Result<(), String> { self.verify_session_owner(upload_id, user_id)?; // Remove from map (~µs) — releases shard immediately let removed = self.sessions.remove(upload_id).map(|(_, s)| s); // Disk I/O happens with NO lock held if let Some(session) = removed && let Err(e) = fs::remove_dir_all(&session.temp_dir).await { tracing::warn!("Failed to cleanup upload {}: {}", upload_id, e); } Ok(()) } /// Cancel an upload and cleanup — disk I/O outside lock. /// /// Returns: /// - `DomainError::NotFound` if no session matches `upload_id` for `user_id` /// (covers both "session missing" and "owned by someone else" — same /// error for anti-enumeration). /// - `DomainError::InternalError` for unexpected disk I/O failures. async fn cancel_upload_inner(&self, upload_id: &str, user_id: &str) -> Result<(), DomainError> { self.verify_session_owner(upload_id, user_id) .map_err(|_| DomainError::not_found("Upload", upload_id))?; // Remove from map (~µs) let removed = self.sessions.remove(upload_id).map(|(_, s)| s); // Disk I/O with NO lock held if let Some(session) = removed { if let Err(e) = fs::remove_dir_all(&session.temp_dir).await { tracing::warn!("Failed to cleanup cancelled upload {}: {}", upload_id, e); } tracing::info!("❌ Cancelled chunked upload: {}", upload_id); } Ok(()) } /// Check if file size qualifies for chunked upload pub fn should_use_chunked(size: u64) -> bool { size as usize >= CHUNKED_UPLOAD_THRESHOLD } /// Get active session count (for monitoring) pub async fn active_sessions(&self) -> usize { self.sessions.len() } } // ─── Port implementation ───────────────────────────────────────────────────── impl ChunkedUploadPort for ChunkedUploadService { async fn create_session( &self, user_id: Uuid, filename: String, folder_id: Option, content_type: String, total_size: u64, chunk_size: Option, ) -> Result { self.create_session_inner( user_id.to_string(), filename, folder_id, content_type, total_size, chunk_size, ) .await .map_err(|e| DomainError::new(ErrorKind::InternalError, "ChunkedUpload", e)) } async fn upload_chunk( &self, upload_id: &str, user_id: Uuid, chunk_index: usize, data: bytes::Bytes, checksum: Option, ) -> Result { self.upload_chunk_inner(upload_id, &user_id.to_string(), chunk_index, data, checksum) .await .map_err(|e| DomainError::new(ErrorKind::InternalError, "ChunkedUpload", e)) } async fn get_status( &self, upload_id: &str, user_id: Uuid, ) -> Result { self.get_status_inner(upload_id, &user_id.to_string()) .await .map_err(|e| DomainError::new(ErrorKind::NotFound, "ChunkedUpload", e)) } async fn complete_upload( &self, upload_id: &str, user_id: Uuid, ) -> Result<(PathBuf, String, Option, String, u64, String), DomainError> { self.complete_upload_inner(upload_id, &user_id.to_string()) .await .map_err(|e| DomainError::new(ErrorKind::InternalError, "ChunkedUpload", e)) } async fn finalize_upload(&self, upload_id: &str, user_id: Uuid) -> Result<(), DomainError> { self.finalize_upload_inner(upload_id, &user_id.to_string()) .await .map_err(|e| DomainError::new(ErrorKind::InternalError, "ChunkedUpload", e)) } async fn cancel_upload(&self, upload_id: &str, user_id: Uuid) -> Result<(), DomainError> { // Inner function now returns DomainError with proper variants // (NotFound for missing/wrong-owner sessions, InternalError otherwise), // so no mapping needed here. self.cancel_upload_inner(upload_id, &user_id.to_string()) .await } fn should_use_chunked(&self, size: u64) -> bool { ChunkedUploadService::should_use_chunked(size) } } // ─── Tests ─────────────────────────────────────────────────────────────────── #[cfg(test)] mod tests { use super::*; #[test] fn test_chunk_count_calculation() { assert_eq!( UploadSession::calculate_chunk_count(10 * 1024 * 1024, 5 * 1024 * 1024), 2 ); assert_eq!( UploadSession::calculate_chunk_count(11 * 1024 * 1024, 5 * 1024 * 1024), 3 ); assert_eq!(UploadSession::calculate_chunk_count(1, 5 * 1024 * 1024), 1); assert_eq!(UploadSession::calculate_chunk_count(0, 5 * 1024 * 1024), 1); } #[test] fn test_should_use_chunked() { assert!(!ChunkedUploadService::should_use_chunked(9 * 1024 * 1024)); assert!(ChunkedUploadService::should_use_chunked(10 * 1024 * 1024)); assert!(ChunkedUploadService::should_use_chunked(100 * 1024 * 1024)); } #[test] fn test_bitmask_roundtrip() { let now = Utc::now(); let mut session = UploadSession { id: "test-id".into(), user_id: "user-1".into(), filename: "file.bin".into(), folder_id: None, content_type: "application/octet-stream".into(), total_size: 15 * 1024 * 1024, chunk_size: 5 * 1024 * 1024, chunks: (0..3) .map(|i| ChunkInfo { index: i, offset: i as u64 * 5 * 1024 * 1024, size: 5 * 1024 * 1024, status: ChunkStatus::Pending, checksum: None, }) .collect(), created_at: now, last_activity: now, temp_dir: PathBuf::from("/tmp/test-id"), bytes_received: 0, }; // Mark chunks 0 and 2 as complete session.chunks[0].status = ChunkStatus::Complete; session.chunks[2].status = ChunkStatus::Complete; session.bytes_received = 10 * 1024 * 1024; let bitmask = session.build_progress_bitmask(); assert_eq!(bitmask, vec![0b00000101]); // bits 0 and 2 set // Reset and re-apply for chunk in &mut session.chunks { chunk.status = ChunkStatus::Pending; } session.bytes_received = 0; session.apply_progress_bitmask(&bitmask); assert_eq!(session.chunks[0].status, ChunkStatus::Complete); assert_eq!(session.chunks[1].status, ChunkStatus::Pending); assert_eq!(session.chunks[2].status, ChunkStatus::Complete); assert_eq!(session.bytes_received, 10 * 1024 * 1024); } #[test] fn test_session_serialisation_roundtrip() { let now = Utc::now(); let session = UploadSession { id: "abc-123".into(), user_id: "user-1".into(), filename: "photo.jpg".into(), folder_id: Some("folder-1".into()), content_type: "image/jpeg".into(), total_size: 1024, chunk_size: 512, chunks: vec![ ChunkInfo { index: 0, offset: 0, size: 512, status: ChunkStatus::Complete, checksum: Some("aabb".into()), }, ChunkInfo { index: 1, offset: 512, size: 512, status: ChunkStatus::Pending, checksum: None, }, ], created_at: now, last_activity: now, temp_dir: PathBuf::from("/tmp/abc-123"), bytes_received: 512, }; let json = serde_json::to_vec(&session).expect("serialise"); let restored: UploadSession = serde_json::from_slice(&json).expect("deserialise"); assert_eq!(restored.id, session.id); assert_eq!(restored.user_id, session.user_id); assert_eq!(restored.filename, session.filename); assert_eq!(restored.folder_id, session.folder_id); assert_eq!(restored.total_size, session.total_size); assert_eq!(restored.chunks.len(), 2); assert_eq!(restored.chunks[0].status, ChunkStatus::Complete); assert_eq!(restored.chunks[1].status, ChunkStatus::Pending); assert_eq!(restored.bytes_received, 512); } #[test] fn test_session_expiry_check() { let mut session = UploadSession { id: "exp-test".into(), user_id: "user-1".into(), filename: "f".into(), folder_id: None, content_type: "x".into(), total_size: 0, chunk_size: 1, chunks: vec![], created_at: Utc::now(), last_activity: Utc::now(), temp_dir: PathBuf::from("/tmp"), bytes_received: 0, }; assert!(!session.is_expired(), "Fresh session must not be expired"); // Simulate old activity session.last_activity = Utc::now() - chrono::Duration::hours(25); assert!(session.is_expired(), "25h-old session must be expired"); } #[tokio::test] async fn test_persist_and_recover_session() { // Use a unique temp dir for this test let base = std::env::temp_dir().join(format!("oxicloud_test_{}", Uuid::new_v4())); let _ = fs::create_dir_all(&base).await; let service = ChunkedUploadService::new(base.clone()).await; // Create a session let resp = service .create_session_inner( "test-user".into(), "bigfile.bin".into(), Some("folder-x".into()), "application/octet-stream".into(), 10 * 1024 * 1024, Some(5 * 1024 * 1024), ) .await .expect("create_session"); let upload_id = resp.upload_id.clone(); // Upload first chunk (5 MB of zeros) let chunk_data = bytes::Bytes::from(vec![0u8; 5 * 1024 * 1024]); service .upload_chunk_inner(&upload_id, "test-user", 0, chunk_data, None) .await .expect("upload_chunk 0"); // Verify files exist on disk let session_dir = base.join(session_dir_name(&upload_id)); assert!(session_dir.join(SESSION_META_FILE).exists()); assert!(session_dir.join(PROGRESS_FILE).exists()); assert!(session_dir.join("chunk_000000").exists()); // Simulate restart: drop service, recover from disk drop(service); let recovered = ChunkedUploadService::recover_sessions(&base).await; assert_eq!(recovered.len(), 1); let session = recovered .get(&upload_id) .expect("session must be recovered"); assert_eq!(session.filename, "bigfile.bin"); assert_eq!(session.folder_id, Some("folder-x".into())); assert_eq!(session.chunks[0].status, ChunkStatus::Complete); assert_eq!(session.chunks[1].status, ChunkStatus::Pending); assert_eq!(session.bytes_received, 5 * 1024 * 1024); // Cleanup let _ = fs::remove_dir_all(&base).await; } #[tokio::test] async fn test_full_upload_lifecycle() { let base = std::env::temp_dir().join(format!("oxicloud_test_{}", Uuid::new_v4())); let service = ChunkedUploadService::new(base.clone()).await; // 1. Create session (1024 bytes, 512 byte chunks → 2 chunks) let resp = service .create_session_inner( "test-user".into(), "test.txt".into(), None, "text/plain".into(), 1024, Some(512), ) .await .expect("create"); assert_eq!(resp.total_chunks, 2); assert_eq!(resp.chunk_size, 512); let id = resp.upload_id; // 2. Upload chunks let chunk0 = bytes::Bytes::from(vec![b'A'; 512]); let r0 = service .upload_chunk_inner(&id, "test-user", 0, chunk0, None) .await .expect("chunk 0"); assert!(!r0.is_complete); let chunk1 = bytes::Bytes::from(vec![b'B'; 512]); let r1 = service .upload_chunk_inner(&id, "test-user", 1, chunk1, None) .await .expect("chunk 1"); assert!(r1.is_complete); assert_eq!(r1.bytes_received, 1024); // 3. Status check let status = service .get_status_inner(&id, "test-user") .await .expect("status"); assert!(status.is_complete); assert_eq!(status.completed_chunks, 2); assert!(status.pending_chunks.is_empty()); // 4. Complete (assemble) let (path, filename, _folder, _ct, size, hash) = service .complete_upload_inner(&id, "test-user") .await .expect("complete"); assert_eq!(filename, "test.txt"); assert_eq!(size, 1024); assert!(!hash.is_empty()); assert!(path.exists()); // 5. Verify assembled content let content = fs::read(&path).await.expect("read assembled"); assert_eq!(&content[..512], &[b'A'; 512]); assert_eq!(&content[512..], &[b'B'; 512]); // 6. Finalize service .finalize_upload_inner(&id, "test-user") .await .expect("finalize"); assert_eq!(service.active_sessions().await, 0); let _ = fs::remove_dir_all(&base).await; } #[tokio::test] async fn test_cancel_removes_files() { let base = std::env::temp_dir().join(format!("oxicloud_test_{}", Uuid::new_v4())); let service = ChunkedUploadService::new(base.clone()).await; let resp = service .create_session_inner( "test-user".into(), "x.bin".into(), None, "application/octet-stream".into(), 512, Some(512), ) .await .expect("create"); let session_dir = base.join(session_dir_name(&resp.upload_id)); assert!(session_dir.exists()); service .cancel_upload_inner(&resp.upload_id, "test-user") .await .expect("cancel"); assert!(!session_dir.exists(), "temp dir must be removed on cancel"); assert_eq!(service.active_sessions().await, 0); let _ = fs::remove_dir_all(&base).await; } #[tokio::test] async fn test_recovery_skips_expired_sessions() { let base = std::env::temp_dir().join(format!("oxicloud_test_{}", Uuid::new_v4())); let _ = fs::create_dir_all(&base).await; // Manually create an expired session on disk. The dir name MUST // carry the `oxi-chunk-` prefix or recovery will (correctly) skip // it as belonging to another writer co-located in chunk_dir. let session_dir = base.join(session_dir_name("expired-session")); let _ = fs::create_dir_all(&session_dir).await; let expired_session = UploadSession { id: "expired-session".into(), user_id: "user-1".into(), filename: "old.bin".into(), folder_id: None, content_type: "application/octet-stream".into(), total_size: 1024, chunk_size: 1024, chunks: vec![ChunkInfo { index: 0, offset: 0, size: 1024, status: ChunkStatus::Pending, checksum: None, }], created_at: Utc::now() - chrono::Duration::hours(48), last_activity: Utc::now() - chrono::Duration::hours(48), temp_dir: session_dir.clone(), bytes_received: 0, }; let json = serde_json::to_vec(&expired_session).unwrap(); fs::write(session_dir.join(SESSION_META_FILE), &json) .await .unwrap(); let recovered = ChunkedUploadService::recover_sessions(&base).await; assert!( recovered.is_empty(), "Expired sessions must not be recovered" ); let _ = fs::remove_dir_all(&base).await; } #[tokio::test] async fn test_recovery_downgrades_missing_chunk_files() { let base = std::env::temp_dir().join(format!("oxicloud_test_{}", Uuid::new_v4())); let _ = fs::create_dir_all(&base).await; let session_dir = base.join(session_dir_name("partial-session")); let _ = fs::create_dir_all(&session_dir).await; let session = UploadSession { id: "partial-session".into(), user_id: "user-1".into(), filename: "file.bin".into(), folder_id: None, content_type: "application/octet-stream".into(), total_size: 1024, chunk_size: 512, chunks: vec![ ChunkInfo { index: 0, offset: 0, size: 512, status: ChunkStatus::Pending, checksum: None, }, ChunkInfo { index: 1, offset: 512, size: 512, status: ChunkStatus::Pending, checksum: None, }, ], created_at: Utc::now(), last_activity: Utc::now(), temp_dir: session_dir.clone(), bytes_received: 0, }; // Write metadata let json = serde_json::to_vec(&session).unwrap(); fs::write(session_dir.join(SESSION_META_FILE), &json) .await .unwrap(); // Write progress marking both chunks complete let bitmask = vec![0b00000011u8]; // bits 0 and 1 fs::write(session_dir.join(PROGRESS_FILE), &bitmask) .await .unwrap(); // But only create chunk_000000 on disk — chunk_000001 is "missing" fs::write(session_dir.join("chunk_000000"), &[0u8; 512]) .await .unwrap(); let recovered = ChunkedUploadService::recover_sessions(&base).await; let s = recovered.get("partial-session").expect("must be recovered"); assert_eq!(s.chunks[0].status, ChunkStatus::Complete); assert_eq!( s.chunks[1].status, ChunkStatus::Pending, "Missing chunk file must be downgraded to Pending" ); assert_eq!(s.bytes_received, 512); let _ = fs::remove_dir_all(&base).await; } }