use std::collections::{HashMap, HashSet};
use std::io::{Seek, SeekFrom, Write};
use std::path::PathBuf;
use std::sync::{Mutex, RwLock};
use std::time::{Duration, Instant};
use crate::fs::sha256::Hasher;
use crate::fs::UPLOAD_DIR;
pub const DEFAULT_CHUNK_SIZE: usize = 4 * 1024 * 1024;
pub const MAX_CHUNK_SIZE: usize = 8 * 1024 * 1024;
pub const SESSION_TTL: Duration = Duration::from_secs(3600);
const MAX_CONCURRENT_UPLOADS: usize = 128;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum UploadError {
NotFound,
OffsetMismatch { expected: u64 },
Conflict,
TooLarge,
SizeExceeded,
TooManySessions,
Checksum {
expected: String,
actual: String,
dest_rel: String,
},
Io {
detail: String,
raw_os_error: Option<i32>,
},
}
impl From<std::io::Error> for UploadError {
fn from(e: std::io::Error) -> Self {
UploadError::Io {
raw_os_error: e.raw_os_error(),
detail: e.to_string(),
}
}
}
#[derive(Debug, Clone)]
pub struct FinishedUpload {
pub dest_rel: String,
pub part_path: PathBuf,
pub bytes: u64,
pub digest: String,
pub expected: String,
}
struct Session {
dest_rel: String,
part_path: PathBuf,
declared_size: u64,
declared_sha256: String,
offset: u64,
hasher: Hasher,
file: std::fs::File,
touched: Instant,
}
pub struct UploadStore {
sessions: RwLock<HashMap<String, Mutex<Session>>>,
claimed: Mutex<HashSet<String>>,
chunk_size: usize,
counter: std::sync::atomic::AtomicU64,
}
impl UploadStore {
pub fn new(chunk_size: usize) -> Self {
Self {
sessions: RwLock::new(HashMap::new()),
claimed: Mutex::new(HashSet::new()),
chunk_size: chunk_size.clamp(1, MAX_CHUNK_SIZE - 1),
counter: std::sync::atomic::AtomicU64::new(0),
}
}
pub fn chunk_size(&self) -> usize {
self.chunk_size
}
pub fn staging_dir(root: &crate::fs::FsRoot) -> PathBuf {
root.path().join(UPLOAD_DIR)
}
pub fn create(
&self,
root: &crate::fs::FsRoot,
dest_rel: String,
size: u64,
sha256: String,
) -> Result<String, UploadError> {
{
let mut claimed = self.claimed.lock().map_err(|_| poisoned())?;
if claimed.contains(&dest_rel) {
return Err(UploadError::Conflict);
}
if claimed.len() >= MAX_CONCURRENT_UPLOADS {
return Err(UploadError::TooManySessions);
}
claimed.insert(dest_rel.clone());
}
let staging = Self::staging_dir(root);
if let Err(e) = std::fs::create_dir_all(&staging) {
self.release(&dest_rel);
return Err(e.into());
}
let serial = self
.counter
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let id = format!("up-{serial:016x}");
let part_path = staging.join(format!("{id}.part"));
let file = match std::fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(&part_path)
{
Ok(file) => file,
Err(e) => {
self.release(&dest_rel);
return Err(e.into());
}
};
let session = Session {
dest_rel: dest_rel.clone(),
part_path,
declared_size: size,
declared_sha256: sha256,
offset: 0,
hasher: Hasher::new(),
file,
touched: Instant::now(),
};
let mut sessions = match self.sessions.write() {
Ok(sessions) => sessions,
Err(_) => {
std::fs::remove_file(&session.part_path).ok();
self.release(&dest_rel);
return Err(poisoned());
}
};
sessions.insert(id.clone(), Mutex::new(session));
drop(sessions);
if let Ok(mut claimed) = self.claimed.lock() {
claimed.insert(dest_rel);
}
Ok(id)
}
pub fn offset(&self, id: &str) -> Option<u64> {
let sessions = self.sessions.read().ok()?;
let session = sessions.get(id)?.lock().ok()?;
Some(session.offset)
}
pub fn append(&self, id: &str, offset: u64, bytes: &[u8]) -> Result<u64, UploadError> {
if bytes.len() > self.chunk_size {
return Err(UploadError::TooLarge);
}
let sessions = self.sessions.read().map_err(|_| poisoned())?;
let cell = sessions.get(id).ok_or(UploadError::NotFound)?;
let mut session = cell.lock().map_err(|_| poisoned())?;
if offset != session.offset {
return Err(UploadError::OffsetMismatch {
expected: session.offset,
});
}
let next_offset = offset.checked_add(bytes.len() as u64);
if next_offset.map_or(true, |next| next > session.declared_size) {
return Err(UploadError::SizeExceeded);
}
session
.file
.seek(SeekFrom::Start(offset))
.map_err(UploadError::from)?;
session.file.write_all(bytes).map_err(UploadError::from)?;
session.hasher.update(bytes);
session.offset += bytes.len() as u64;
session.touched = Instant::now();
Ok(session.offset)
}
pub fn take_for_complete(&self, id: &str) -> Result<FinishedUpload, UploadError> {
let cell = self
.sessions
.write()
.map_err(|_| poisoned())?
.remove(id)
.ok_or(UploadError::NotFound)?;
let session = cell.into_inner().map_err(|_| poisoned())?;
let Session {
dest_rel,
part_path,
declared_size,
declared_sha256,
offset,
hasher,
file,
..
} = session;
drop(file);
let digest = hasher.finish();
if declared_size != offset || digest != declared_sha256 {
std::fs::remove_file(&part_path).ok();
self.release(&dest_rel);
return Err(UploadError::Checksum {
expected: declared_sha256,
actual: digest,
dest_rel,
});
}
Ok(FinishedUpload {
dest_rel,
part_path,
bytes: offset,
digest: digest.clone(),
expected: declared_sha256,
})
}
pub fn release_destination(&self, dest_rel: &str) {
self.release(dest_rel);
}
pub fn cancel(&self, id: &str) -> Option<(String, u64)> {
let Ok(mut sessions) = self.sessions.write() else {
return None;
};
let cell = sessions.remove(id)?;
drop(sessions);
let Ok(session) = cell.into_inner() else {
return None;
};
self.release(&session.dest_rel);
std::fs::remove_file(&session.part_path).ok();
Some((session.dest_rel, session.offset))
}
pub fn sweep(&self, ttl: Duration) -> Vec<(String, String, u64)> {
let mut expired = Vec::new();
let Ok(sessions) = self.sessions.read() else {
return expired;
};
let stale: Vec<String> = sessions
.iter()
.filter(|(_, cell)| {
cell.lock()
.map(|s| s.touched.elapsed() >= ttl)
.unwrap_or(false)
})
.map(|(id, _)| id.clone())
.collect();
drop(sessions);
for id in stale {
let Ok(mut sessions) = self.sessions.write() else {
break;
};
let Some(cell) = sessions.remove(&id) else {
continue;
};
drop(sessions);
if let Ok(session) = cell.into_inner() {
self.release(&session.dest_rel);
std::fs::remove_file(&session.part_path).ok();
expired.push((id, session.dest_rel, session.offset));
}
}
expired
}
fn release(&self, dest_rel: &str) {
if let Ok(mut claimed) = self.claimed.lock() {
claimed.remove(dest_rel);
}
}
}
impl std::fmt::Debug for UploadStore {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("UploadStore")
.field("chunk_size", &self.chunk_size)
.finish_non_exhaustive()
}
}
fn poisoned() -> UploadError {
UploadError::Io {
detail: "internal lock poisoned".to_string(),
raw_os_error: None,
}
}
pub fn sweep_orphan_parts(root: &crate::fs::FsRoot) -> Vec<(String, u64)> {
let staging = UploadStore::staging_dir(root);
let Ok(entries) = std::fs::read_dir(&staging) else {
return Vec::new();
};
let mut removed = Vec::new();
for entry in entries.flatten() {
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) != Some("part") {
continue;
}
let bytes = std::fs::metadata(&path).map(|m| m.len()).unwrap_or(0);
let Some(id) = path.file_stem().and_then(|s| s.to_str()) else {
std::fs::remove_file(&path).ok();
continue;
};
let id = id.to_string();
if std::fs::remove_file(&path).is_ok() {
removed.push((id, bytes));
}
}
removed
}
#[cfg(test)]
mod tests {
use super::*;
use crate::fs::FsRoot;
fn store() -> (tempfile::TempDir, FsRoot, UploadStore) {
let dir = tempfile::tempdir().expect("tempdir");
let root = FsRoot::new(dir.path()).expect("root");
let store = UploadStore::new(DEFAULT_CHUNK_SIZE);
(dir, root, store)
}
const HELLO_DIGEST: &str = "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9";
#[test]
fn a_session_starts_at_offset_zero() {
let (_dir, root, store) = store();
let id = store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.expect("create");
assert_eq!(store.offset(&id), Some(0));
}
#[test]
fn chunks_advance_the_offset() {
let (_dir, root, store) = store();
let id = store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.expect("create");
assert_eq!(store.append(&id, 0, b"hello ").expect("first"), 6);
assert_eq!(store.append(&id, 6, b"world").expect("second"), 11);
}
#[test]
fn a_chunk_at_the_wrong_offset_is_refused_with_the_expected_one() {
let (_dir, root, store) = store();
let id = store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.expect("create");
store.append(&id, 0, b"hello ").expect("first");
assert_eq!(
store.append(&id, 0, b"again"),
Err(UploadError::OffsetMismatch { expected: 6 })
);
}
#[test]
fn two_sessions_may_not_target_the_same_path() {
let (_dir, root, store) = store();
store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.expect("first");
assert_eq!(
store.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into()),
Err(UploadError::Conflict)
);
}
#[test]
fn a_matching_checksum_completes() {
let (_dir, root, store) = store();
let id = store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.expect("create");
store.append(&id, 0, b"hello world").expect("append");
let finished = store.take_for_complete(&id).expect("complete");
assert_eq!(finished.bytes, 11);
assert_eq!(finished.digest, HELLO_DIGEST);
assert_eq!(finished.dest_rel, "out.bin");
}
#[test]
fn a_mismatched_checksum_is_refused() {
let (_dir, root, store) = store();
let wrong = "0".repeat(64);
let id = store
.create(&root, "out.bin".into(), 11, wrong.clone())
.expect("create");
store.append(&id, 0, b"hello world").expect("append");
match store.take_for_complete(&id) {
Err(UploadError::Checksum {
expected,
actual,
dest_rel,
}) => {
assert_eq!(expected, wrong);
assert_eq!(actual, HELLO_DIGEST);
assert_eq!(dest_rel, "out.bin");
}
other => panic!("expected a checksum refusal, got {other:?}"),
}
assert_eq!(store.offset(&id), None);
}
#[test]
fn a_chunk_above_the_ceiling_is_refused() {
let (_dir, root, store) = store();
let id = store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.expect("create");
let oversized = vec![0_u8; DEFAULT_CHUNK_SIZE + 1];
assert_eq!(store.append(&id, 0, &oversized), Err(UploadError::TooLarge));
}
#[test]
fn a_chunk_that_would_exceed_the_declared_size_is_refused() {
let (_dir, root, store) = store();
let id = store
.create(&root, "out.bin".into(), 5, HELLO_DIGEST.into())
.expect("create");
assert_eq!(
store.append(&id, 0, b"hello world"),
Err(UploadError::SizeExceeded)
);
assert_eq!(store.offset(&id), Some(0));
}
#[test]
fn a_chunk_landing_exactly_on_the_declared_size_is_accepted() {
let (_dir, root, store) = store();
let id = store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.expect("create");
assert_eq!(store.append(&id, 0, b"hello world").expect("append"), 11);
}
#[test]
fn cancelling_removes_the_session_and_frees_the_destination() {
let (_dir, root, store) = store();
let id = store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.expect("create");
store.append(&id, 0, b"hello ").expect("append");
let (destination, bytes) = store.cancel(&id).expect("session existed");
assert_eq!(destination, "out.bin");
assert_eq!(bytes, 6);
assert_eq!(store.offset(&id), None);
assert!(store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.is_ok());
}
#[test]
fn sweeping_drops_sessions_past_their_ttl() {
let (_dir, root, store) = store();
let id = store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.expect("create");
assert_eq!(store.sweep(Duration::ZERO).len(), 1);
assert_eq!(store.offset(&id), None);
}
#[test]
fn create_does_not_sweep_expired_sessions_itself() {
let (_dir, root, store) = store();
let id = store
.create(&root, "old.bin".into(), 11, HELLO_DIGEST.into())
.expect("create");
store
.create(&root, "new.bin".into(), 11, HELLO_DIGEST.into())
.expect("second create");
assert_eq!(
store.offset(&id),
Some(0),
"create must not silently reclaim a stale session; only an explicit sweep call may"
);
}
#[test]
fn a_pre_existing_symlink_at_the_predicted_staging_path_cannot_be_written_through() {
let outer = tempfile::tempdir().expect("outer tempdir");
let root_dir = outer.path().join("root");
std::fs::create_dir_all(&root_dir).expect("mkdir root");
let root = FsRoot::new(&root_dir).expect("root");
let store = UploadStore::new(DEFAULT_CHUNK_SIZE);
let secret = outer.path().join("secret.txt");
std::fs::write(&secret, b"outside-secret").expect("write secret");
let staging = UploadStore::staging_dir(&root);
std::fs::create_dir_all(&staging).expect("mkdir staging");
let predicted = staging.join("up-0000000000000000.part");
#[cfg(unix)]
let linked = std::os::unix::fs::symlink(&secret, &predicted).is_ok();
#[cfg(windows)]
let linked = std::os::windows::fs::symlink_file(&secret, &predicted).is_ok();
#[cfg(not(any(unix, windows)))]
let linked = false;
if !linked {
return; }
let result = store.create(&root, "app-new.bin".into(), 11, HELLO_DIGEST.into());
assert!(
matches!(result, Err(UploadError::Io { .. })),
"create_new must refuse a pre-existing symlink at the staging path \
rather than follow it, got {result:?}"
);
assert_eq!(
std::fs::read(&secret).expect("read secret"),
b"outside-secret",
"the outside target must be untouched: the open must fail before \
any write reaches it"
);
}
#[test]
fn a_cap_limits_concurrent_sessions_and_releasing_one_frees_a_slot() {
let (_dir, root, store) = store();
let mut ids = Vec::with_capacity(MAX_CONCURRENT_UPLOADS);
for i in 0..MAX_CONCURRENT_UPLOADS {
let id = store
.create(&root, format!("f{i}.bin"), 1, HELLO_DIGEST.into())
.unwrap_or_else(|e| panic!("session {i} should fit under the cap: {e:?}"));
ids.push(id);
}
assert_eq!(
store.create(&root, "one-too-many.bin".into(), 1, HELLO_DIGEST.into()),
Err(UploadError::TooManySessions)
);
assert!(store.cancel(&ids[0]).is_some());
assert!(store
.create(&root, "one-too-many.bin".into(), 1, HELLO_DIGEST.into())
.is_ok());
}
#[test]
fn completing_an_upload_keeps_the_destination_claimed_until_explicitly_released() {
let (_dir, root, store) = store();
let id = store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.expect("create");
store.append(&id, 0, b"hello world").expect("append");
let finished = store.take_for_complete(&id).expect("complete");
assert_eq!(
store.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into()),
Err(UploadError::Conflict)
);
store.release_destination(&finished.dest_rel);
assert!(store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.is_ok());
}
#[test]
fn sweep_orphan_parts_removes_leftover_part_files_and_nothing_else() {
let dir = tempfile::tempdir().expect("tempdir");
let root = FsRoot::new(dir.path()).expect("root");
let staging = UploadStore::staging_dir(&root);
std::fs::create_dir_all(&staging).expect("mkdir staging");
std::fs::write(staging.join("up-0000000000000000.part"), b"leftover")
.expect("write orphan");
std::fs::write(staging.join("up-0000000000000001.part"), b"leftover2")
.expect("write second orphan");
std::fs::write(staging.join("keep.txt"), b"not a part file").expect("write keep");
let mut removed = sweep_orphan_parts(&root);
removed.sort();
assert_eq!(
removed,
vec![
("up-0000000000000000".to_string(), 8),
("up-0000000000000001".to_string(), 9),
],
"each orphan must be reported by its id (the filename stem) and the bytes it held, so a caller can audit it"
);
assert!(!staging.join("up-0000000000000000.part").exists());
assert!(!staging.join("up-0000000000000001.part").exists());
assert!(
staging.join("keep.txt").exists(),
"only .part files are orphans; anything else in staging must survive"
);
}
#[test]
fn a_poisoned_sessions_lock_does_not_leak_the_claim_or_the_staging_file() {
let (_dir, root, store) = store();
let poisoned = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let _guard = store.sessions.write().expect("lock not yet poisoned");
panic!("poison it");
}));
assert!(
poisoned.is_err(),
"the closure must have panicked while holding the write guard"
);
let outcome = store.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into());
assert!(
matches!(outcome, Err(UploadError::Io { .. })),
"a poisoned sessions lock must surface as an Io error, got {outcome:?}"
);
store.sessions.clear_poison();
assert!(
store
.create(&root, "out.bin".into(), 11, HELLO_DIGEST.into())
.is_ok(),
"the destination must not still be claimed by the failed attempt"
);
let staging = UploadStore::staging_dir(&root);
let leftover_parts = std::fs::read_dir(&staging)
.expect("staging dir")
.flatten()
.filter(|e| e.path().extension().and_then(|x| x.to_str()) == Some("part"))
.count();
assert_eq!(
leftover_parts, 1,
"only the second, successful session's staging file should remain"
);
}
}