hayson_engine/database/

pack.rs

use std::collections::{BTreeSet, HashMap, HashSet};
use std::fs;
use std::io::{self, Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};

use serde_json::json;
use zip::write::SimpleFileOptions;
use zip::{CompressionMethod, ZipArchive, ZipWriter};

use super::backup::{
    check_relative, emptied_logs, hold, owned_files, parse_manifest, read_catalog_of, refuse, restore_staged, sibling,
    ManifestEntry, RestoreResult, BACKUP_MANIFEST,
};
use super::database::{Database, DatabaseError, LockMode};
use super::layout::LayoutOptions;

const MIB: u64 = 1024 * 1024;
const MAX_NAME_BYTES: usize = 1024;
const MAX_COMPONENT_BYTES: usize = 255;
const DEVICE_NAMES: [&str; 22] = [
    "con", "prn", "aux", "nul", "com1", "com2", "com3", "com4", "com5", "com6", "com7", "com8", "com9", "lpt1", "lpt2",
    "lpt3", "lpt4", "lpt5", "lpt6", "lpt7", "lpt8", "lpt9",
];

/// What a pack may hold and what an unpack accepts. The same limits guard both directions, so a
/// pack this engine writes is one it reads back.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PackLimits {
    /// Size of the archive file itself.
    pub max_archive_bytes: u64,
    /// Sum of the declared sizes of every entry.
    pub max_total_bytes: u64,
    /// Declared size of a single entry.
    pub max_entry_bytes: u64,
    pub max_entries: usize,
}

impl Default for PackLimits {
    fn default() -> Self {
        PackLimits {
            max_archive_bytes: 256 * MIB,
            max_total_bytes: 256 * MIB,
            max_entry_bytes: 128 * MIB,
            max_entries: 10_000,
        }
    }
}

#[derive(Debug, Clone)]
pub struct PackOptions {
    /// How the writer lock is taken for the length of the read. As in [`super::BackupOptions`].
    pub lock: LockMode,
    /// Also pack the `<table>.idx` snapshots (caches; off by default).
    pub indexes: bool,
    /// Tables whose data is left out of the archive and its manifest — see
    /// [`super::BackupOptions::exclude_tables`].
    pub exclude_tables: Vec<String>,
    /// Files of the caller's own to store next to the database files (a manifest of the application's
    /// format, say), as `(name inside the archive, bytes)`. Read back with [`read_pack_entry`].
    pub extra: Vec<(String, Vec<u8>)>,
    pub limits: PackLimits,
}

