use crate::{city::safe_id, protocol::RoomSave}; use prost::Message; use std::{ fs, io::Write, path::{Path, PathBuf}, }; pub fn save(root: &Path, snapshot: &RoomSave, keep: usize) -> Result<(), String> { save_with_quota(root, snapshot, keep, u64::MAX) } // Serialize quota checks and writes across room writer threads. Never evict the // last snapshot of another room to make space for the current one. static WRITES: std::sync::Mutex<()> = std::sync::Mutex::new(()); pub fn disk_bytes(root: &Path) -> Result { let mut bytes = 0u64; for entry in fs::read_dir(root).map_err(|e| e.to_string())? { let entry = entry.map_err(|e| e.to_string())?; let kind = entry.file_type().map_err(|e| e.to_string())?; if kind.is_symlink() { continue; } bytes = bytes.saturating_add(if kind.is_dir() { disk_bytes(&entry.path())? } else { entry.metadata().map_err(|e| e.to_string())?.len() }); } Ok(bytes) } pub fn save_with_quota( root: &Path, snapshot: &RoomSave, keep: usize, quota: u64, ) -> Result<(), String> { let _guard = WRITES.lock().map_err(|_| "Snapshot lock poisoned")?; if !safe_id(&snapshot.id) { return Err("Invalid room ID".into()); } let folder = root.join(&snapshot.id); fs::create_dir_all(&folder).map_err(|e| e.to_string())?; let bytes = snapshot.encode_to_vec(); if bytes.len() > 128 * 1024 * 1024 { return Err("Snapshot exceeds 128 MiB".into()); } if disk_bytes(root)?.saturating_add(bytes.len() as u64) > quota { return Err("Snapshot disk quota exceeded; previous snapshots preserved".into()); } let name = format!("{:020}.pb", snapshot.tick); let file = folder.join(&name); // Replace through a complete temporary file, including control changes that // happened between two ticks. Never truncate the last committed snapshot. { let temporary = folder.join(format!("{name}.tmp")); let mut output = fs::File::create(&temporary).map_err(|e| e.to_string())?; output.write_all(&bytes).map_err(|e| e.to_string())?; output.sync_all().map_err(|e| e.to_string())?; drop(output); fs::rename(temporary, file).map_err(|e| e.to_string())?; } let mut files = snapshots(&folder)?; files.sort(); let excess = files.len().saturating_sub(keep.max(1)); for path in files.into_iter().take(excess) { fs::remove_file(path).map_err(|e| e.to_string())?; } Ok(()) } fn snapshots(folder: &Path) -> Result, String> { Ok(fs::read_dir(folder) .map_err(|e| e.to_string())? .filter_map(Result::ok) .map(|e| e.path()) .filter(|p| p.extension().is_some_and(|e| e == "pb")) .collect()) } pub fn load(root: &Path, id: &str) -> Result { if !safe_id(id) { return Err("Invalid room ID".into()); } let mut files = snapshots(&root.join(id))?; files.sort(); files.reverse(); // Try an older intact file if the latest is damaged. for file in files { if fs::metadata(&file).map_or(true, |m| m.len() > 128 * 1024 * 1024) { continue; } if let Ok(bytes) = fs::read(file) { if let Ok(snapshot) = RoomSave::decode(&bytes[..]) { if snapshot.id != id { continue; } if ![1, 2, 3].contains(&snapshot.version) || snapshot.city_version != 1 { return Err("Incompatible room snapshot".into()); } return Ok(snapshot); } } } Err("No compatible room snapshot".into()) } #[cfg(test)] mod tests { use super::*; #[test] fn atomic_save_rotation_and_corrupt_fallback() { let dir = tempfile::tempdir().unwrap(); let mut s = RoomSave { version: 1, city_version: 1, id: "r".into(), ..Default::default() }; for tick in 0..5 { s.tick = tick; save(dir.path(), &s, 3).unwrap(); } assert_eq!(snapshots(&dir.path().join("r")).unwrap().len(), 3); assert_eq!(load(dir.path(), "r").unwrap().tick, 4); s.speed = 4; save(dir.path(), &s, 3).unwrap(); assert_eq!(load(dir.path(), "r").unwrap().speed, 4); fs::write(dir.path().join("r/00000000000000000004.pb"), [255, 255]).unwrap(); assert_eq!(load(dir.path(), "r").unwrap().tick, 3); assert!(load(dir.path(), "../secret").is_err()); } #[test] fn quota_preserves_committed_snapshot() { let dir = tempfile::tempdir().unwrap(); let mut snapshot = RoomSave { version: 2, city_version: 1, id: "r".into(), ..Default::default() }; save(dir.path(), &snapshot, 3).unwrap(); let size = disk_bytes(dir.path()).unwrap(); snapshot.tick = 1; assert!(save_with_quota(dir.path(), &snapshot, 3, size).is_err()); assert_eq!(load(dir.path(), "r").unwrap().tick, 0); assert_eq!(disk_bytes(dir.path()).unwrap(), size); } }