znippy-common 0.9.9

//! ZnippyArchive — trait and implementation for reading znippy archives.
//!
//! Provides selective file extraction by path (serve individual artifacts
//! on demand from a single .znippy archive).

use std::collections::HashMap;
use std::fs::File;
use std::os::unix::fs::FileExt;
use std::path::{Path, PathBuf};
use std::sync::{Arc, OnceLock};

use anyhow::{anyhow, Result};
use arrow::record_batch::RecordBatch;
use arrow_array::{BooleanArray, FixedSizeBinaryArray, StringArray, UInt64Array};

use crate::codec;
use crate::index::read_znippy_index;
use crate::views::{
    build_conda_view, build_deb_view, build_gem_view, build_maven_view, build_npm_view,
    build_python_view, build_rpm_view, build_rust_view, CondaView, DebView, GemView, MavenView,
    NpmView, PythonView, RpmView, RustView,
};

/// Trait for reading from a znippy archive.
pub trait ZnippyReader: Send + Sync {
    fn list_files(&self) -> Result<Vec<String>>;
    fn extract_file(&self, relative_path: &str) -> Result<Vec<u8>>;
    fn contains(&self, relative_path: &str) -> bool;
    fn file_size(&self, relative_path: &str) -> Option<u64>;

    /// Batch extract multiple files. Default impl calls extract_file sequentially.
    fn extract_files(&self, paths: &[&str]) -> Vec<Result<Vec<u8>>> {
        paths.iter().map(|p| self.extract_file(p)).collect()
    }
}

struct ChunkInfo {
    blob_offset: u64,
    blob_size: u64,
    fdata_offset: u64,
    compressed: bool,
    /// Per-chunk blake3 over the UNCOMPRESSED bytes (Design law 3). Carried so
    /// the opt-in [`ZnippyArchive::extract_file_verified`] read can integrity-
    /// check without a second index pass. All-zero if the index predates the
    /// checksum column (older archives), in which case verification is skipped.
    checksum: [u8; 32],
}

struct FileEntry {
    uncompressed_size: u64,
    chunks: Vec<ChunkInfo>,
}

/// A znippy archive opened for random-access reads.
/// Loads only the Arrow IPC index on open; blobs are pread on demand. The
/// archive fd is shared (`Arc<File>`) and read via positioned I/O, so
/// `extract_file` is safe to call concurrently from many threads.
pub struct ZnippyArchive {
    archive: Arc<File>,
    /// Real byte length of the archive file, cached at open. Every `pread`
    /// bounds-checks the (index-declared, therefore untrusted) `blob_offset +
    /// blob_size` against this BEFORE allocating, so a corrupt/malicious index
    /// can never force a giant zero-fill allocation (DoS) — mirrors the check
    /// already in `decompress::get_file`.
    archive_len: u64,
    file_index: HashMap<String, FileEntry>,
    /// Archive path — kept so the typed views can do the one-time filtered
    /// sub-index read at view construction.
    path: PathBuf,
    /// Per-`pkg_type` typed view caches. Built once on first `as_*()` call and
    /// reused (the HARD perf contract: repeated `as_maven()` is free). `None`
    /// inside the `Option` means "no sub-index of that type in this archive".
    rust_view: OnceLock<Option<RustView>>,
    maven_view: OnceLock<Option<MavenView>>,
    python_view: OnceLock<Option<PythonView>>,
    npm_view: OnceLock<Option<NpmView>>,
    gem_view: OnceLock<Option<GemView>>,
    conda_view: OnceLock<Option<CondaView>>,
    rpm_view: OnceLock<Option<RpmView>>,
    deb_view: OnceLock<Option<DebView>>,
}

impl ZnippyArchive {
    pub fn open(path: &Path) -> Result<Self> {
        let (_, batches) = read_znippy_index(path)?;
        let file_index = Self::build_file_index(&batches)?;
        let file = File::open(path)?;
        let archive_len = file.metadata()?.len();
        let archive = Arc::new(file);
        Ok(Self {
            archive,
            archive_len,
            file_index,
            path: path.to_path_buf(),
            rust_view: OnceLock::new(),
            maven_view: OnceLock::new(),
            python_view: OnceLock::new(),
            npm_view: OnceLock::new(),
            gem_view: OnceLock::new(),
            conda_view: OnceLock::new(),
            rpm_view: OnceLock::new(),
            deb_view: OnceLock::new(),
        })
    }

