citygen/server/src/storage.rs

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<u64, String> {
    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<Vec<PathBuf>, 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<RoomSave, String> {
    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);
    }
}