impl Default for PackOptions {
    fn default() -> Self {
        PackOptions {
            lock: LockMode::default(),
            indexes: false,
            exclude_tables: Vec::new(),
            extra: Vec::new(),
            limits: PackLimits::default(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PackResult {
    pub dest: PathBuf,
    /// Database files in the archive (not counting the manifest and the caller's extra entries).
    pub files: usize,
    /// Their size.
    pub bytes: u64,
    /// Size of the archive file.
    pub archive_bytes: u64,
}

#[derive(Debug, Clone, Default)]
pub struct UnpackOptions {
    pub lock: LockMode,
    /// Refuses a pack holding an image file larger than this, after extraction but before anything
    /// in the target is replaced.
    pub max_image_bytes: Option<u64>,
    pub limits: PackLimits,
}

fn zip_error(error: zip::result::ZipError) -> DatabaseError {
    match error {
        zip::result::ZipError::Io(error) => DatabaseError::Io(error),
        other => refuse(format!("not a readable archive: {other}")),
    }
}

/// An entry name this engine writes or extracts: relative, `/`-separated, no empty, `.` or `..`
/// component, no control character, no Windows device name, no trailing dot or space.
fn check_entry_name(name: &str) -> Result<(), DatabaseError> {
    let bad = |why: &str| refuse(format!("archive entry name {name:?} is refused: {why}"));
    if name.is_empty() || name.len() > MAX_NAME_BYTES {
        return Err(bad("empty or too long"));
    }
    if name.starts_with('/') || name.contains('\\') {
        return Err(bad("absolute path or backslash"));
    }
    if name.chars().any(|c| c.is_control()) {
        return Err(bad("control character"));
    }
    for component in name.split('/') {
        if component.is_empty() || component == "." || component == ".." {
            return Err(bad("empty, \".\" or \"..\" component"));
        }
        if component.len() > MAX_COMPONENT_BYTES {
            return Err(bad("component too long"));
        }
        if component.ends_with('.') || component.ends_with(' ') {
            return Err(bad("component ends with a dot or a space"));
        }
        if component.contains(':') {
            return Err(bad("drive letter or stream marker"));
        }
        let stem = component.split('.').next().unwrap_or("").trim_end().to_lowercase();
        if DEVICE_NAMES.contains(&stem.as_str()) {
            return Err(bad("device name"));
        }
    }
    check_relative(Path::new(name))
}

fn fold(name: &str) -> String {
    name.to_lowercase()
}

/// The result of reading an archive's central directory and judging every entry before any byte is
/// extracted.
struct Inspected {
    archive: ZipArchive<fs::File>,
    /// Non-directory entries: name → index.
    files: HashMap<String, usize>,
    sizes: HashMap<String, u64>,
}

/// How many entries the end-of-central-directory record declares, after making sure that record is
/// the last thing in the file: the zip reader tolerates bytes after it, and collapses entries that
/// share a name, so both are read off the record itself.
fn declared_entries(path: &Path) -> Result<usize, DatabaseError> {
    let mut file = fs::File::open(path)?;
    let length = file.metadata()?.len();
    let window = length.min(22 + u16::MAX as u64);
    file.seek(SeekFrom::Start(length - window))?;
    let mut tail = vec![0u8; window as usize];
    file.read_exact(&mut tail)?;
    for at in (0..tail.len().saturating_sub(21)).rev() {
        if &tail[at..at + 4] != b"PK\x05\x06" {
            continue;
        }
        let comment = u16::from_le_bytes([tail[at + 20], tail[at + 21]]) as usize;
        if at + 22 + comment == tail.len() {
            return Ok(u16::from_le_bytes([tail[at + 10], tail[at + 11]]) as usize);
        }
    }
    Err(refuse(
        "the archive does not end with its end-of-central-directory record — stray bytes follow it, or it is damaged"
            .to_string(),
    ))
}

/// The name the local header of the entry at `header_start` carries.
fn local_name(file: &mut fs::File, header_start: u64) -> Result<Vec<u8>, DatabaseError> {
    let mut fixed = [0u8; 30];
    file.seek(SeekFrom::Start(header_start))?;
    file.read_exact(&mut fixed)
        .map_err(|_| refuse("an archive entry has a truncated local header".to_string()))?;
    if &fixed[..4] != b"PK\x03\x04" {
        return Err(refuse(
            "an archive entry has no local header where the directory says".to_string(),
        ));
    }
    let length = u16::from_le_bytes([fixed[26], fixed[27]]) as usize;
    let mut name = vec![0u8; length];
    file.read_exact(&mut name)
        .map_err(|_| refuse("an archive entry has a truncated local header".to_string()))?;
    Ok(name)
}

fn inspect(path: &Path, limits: &PackLimits) -> Result<Inspected, DatabaseError> {
    let meta = fs::symlink_metadata(path)
        .map_err(|error| refuse(format!("\"{}\" cannot be read: {error}", path.display())))?;
    if !meta.is_file() {
        return Err(refuse(format!("\"{}\" is not a regular file", path.display())));
    }
    if meta.len() > limits.max_archive_bytes {
        return Err(refuse(format!(
            "\"{}\" is {} bytes, over the {} byte limit",
            path.display(),
            meta.len(),
            limits.max_archive_bytes
        )));
    }
    let declared = declared_entries(path)?;
    let mut archive = ZipArchive::new(fs::File::open(path)?).map_err(zip_error)?;
    if archive.len() != declared {
        return Err(refuse(format!(
            "the archive declares {declared} entries but {} can be told apart — two share a name",
            archive.len()
        )));
    }
    let mut headers = fs::File::open(path)?;
    if archive.len() > limits.max_entries {
        return Err(refuse(format!(
            "the archive has {} entries, over the limit of {}",
            archive.len(),
            limits.max_entries
        )));
    }
    let mut files = HashMap::new();
    let mut sizes = HashMap::new();
    let mut seen: HashSet<String> = HashSet::new();
    let mut total = 0u64;
    for index in 0..archive.len() {
        let entry = archive.by_index_raw(index).map_err(zip_error)?;
        let raw_name = entry.name().to_string();
        if std::str::from_utf8(entry.name_raw()) != Ok(raw_name.as_str()) {
            return Err(refuse(format!(
                "archive entry {raw_name:?} has a name that is not UTF-8"
            )));
        }
        if local_name(&mut headers, entry.header_start())? != entry.name_raw() {
            return Err(refuse(format!(
                "archive entry {raw_name:?} is named differently in its local header and in the directory"
            )));
        }
        let is_dir = raw_name.ends_with('/');
        check_entry_name(raw_name.trim_end_matches('/'))?;
        if !seen.insert(fold(raw_name.trim_end_matches('/'))) {
            return Err(refuse(format!(
                "archive entry {raw_name:?} appears twice (names are compared ignoring case)"
            )));
        }
        if entry.encrypted() {
            return Err(refuse(format!("archive entry {raw_name:?} is encrypted")));
        }
        if let Some(mode) = entry.unix_mode() {
            let kind = mode & 0o170000;
            if kind == 0o120000 {
                return Err(refuse(format!("archive entry {raw_name:?} is a symbolic link")));
            }
            if kind != 0 && kind != 0o100000 && kind != 0o040000 {
                return Err(refuse(format!("archive entry {raw_name:?} is not a regular file")));
            }
        }
        if is_dir {
            continue;
        }
        if entry.compression() != CompressionMethod::Stored {
            return Err(refuse(format!(
                "archive entry {raw_name:?} is compressed — packs are stored, never compressed"
            )));
        }
        if entry.size() != entry.compressed_size() {
            return Err(refuse(format!(
                "archive entry {raw_name:?} declares inconsistent sizes"
            )));
        }
        if entry.size() > limits.max_entry_bytes {
            return Err(refuse(format!(
                "archive entry {raw_name:?} declares {} bytes, over the {} byte limit",
                entry.size(),
                limits.max_entry_bytes
            )));
        }
        total = total.saturating_add(entry.size());
        if total > limits.max_total_bytes {
            return Err(refuse(format!(
                "the archive declares more than {} bytes in total",
                limits.max_total_bytes
            )));
        }
        sizes.insert(raw_name.clone(), entry.size());
        files.insert(raw_name, index);
    }
    Ok(Inspected { archive, files, sizes })
}

fn read_entry(inspected: &mut Inspected, name: &str) -> Result<Option<Vec<u8>>, DatabaseError> {
    let Some(&index) = inspected.files.get(name) else {
        return Ok(None);
    };
    let declared = inspected.sizes[name];
    let mut entry = inspected.archive.by_index(index).map_err(zip_error)?;
    let mut bytes = Vec::with_capacity(declared as usize);
    (&mut entry).take(declared + 1).read_to_end(&mut bytes)?;
    if bytes.len() as u64 != declared {
        return Err(refuse(format!("archive entry {name:?} is not the size it declares")));
    }
    Ok(Some(bytes))
}

/// Reads one entry of a pack — the caller's own manifest, say — after the same checks an unpack makes.
/// `None` when the archive has no entry of that name.
pub fn read_pack_entry(archive: &Path, name: &str, limits: &PackLimits) -> Result<Option<Vec<u8>>, DatabaseError> {
    let mut inspected = inspect(archive, limits)?;
    read_entry(&mut inspected, name)
}

/// Packs a database directory into one zip archive at `dest`: the files the database accounts for,
/// stored without compression, plus the `hayson-backup.json` manifest and the caller's `extra` entries.
///
/// `dest` must not exist; the archive is built next to it and renamed into place. The directory is
/// held under the writer lock while it is read, and a database with an unfinished unit of work is
/// refused, as in [`super::backup_database`]. Limits are enforced here too, so what is packed can be
/// unpacked. History kept in a database of its own is not part of this database's files.
pub fn pack_database(dir: &Path, dest: &Path, options: &PackOptions) -> Result<PackResult, DatabaseError> {
    let layout = Database::resolve_layout(dir, &LayoutOptions::default())?;
    let catalog = read_catalog_of(&layout)?;
    if dest.exists() {
        return Err(refuse(format!("\"{}\" already exists", dest.display())));
    }
    let mut taken: BTreeSet<String> = BTreeSet::new();
    taken.insert(fold(BACKUP_MANIFEST));
    for (name, bytes) in &options.extra {
        check_entry_name(name)?;
        if !taken.insert(fold(name)) {
            return Err(refuse(format!("extra entry {name:?} appears twice or is reserved")));
        }
        if bytes.len() as u64 > options.limits.max_entry_bytes {
            return Err(refuse(format!("extra entry {name:?} is over the entry size limit")));
        }
    }
    let partial = sibling(dest, ".partial");
    if partial.exists() {
        fs::remove_file(&partial)?;
    }

    let held = hold(dir, &layout, &options.lock)?;
    let outcome = (|| -> Result<PackResult, DatabaseError> {
        if layout.journal_path().exists() {
            return Err(refuse(format!(
                "\"{}\" holds an unfinished unit of work — open the database for writing once so it is finished, then pack it",
                layout.journal_path().display()
            )));
        }
        let mut files = owned_files(dir, &layout, &catalog, options.indexes, &options.exclude_tables)?;
        let emptied = emptied_logs(dir, &layout, &catalog, &options.exclude_tables);
        files.extend(emptied.iter().cloned());
        files.sort();
        if files.len() + options.extra.len() + 1 > options.limits.max_entries {
            return Err(refuse(
                "the database has more files than the entry limit allows".to_string(),
            ));
        }
        let stored = SimpleFileOptions::default()
            .compression_method(CompressionMethod::Stored)
            .unix_permissions(0o644)
            .large_file(false);
        let mut writer = ZipWriter::new(fs::File::create(&partial)?);
        let mut listed = Vec::new();
        let mut total = 0u64;
        for relative in &files {
            let name = relative.to_string_lossy().replace('\\', "/");
            check_entry_name(&name)?;
            if !taken.insert(fold(&name)) {
                return Err(refuse(format!("database file {name:?} collides with an extra entry")));
            }
            writer.start_file(name.as_str(), stored).map_err(zip_error)?;
            let bytes = if emptied.contains(relative) {
                0
            } else {
                let mut source = fs::File::open(dir.join(relative))?;
                io::copy(&mut (&mut source).take(options.limits.max_entry_bytes + 1), &mut writer)?
            };
            if bytes > options.limits.max_entry_bytes {
                return Err(refuse(format!(
                    "\"{name}\" is over the {} byte entry limit",
                    options.limits.max_entry_bytes
                )));
            }
            total += bytes;
            if total > options.limits.max_total_bytes {
                return Err(refuse(format!(
                    "the database is over the {} byte limit",
                    options.limits.max_total_bytes
                )));
            }
            listed.push(json!({ "path": name, "bytes": bytes }));
        }
        let manifest = json!({
            "format": 1,
            "createdAt": chrono::Utc::now().timestamp_millis(),
            "engine": env!("CARGO_PKG_VERSION"),
            "metadataFile": layout.metadata_file(),
            "files": listed,
        });
        let mut text = serde_json::to_string_pretty(&manifest).map_err(|e| DatabaseError::Catalog(e.to_string()))?;
        text.push('\n');
        writer.start_file(BACKUP_MANIFEST, stored).map_err(zip_error)?;
        writer.write_all(text.as_bytes())?;
        for (name, bytes) in &options.extra {
            writer.start_file(name.as_str(), stored).map_err(zip_error)?;
            writer.write_all(bytes)?;
        }
        let file = writer.finish().map_err(zip_error)?;
        file.sync_all()?;
        let archive_bytes = file.metadata()?.len();
        if archive_bytes > options.limits.max_archive_bytes {
            return Err(refuse(format!(
                "the archive would be {archive_bytes} bytes, over the {} byte limit",
                options.limits.max_archive_bytes
            )));
        }
        Ok(PackResult {
            dest: dest.to_path_buf(),
            files: files.len(),
            bytes: total,
            archive_bytes,
        })
    })();
    drop(held);
    match outcome {
        Ok(result) => {
            fs::rename(&partial, dest)?;
            Ok(result)
        }
        Err(error) => {
            let _ = fs::remove_file(&partial);
            Err(error)
        }
    }
}

/// Unpacks an archive made by [`pack_database`] into `dir`, replacing the database that is there,
/// with the recoverable procedure of [`super::restore_database`] (stage, verify, swap, roll back).
///
/// The archive is judged before anything is written: size and entry-count limits, names (no `..`,
/// absolute or drive paths, backslashes, control characters, device names, duplicates ignoring case),
/// no symbolic links, no compressed or encrypted entries, declared sizes against the limits and against
/// what each entry really holds. Entries the manifest does not list — the caller's extras — are not
/// extracted.
pub fn unpack_database(archive: &Path, dir: &Path, options: &UnpackOptions) -> Result<RestoreResult, DatabaseError> {
    let mut inspected = inspect(archive, &options.limits)?;
    let manifest = read_entry(&mut inspected, BACKUP_MANIFEST)?.ok_or_else(|| {
        refuse(format!(
            "the archive has no \"{BACKUP_MANIFEST}\" — it is not a hayson pack"
        ))
    })?;
    let text = String::from_utf8(manifest).map_err(|_| refuse(format!("\"{BACKUP_MANIFEST}\" is not UTF-8")))?;
    let (metadata_file, entries) = parse_manifest(&text, BACKUP_MANIFEST)?;
    let mut listed: HashSet<String> = HashSet::new();
    for entry in &entries {
        let name = entry.relative.to_string_lossy().replace('\\', "/");
        check_entry_name(&name)?;
        if !listed.insert(fold(&name)) {
            return Err(refuse(format!("the manifest lists {name:?} twice")));
        }
        match inspected.sizes.get(&name) {
            Some(size) if *size == entry.bytes => {}
            Some(_) => {
                return Err(refuse(format!(
                    "archive entry {name:?} does not match its manifest entry ({} bytes expected)",
                    entry.bytes
                )))
            }
            None => return Err(refuse(format!("archive entry {name:?} is missing"))),
        }
    }
    restore_staged(
        dir,
        &metadata_file,
        &entries,
        &options.lock,
        options.max_image_bytes,
        |staged| extract(&mut inspected, &entries, staged),
    )
}

fn extract(inspected: &mut Inspected, entries: &[ManifestEntry], staged: &Path) -> Result<u64, DatabaseError> {
    let mut total = 0u64;
    for entry in entries {
        let name = entry.relative.to_string_lossy().replace('\\', "/");
        let index = inspected.files[&name];
        let target = staged.join(&entry.relative);
        if let Some(parent) = target.parent() {
            fs::create_dir_all(parent)?;
        }
        let mut source = inspected.archive.by_index(index).map_err(zip_error)?;
        let mut out = fs::OpenOptions::new().write(true).create_new(true).open(&target)?;
        let written = io::copy(&mut (&mut source).take(entry.bytes + 1), &mut out)?;
        if written != entry.bytes {
            return Err(refuse(format!("archive entry {name:?} is not the size it declares")));
        }
        out.sync_all()?;
        total += written;
    }
    Ok(total)
}

#[cfg(test)]
mod tests {
    use super::check_entry_name;

    #[test]
    fn refuses_names_that_escape_or_confuse() {
        for bad in [
            "",
            "/etc/passwd",
            "../x",
            "a/../x",
            "a//b",
            "./a",
            "a\\b",
            "C:/x",
            "a\u{0}b",
            "a\nb",
            "con",
            "tables/NUL.txt",
            "a. ",
            "a.",
        ] {
            assert!(check_entry_name(bad).is_err(), "{bad:?} should be refused");
        }
        for good in [
            "hayson.json",
            "tables/users.jsonl",
            "tables/users.images/ab/cd.png",
            "gly.json",
        ] {
            assert!(check_entry_name(good).is_ok(), "{good:?} should be accepted");
        }
    }
}