    pub fn file_count(&self) -> usize {
        self.file_index.len()
    }

    /// Typed **rust/cargo** view of this archive (coords → crate). Built ONCE on
    /// first call from the rust sub-index and cached; subsequent calls are free.
    /// Returns `None` if the archive has no rust sub-index.
    pub fn as_rust(&self) -> Option<&RustView> {
        self.rust_view
            .get_or_init(|| build_rust_view(&self.path, Arc::clone(&self.archive)).unwrap_or(None))
            .as_ref()
    }

    /// Typed **maven** view (GAV[+classifier] → artifact). Built ONCE and cached.
    /// Returns `None` if the archive has no maven sub-index.
    pub fn as_maven(&self) -> Option<&MavenView> {
        self.maven_view
            .get_or_init(|| build_maven_view(&self.path, Arc::clone(&self.archive)).unwrap_or(None))
            .as_ref()
    }

    /// Typed **python** view (name, version → wheel/sdist). Built ONCE and cached.
    /// Returns `None` if the archive has no python sub-index.
    pub fn as_python(&self) -> Option<&PythonView> {
        self.python_view
            .get_or_init(|| build_python_view(&self.path, Arc::clone(&self.archive)).unwrap_or(None))
            .as_ref()
    }

    /// Typed **npm** view (name[, incl. @scope], version → tarball). Built ONCE
    /// and cached. Returns `None` if the archive has no npm sub-index.
    pub fn as_npm(&self) -> Option<&NpmView> {
        self.npm_view
            .get_or_init(|| build_npm_view(&self.path, Arc::clone(&self.archive)).unwrap_or(None))
            .as_ref()
    }

    /// Typed **gem** view (name, version[, platform] → gem). Built ONCE and
    /// cached. Returns `None` if the archive has no gem sub-index.
    pub fn as_gem(&self) -> Option<&GemView> {
        self.gem_view
            .get_or_init(|| build_gem_view(&self.path, Arc::clone(&self.archive)).unwrap_or(None))
            .as_ref()
    }

    /// Typed **conda** view (name, version[, build, subdir] → package). Built ONCE
    /// and cached. Returns `None` if the archive has no conda sub-index.
    pub fn as_conda(&self) -> Option<&CondaView> {
        self.conda_view
            .get_or_init(|| build_conda_view(&self.path, Arc::clone(&self.archive)).unwrap_or(None))
            .as_ref()
    }

    /// Typed **rpm** view (name, version, release, arch → rpm; authoritative
    /// NEVRA incl. `epoch` from the header). Built ONCE and cached. Returns `None`
    /// if the archive has no rpm sub-index.
    pub fn as_rpm(&self) -> Option<&RpmView> {
        self.rpm_view
            .get_or_init(|| build_rpm_view(&self.path, Arc::clone(&self.archive)).unwrap_or(None))
            .as_ref()
    }

    /// Typed **deb** view (name, version, arch → deb; authoritative coords + the
    /// raw `control` stanza). Built ONCE and cached. Returns `None` if the archive
    /// has no deb sub-index.
    pub fn as_deb(&self) -> Option<&DebView> {
        self.deb_view
            .get_or_init(|| build_deb_view(&self.path, Arc::clone(&self.archive)).unwrap_or(None))
            .as_ref()
    }

    fn build_file_index(batches: &[RecordBatch]) -> Result<HashMap<String, FileEntry>> {
        let mut index: HashMap<String, FileEntry> = HashMap::new();

        for batch in batches {
            let paths = batch
                .column_by_name("relative_path")
                .ok_or_else(|| anyhow!("missing relative_path column"))?
                .as_any()
                .downcast_ref::<StringArray>()
                .ok_or_else(|| anyhow!("relative_path not StringArray"))?;
            let compressed_col = batch
                .column_by_name("compressed")
                .ok_or_else(|| anyhow!("missing compressed column"))?
                .as_any()
                .downcast_ref::<BooleanArray>()
                .ok_or_else(|| anyhow!("compressed not BooleanArray"))?;
            let sizes = batch
                .column_by_name("uncompressed_size")
                .ok_or_else(|| anyhow!("missing uncompressed_size column"))?
                .as_any()
                .downcast_ref::<UInt64Array>()
                .ok_or_else(|| anyhow!("uncompressed_size not UInt64Array"))?;
            let blob_offset_col = batch
                .column_by_name("blob_offset")
                .ok_or_else(|| anyhow!("missing blob_offset column"))?
                .as_any()
                .downcast_ref::<UInt64Array>()
                .ok_or_else(|| anyhow!("blob_offset not UInt64Array"))?;
            let blob_size_col = batch
                .column_by_name("blob_size")
                .ok_or_else(|| anyhow!("missing blob_size column"))?
                .as_any()
                .downcast_ref::<UInt64Array>()
                .ok_or_else(|| anyhow!("blob_size not UInt64Array"))?;
            let fdata_offset_col = batch
                .column_by_name("fdata_offset")
                .ok_or_else(|| anyhow!("missing fdata_offset column"))?
                .as_any()
                .downcast_ref::<UInt64Array>()
                .ok_or_else(|| anyhow!("fdata_offset not UInt64Array"))?;
            // Optional: older archives may lack a valid 32-byte checksum column.
            // When absent (or the wrong width), verification is simply skipped —
            // the fast `extract_file` path never touched it, so this is additive.
            let checksum_col = batch
                .column_by_name("checksum")
                .and_then(|c| c.as_any().downcast_ref::<FixedSizeBinaryArray>())
                .filter(|c| c.value_length() == 32);

            for row in 0..batch.num_rows() {
                let path = paths.value(row).to_string();
                let compressed = compressed_col.value(row);
                let uncompressed_size = sizes.value(row);
                let blob_offset = blob_offset_col.value(row);
                let blob_size = blob_size_col.value(row);
                let fdata_offset = fdata_offset_col.value(row);
                let mut checksum = [0u8; 32];
                if let Some(col) = checksum_col {
                    checksum.copy_from_slice(col.value(row));
                }

                let entry = index.entry(path).or_insert_with(|| FileEntry {
                    uncompressed_size: 0,
                    chunks: Vec::new(),
                });
                entry.uncompressed_size += uncompressed_size;
                entry.chunks.push(ChunkInfo {
                    blob_offset,
                    blob_size,
                    fdata_offset,
                    compressed,
                    checksum,
                });
            }
        }

        for entry in index.values_mut() {
            entry.chunks.sort_by_key(|c| c.fdata_offset);
        }

        Ok(index)
    }

    /// Shared read loop for [`ZnippyReader::extract_file`] (fast, `verify=false`)
    /// and [`Self::extract_file_verified`] (`verify=true`). Every chunk's
    /// index-declared `blob_offset + blob_size` is bounds-checked against the
    /// real archive length before the `resize`, so a corrupt/hostile index can
    /// never force a giant allocation. With `verify` set, each chunk's
    /// reconstructed bytes are blake3-checked against the per-chunk `checksum`.
    fn extract_inner(&self, relative_path: &str, verify: bool) -> Result<Vec<u8>> {
        let entry = self
            .file_index
            .get(relative_path)
            .ok_or_else(|| anyhow!("file not found in archive: {}", relative_path))?;

        let mut result = Vec::with_capacity(entry.uncompressed_size as usize);
        let mut blob = Vec::new(); // reused across chunks
        let mut decomp = Vec::new(); // reused across compressed chunks

        for chunk in &entry.chunks {
            // Bounds-check the untrusted (index-declared) blob extent against the
            // real file length BEFORE allocating — a malformed index must not be
            // able to drive a multi-GB zero-fill (DoS). Matches get_file.
            if chunk.blob_size > 0 {
                let in_bounds = chunk
                    .blob_offset
                    .checked_add(chunk.blob_size)
                    .is_some_and(|end| end <= self.archive_len);
                if !in_bounds {
                    return Err(anyhow!(
                        "blob for {} out of bounds (offset={}, size={}, archive_len={})",
                        relative_path,
                        chunk.blob_offset,
                        chunk.blob_size,
                        self.archive_len
                    ));
                }
            }
            blob.resize(chunk.blob_size as usize, 0);
            // Positioned read — no shared seek, safe under concurrent calls.
            self.archive.read_exact_at(&mut blob, chunk.blob_offset)?;

            let bytes: &[u8] = if chunk.compressed {
                codec::decompress_into(&blob, &mut decomp)?;
                &decomp
            } else {
                &blob
            };

            if verify && chunk.checksum != [0u8; 32] {
                let computed = blake3::hash(bytes);
                if computed.as_bytes()[..] != chunk.checksum[..] {
                    return Err(anyhow!(
                        "checksum mismatch for {} at fdata_offset {}",
                        relative_path,
                        chunk.fdata_offset
                    ));
                }
            }

            result.extend_from_slice(bytes);
        }

        Ok(result)
    }

    /// Random-access read of `relative_path` that **blake3-verifies** every chunk
    /// against the per-chunk checksum in the index before returning (Design law 3).
    ///
    /// This is the integrity-checked counterpart to the fast
    /// [`ZnippyReader::extract_file`], which — by deliberate design, to keep the
    /// artifact-serving hot path allocation-light — does NOT re-hash. Callers that
    /// serve untrusted or long-lived archives (e.g. a registry) should prefer this.
    /// Errors if any chunk's bytes do not match, or if the file is absent. Archives
    /// written before the checksum column skip the hash silently (nothing to check).
    pub fn extract_file_verified(&self, relative_path: &str) -> Result<Vec<u8>> {
        self.extract_inner(relative_path, true)
    }
}

impl ZnippyReader for ZnippyArchive {
    fn list_files(&self) -> Result<Vec<String>> {
        Ok(self.file_index.keys().cloned().collect())
    }

    fn extract_file(&self, relative_path: &str) -> Result<Vec<u8>> {
        self.extract_inner(relative_path, false)
    }

    fn contains(&self, relative_path: &str) -> bool {
        self.file_index.contains_key(relative_path)
    }

    fn file_size(&self, relative_path: &str) -> Option<u64> {
        self.file_index.get(relative_path).map(|e| e.uncompressed_size)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::codec::CompressCtx;
    use crate::index::{build_metadata_batch, lookup_schema};
    use crate::meta::{BlobMeta, ChunkMeta};
    use crate::meta_sink::{ArchiveMetaSink, ArrowIpcSink, GroupKey};
    use std::os::unix::fs::FileExt;
    use std::time::{SystemTime, UNIX_EPOCH};

    fn tmp(tag: &str) -> PathBuf {
        let ns = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos();
        let d = std::env::temp_dir().join(format!("znippy_archive_{tag}_{ns}"));
        std::fs::create_dir_all(&d).unwrap();
        d
    }

    /// Write a one-blob-per-file sealed archive. `size_override` lets a test
    /// declare a `blob_size` in the index that differs from the bytes actually on
    /// disk (used to synthesize a corrupt/hostile index for the bounds-check test).
    fn write_archive(
        path: &Path,
        files: &[(String, Vec<u8>)],
        size_override: Option<u64>,
    ) -> u64 {
        let file = Arc::new(File::create(path).unwrap());
        let mut ctx = CompressCtx::new(3).unwrap();
        let mut blobs = Vec::new();
        let mut paths = Vec::new();
        let mut cursor = 0u64;
        for (fi, (rel, bytes)) in files.iter().enumerate() {
            let checksum = *blake3::hash(bytes).as_bytes();
            let frame = ctx.compress(bytes).unwrap();
            let (on_disk, compressed): (&[u8], bool) =
                if frame.len() < bytes.len() { (&frame, true) } else { (bytes, false) };
            file.write_all_at(on_disk, cursor).unwrap();
            let blob_offset = cursor;
            cursor += on_disk.len() as u64;
            paths.push(rel.clone());
            blobs.push(BlobMeta {
                blob_offset,
                blob_size: size_override.unwrap_or(on_disk.len() as u64),
                chunk_meta: ChunkMeta {
                    fdata_offset: 0,
                    file_index: fi as u64,
                    chunk_seq: 0,
                    checksum,
                    compressed,
                    uncompressed_size: bytes.len() as u64,
                    compressed_size: on_disk.len() as u64,
                },
            });
        }
        let resolver = { let p = paths.clone(); move |fi: u64| p[fi as usize].clone() };
        let batch = build_metadata_batch(&blobs, resolver, &[], &[]).unwrap();
        let schema = lookup_schema();
        let mut sink = ArrowIpcSink::new(file.clone(), cursor);
        sink.push_subindex(
            schema.as_ref(),
            &[batch],
            GroupKey { pkg_type: 0, repo: String::new(), module_name: String::new() },
        )
        .unwrap();
        Box::new(sink).finish().unwrap()
    }

    /// A verified read returns the exact bytes; a blob corrupted on disk is caught
    /// by `extract_file_verified` (blake3 mismatch) while the fast `extract_file`
    /// hands the corrupt bytes back unflagged — proving the verified path adds the
    /// integrity check that the deliberately-fast path omits.
    #[test]
    fn extract_file_verified_catches_corruption_fast_path_does_not() {
        let dir = tmp("verify");
        let archive = dir.join("a.znippy");
        // High-entropy (incompressible) payloads → stored raw, so flipping one
        // byte changes the reconstructed bytes (checksum mismatch) without
        // corrupting a compressed frame the decoder would then reject. A
        // splitmix64 stream gives per-byte pseudo-random fill that zstd/OpenZL
        // cannot shrink, so write_archive keeps it uncompressed.
        let files: Vec<(String, Vec<u8>)> = (0..8)
            .map(|i| {
                let mut s = 0x9e37_79b9_7f4a_7c15u64 ^ (i as u64);
                let body: Vec<u8> = (0..4096u32)
                    .map(|_| {
                        s = s.wrapping_add(0x9e37_79b9_7f4a_7c15);
                        let mut z = s;
                        z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
                        z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
                        (z ^ (z >> 31)) as u8
                    })
                    .collect();
                (format!("repo/f{i:03}.bin"), body)
            })
            .collect();
        write_archive(&archive, &files, None);

        // Clean archive: both paths agree and verification passes.
        let ar = ZnippyArchive::open(&archive).unwrap();
        for (p, bytes) in &files {
            assert_eq!(&ar.extract_file(p).unwrap(), bytes);
            assert_eq!(&ar.extract_file_verified(p).unwrap(), bytes, "clean verify for {p}");
        }
        drop(ar);

        // Flip one byte inside the first blob region (offset 0).
        {
            let f = std::fs::OpenOptions::new().read(true).write(true).open(&archive).unwrap();
            let mut b = [0u8; 1];
            f.read_exact_at(&mut b, 0).unwrap();
            b[0] ^= 0xFF;
            f.write_all_at(&b, 0).unwrap();
        }

        let ar = ZnippyArchive::open(&archive).unwrap();
        let target = &files[0].0;
        // Fast path returns the corrupt bytes without complaint (by design).
        assert_ne!(&ar.extract_file(target).unwrap(), &files[0].1);
        // Verified path rejects them.
        let err = ar.extract_file_verified(target).unwrap_err();
        assert!(
            err.to_string().contains("checksum mismatch"),
            "expected checksum mismatch, got: {err}"
        );

        let _ = std::fs::remove_dir_all(&dir);
    }

    /// A corrupt/hostile index that declares a multi-GB `blob_size` for a tiny
    /// on-disk blob must yield a clean bounds Err — never a giant zero-fill
    /// allocation (DoS). Guards the `checked_add`/`archive_len` check in the read
    /// loop shared by both extract paths.
    #[test]
    fn extract_file_rejects_out_of_bounds_blob_without_huge_alloc() {
        let dir = tmp("bounds");
        let archive = dir.join("b.znippy");
        let files = vec![("repo/small.bin".to_string(), b"hello znippy".to_vec())];
        // Declare an 8 GiB blob_size while writing only a few real bytes.
        write_archive(&archive, &files, Some(8 * 1024 * 1024 * 1024));

        let ar = ZnippyArchive::open(&archive).unwrap();
        let err = ar.extract_file("repo/small.bin").unwrap_err();
        assert!(
            err.to_string().contains("out of bounds"),
            "expected out-of-bounds Err, got: {err}"
        );
        // The verified path funnels through the same guard.
        assert!(ar.extract_file_verified("repo/small.bin").is_err());

        let _ = std::fs::remove_dir_all(&dir);
    }
}