znippy-common 0.9.9

use anyhow::Result;
use arrow_array::{BooleanArray, FixedSizeBinaryArray, StringArray, UInt64Array};
use std::{
    cell::RefCell,
    collections::{HashMap, HashSet},
    fs::{File, OpenOptions},
    os::unix::fs::FileExt,
    path::{Path, PathBuf},
    sync::Arc,
};

use znippy_zoomies::gatling_forkjoin::gatling_for_each_balanced;

use crate::{
    common_config::CONFIG,
    index::{ChunkLoc, VerifyReport},
};

/// Join an archive-relative path onto `out_dir`, rejecting any path that is
/// absolute or escapes `out_dir` via `..`. Returns the normalized destination,
/// guaranteed to stay lexically within `out_dir`.
///
/// The relative path is read straight from an untrusted archive index, so this
/// is the zip-slip / path-traversal gate: a malicious `..` or absolute path
/// must never cause a write outside the extraction root.
fn safe_output_path(out_dir: &Path, rel_path: &str) -> Result<PathBuf> {
    use std::path::Component;
    let mut normalized = PathBuf::new();
    for comp in Path::new(rel_path).components() {
        match comp {
            Component::Normal(c) => normalized.push(c),
            Component::CurDir => {}
            Component::ParentDir => {
                anyhow::bail!("unsafe archive path escapes output dir: {:?}", rel_path)
            }
            Component::RootDir | Component::Prefix(_) => {
                anyhow::bail!("unsafe absolute archive path: {:?}", rel_path)
            }
        }
    }
    anyhow::ensure!(!normalized.as_os_str().is_empty(), "empty archive path");
    Ok(out_dir.join(normalized))
}

/// Fetch a required index column and downcast it, returning an `Err` (never a
/// panic) when the column is absent or has the wrong Arrow type.
///
/// The index schema is read straight from an untrusted archive: an attacker can
/// ship a *valid* Arrow IPC sub-index whose schema is missing a required column
/// (or gives it the wrong type). A bare `column_by_name(..).unwrap()` /
/// `downcast_ref().unwrap()` on that would panic the process (DoS). This gate
/// turns a hostile/corrupt schema into a clean error.
fn required_col<'a, T: 'static>(
    batch: &'a arrow::record_batch::RecordBatch,
    name: &str,
) -> Result<&'a T> {
    let array = batch
        .column_by_name(name)
        .ok_or_else(|| anyhow::anyhow!("archive index missing required column `{name}`"))?;
    array
        .as_any()
        .downcast_ref::<T>()
        .ok_or_else(|| anyhow::anyhow!("archive index column `{name}` has unexpected type"))
}

/// Decompress (and optionally extract) a znippy archive.
///
/// No-coordination read pipeline: every chunk in the index is independent work,
/// so it rides the shared no-barrier **gatling** fork-join pool
/// ([`gatling_for_each_balanced`]) — the same engine the compress path uses, no
/// hand-rolled `thread::spawn` + atomic cursor. Each row is self-dispatched to a
/// worker (heaviest blob first, so a giant chunk never idles siblings), which
/// `pread`s its blob straight from the archive (positioned I/O, no shared seek),
/// decompresses-or-skips into per-worker reusable buffers (thread-local, so the
/// hot loop stays allocation-free — Law 2), blake3-verifies against the per-chunk
/// checksum, and `pwrite`s to its output file at `fdata_offset`. Reads and writes
/// are all positioned, so there is no reader/writer bottleneck and no shared
/// mutable state.
pub fn decompress_archive(
    index_path: &Path,
    save_data: bool,
    out_dir: &Path,
) -> Result<VerifyReport> {
    decompress_archive_filtered(index_path, save_data, out_dir, &crate::index::IndexFilter::default())
}

/// Like [`decompress_archive`] but extracts only the files whose `(pkg_type,
/// repo)` sub-index matches `filter` — a selective extract. An empty filter is
/// equivalent to [`decompress_archive`].
pub fn decompress_archive_filtered(
    index_path: &Path,
    save_data: bool,
    out_dir: &Path,
    filter: &crate::index::IndexFilter,
) -> Result<VerifyReport> {
    let (schema, batches) = crate::index::read_znippy_index_filtered(index_path, filter)?;

    let batch = Arc::new(match batches.len() {
        0 => arrow::record_batch::RecordBatch::new_empty(Arc::new(
            crate::index::ZNIPPY_INDEX_SCHEMA.as_ref().clone(),
        )),
        1 => batches.into_iter().next().unwrap(),
        _ => arrow_select::concat::concat_batches(&schema, batches.iter())
            .map_err(|e| anyhow::anyhow!("failed to merge index batches: {}", e))?,
    });

    let total_rows = batch.num_rows();

    // Downcast every required column once, up front, on the calling thread. The
    // schema comes from an untrusted archive, so a missing/mistyped column must
    // surface as an `Err` here — never an `unwrap` panic — and BEFORE we touch
    // the filesystem creating output files. The gatling closure then reads these
    // by shared borrow (they are `Sync`) — no per-worker re-downcast.
    let paths_col = required_col::<StringArray>(&batch, "relative_path")?;
    let blob_offset_col = required_col::<UInt64Array>(&batch, "blob_offset")?;
    let blob_size_col = required_col::<UInt64Array>(&batch, "blob_size")?;
    let fdata_offset_col = required_col::<UInt64Array>(&batch, "fdata_offset")?;
    let compressed_col = required_col::<BooleanArray>(&batch, "compressed")?;
    let checksum_col = required_col::<FixedSizeBinaryArray>(&batch, "checksum")?;

    let mut unique_files = HashSet::new();
    for i in 0..total_rows {
        unique_files.insert(paths_col.value(i));
    }
    let total_files = unique_files.len();
    drop(unique_files);

    // Pre-create output files, indexed by row for O(1) lock-free lookup in
    // workers. Rows sharing a path share one Arc<File> (pwrite is positioned, so
    // many threads can write disjoint regions of the same file concurrently).
    let output_files: Arc<Vec<Option<Arc<File>>>> = if save_data {
        let mut path_to_file: HashMap<&str, Arc<File>> = HashMap::new();
        let mut created_dirs: HashSet<PathBuf> = HashSet::new();
        let mut files: Vec<Option<Arc<File>>> = Vec::with_capacity(total_rows);
        for i in 0..total_rows {
            let rel_path = paths_col.value(i);
            let f = match path_to_file.get(rel_path) {
                Some(f) => Arc::clone(f),
                None => {
                    // Reject zip-slip / absolute paths before touching the FS —
                    // on any unsafe entry we return Err and extract nothing.
                    let full = safe_output_path(out_dir, rel_path)?;
                    if let Some(parent) = full.parent() {
                        if created_dirs.insert(parent.to_path_buf()) {
                            std::fs::create_dir_all(parent).map_err(|e| {
                                anyhow::anyhow!(
                                    "failed to create output dir {}: {}",
                                    parent.display(),
                                    e
                                )
                            })?;
                        }
                    }
                    let file = Arc::new(
                        OpenOptions::new()
                            .create(true)
                            .write(true)
                            .truncate(true)
                            .open(&full)
                            .map_err(|e| {
                                anyhow::anyhow!(
                                    "failed to open output file {}: {}",
                                    full.display(),
                                    e
                                )
                            })?,
                    );
                    path_to_file.insert(rel_path, Arc::clone(&file));
                    file
                }
            };
            files.push(Some(f));
        }
        Arc::new(files)
    } else {
        Arc::new(vec![])
    };

    let archive = Arc::new(File::open(index_path)?);
    let archive_len = archive.metadata()?.len();
    let num_workers = (CONFIG.max_core_in_flight as usize).max(1);

    // Per-worker reusable buffers (Law 2: no allocation in the hot loop). The
    // gatling closure is a shared `Fn`, so the reuse lives in thread-locals that
    // persist across every row a given worker thread handles — exactly what the
    // old per-thread `read_buf`/`out_buf` did.
    thread_local! {
        static DECOMP_TLS: RefCell<(Vec<u8>, Vec<u8>)> =
            const { RefCell::new((Vec::new(), Vec::new())) };
    }

    // Outcome of decoding one row. `corrupt_row` is `Some(row)` when the chunk
    // failed to decode or its checksum mismatched (surfaced, never a silent hole).
    struct RowStat {
        written: u64,
        verified: u64,
        corrupt: u64,
        corrupt_row: Option<u64>,
    }

    let archive_ref = &archive;
    let output_files_ref = &output_files;

    // One independent unit per row, self-dispatched heaviest-blob-first so a giant
    // chunk is claimed early and never idles siblings. Positioned pread/pwrite mean
    // the units share nothing mutable — pure gatling fork-join.
    let outcomes: Vec<Result<RowStat>> = gatling_for_each_balanced(
        total_rows,
        num_workers,
        /* batch */ 1,
        |row| blob_size_col.value(row),
        |row| -> Result<RowStat> {
            let blob_offset = blob_offset_col.value(row);
            let blob_size = blob_size_col.value(row) as usize;
            let fdata_offset = fdata_offset_col.value(row);
            let compressed = compressed_col.value(row);

            // Bounds-check the attacker-controlled offset/size against the real
            // archive length BEFORE allocating, so a malformed index yields a
            // clean Err instead of a giant zero-fill allocation that aborts the
            // process (DoS).
            if blob_size > 0 {
                let in_bounds = blob_offset
                    .checked_add(blob_size as u64)
                    .is_some_and(|end| end <= archive_len);
                anyhow::ensure!(
                    in_bounds,
                    "blob for row {} out of bounds (offset={}, size={}, archive_len={})",
                    row,
                    blob_offset,
                    blob_size,
                    archive_len
                );
            }

            DECOMP_TLS.with(|cell| -> Result<RowStat> {
                let (read_buf, out_buf) = &mut *cell.borrow_mut();

                // pread the blob into the reusable read buffer (positioned read).
                read_buf.resize(blob_size, 0);
                if blob_size > 0 {
                    archive_ref.read_exact_at(read_buf, blob_offset).map_err(|e| {
                        anyhow::anyhow!("failed to read blob for row {}: {}", row, e)
                    })?;
                }

                // Decompress into out_buf, or use the raw blob bytes (skip path).
                let out: &[u8] = if compressed {
                    match crate::codec::decompress_into(read_buf, out_buf) {
                        Ok(_) => out_buf,
                        Err(e) => {
                            // Count the failed chunk as corrupt so it is surfaced
                            // (never silently dropped, leaving a hole in output).
                            log::error!("[decomp] row {} error={}", row, e);
                            return Ok(RowStat {
                                written: 0,
                                verified: 0,
                                corrupt: 0,
                                corrupt_row: Some(row as u64),
                            });
                        }
                    }
                } else {
                    read_buf
                };

                let len = out.len() as u64;

                // Per-chunk verify: checksum is over the UNCOMPRESSED bytes.
                let computed = blake3::hash(out);
                let verified = &computed.as_bytes()[..] == checksum_col.value(row);
                if !verified {
                    log::error!(
                        "[verify] MISMATCH row={} expected={} got={}",
                        row,
                        hex::encode(checksum_col.value(row)),
                        hex::encode(computed.as_bytes()),
                    );
                    return Ok(RowStat {
                        written: len,
                        verified: 0,
                        corrupt: len,
                        corrupt_row: Some(row as u64),
                    });
                }

                // Only write bytes that verify — never present corrupt output as
                // if it were good. A mismatch is surfaced via corrupt_row above.
                if let Some(Some(file)) = output_files_ref.get(row) {
                    file.write_all_at(out, fdata_offset).map_err(|e| {
                        anyhow::anyhow!("failed to write output for row {}: {}", row, e)
                    })?;
                }

                Ok(RowStat { written: len, verified: len, corrupt: 0, corrupt_row: None })
            })
        },
    );

    let mut total_written_bytes = 0u64;
    let mut verified_bytes = 0u64;
    let mut corrupt_bytes = 0u64;
    let mut corrupt_rows: HashSet<u64> = HashSet::new();
    for outcome in outcomes {
        let st = outcome?;
        total_written_bytes += st.written;
        verified_bytes += st.verified;
        corrupt_bytes += st.corrupt;
        if let Some(r) = st.corrupt_row {
            corrupt_rows.insert(r);
        }
    }
    let total_chunks = total_rows as u64;
    let corrupt_files = corrupt_rows.len();
    let verified_files = total_files.saturating_sub(corrupt_files);

    // ── test-matrix emit ──────────────────────────────────────────────────────
    // Two honest rows for the read path. Reaching here means every self-dispatched
    // gatling unit returned Ok (a hard error is propagated by `outcome?` above and
    // never reaches this point), so the fork-join drain is genuinely green; the
    // blake3 CAS verdict is separate and green ONLY when zero chunks failed the
    // per-chunk checksum. The whole block is stripped in the default build.
    #[cfg(feature = "testmatrix")]
    {
        crate::functional_status(
            "znippy-common/decompress",
            "gatling_for_each_balanced",
            true,
            &format!("workers={num_workers} chunks={total_chunks} files={total_files} drained clean"),
        );
        crate::functional_status(
            "znippy-common/decompress",
            "blake3_cas_verify",
            corrupt_bytes == 0,
            &format!(
                "verified_bytes={verified_bytes} corrupt_bytes={corrupt_bytes} corrupt_files={corrupt_files}"
            ),
        );
    }

    Ok(VerifyReport {
        total_files,
        verified_files,
        corrupt_files,
        total_bytes: total_written_bytes,
        verified_bytes,
        corrupt_bytes,
        chunks: total_chunks,
    })
}

/// Random-access read of a single file from the archive by its relative path.
///
/// Uses the lookup sub-index / trie to find the file's chunks in O(log n)/O(key)
/// (falling back to an index scan for pre-lookup archives), then `pread`s each
/// chunk's blob, decompresses-or-copies it, blake3-verifies against the per-chunk
/// checksum, and reassembles the file by `fdata_offset`. Returns the file bytes,
/// or an error if `target` is not present.
pub fn get_file(archive_path: &Path, target: &str) -> Result<Vec<u8>> {
    let chunks = crate::index::locate_file(archive_path, target)?;
    anyhow::ensure!(!chunks.is_empty(), "file not found in archive: {target}");

    let archive = File::open(archive_path)?;
    let archive_len = archive.metadata()?.len();
    reassemble_file(&archive, archive_len, target, &chunks)
}

/// Read, decompress-or-copy, blake3-verify and reassemble a single file's
/// `chunks` from an already-open archive handle. Shared by the free
/// [`get_file`] (which re-locates per call) and [`crate::index::ArchiveReader`]
/// (which holds one open handle + cached locate across many calls — the
/// selective-restore path). Positioned `pread`s only, so `&File` is enough and
/// the handle can be shared read-only.
pub(crate) fn reassemble_file(
    archive: &File,
    archive_len: u64,
    target: &str,
    chunks: &[ChunkLoc],
) -> Result<Vec<u8>> {
    let total: u64 = chunks.iter().map(|c| c.fdata_offset + c.uncompressed_size).max().unwrap_or(0);
    let mut out = vec![0u8; total as usize];

    let mut read_buf: Vec<u8> = Vec::new();
    let mut dec_buf: Vec<u8> = Vec::new();
    for c in chunks {
        // Bounds-check the attacker-controlled offset/size against the real
        // archive length BEFORE allocating, so a malformed index can't force a
        // giant zero-fill allocation that aborts the process (DoS).
        if c.blob_size > 0 {
            let in_bounds = c
                .blob_offset
                .checked_add(c.blob_size as u64)
                .is_some_and(|end| end <= archive_len);
            anyhow::ensure!(
                in_bounds,
                "blob for {target} out of bounds (offset={}, size={}, archive_len={})",
                c.blob_offset,
                c.blob_size,
                archive_len
            );
        }
        read_buf.resize(c.blob_size as usize, 0);
        if c.blob_size > 0 {
            archive.read_exact_at(&mut read_buf, c.blob_offset)?;
        }
        let bytes: &[u8] = if c.compressed {
            crate::codec::decompress_into(&read_buf, &mut dec_buf)?;
            &dec_buf
        } else {
            &read_buf
        };

        let computed = blake3::hash(bytes);
        anyhow::ensure!(
            computed.as_bytes()[..] == c.checksum[..],
            "checksum mismatch for {target} chunk {}",
            c.chunk_seq
        );

        let start = c.fdata_offset as usize;
        let end = start + bytes.len();
        anyhow::ensure!(end <= out.len(), "chunk overruns file length for {target}");
        out[start..end].copy_from_slice(bytes);
    }
    Ok(out)
}

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

    #[test]
    fn safe_output_path_accepts_normal_relative_paths() {
        let out = Path::new("/tmp/out");
        assert_eq!(
            safe_output_path(out, "a/b/c.txt").unwrap(),
            out.join("a/b/c.txt")
        );
        // `.` components are normalized away but stay inside the root.
        assert_eq!(
            safe_output_path(out, "./a/./b.txt").unwrap(),
            out.join("a/b.txt")
        );
    }

    #[test]
    fn safe_output_path_rejects_zip_slip_and_absolute_paths() {
        let out = Path::new("/tmp/out");
        // `..` traversal must be rejected (zip-slip).
        assert!(safe_output_path(out, "../evil").is_err());
        assert!(safe_output_path(out, "a/../../evil").is_err());
        assert!(safe_output_path(out, "a/b/../../../evil").is_err());
        // Absolute paths must be rejected.
        assert!(safe_output_path(out, "/etc/passwd").is_err());
        // An empty path is rejected.
        assert!(safe_output_path(out, "").is_err());
    }

    #[test]
    fn decompress_corrupt_archive_errors_not_panics() {
        // A malformed `.znippy` (random bytes) must yield a clean Err, never a
        // panic / process abort.
        let dir = tempfile::tempdir().unwrap();
        let archive = dir.path().join("bad.znippy");
        std::fs::write(&archive, b"not a real znippy archive, just junk bytes").unwrap();
        let out_dir = dir.path().join("out");
        let res = decompress_archive(&archive, true, &out_dir);
        assert!(res.is_err(), "corrupt archive should return Err");
    }

    #[test]
    fn decompress_archive_with_wrong_index_schema_errors_not_panics() {
        use arrow::array::StringBuilder;
        use arrow::datatypes::{DataType, Field, Schema};
        use arrow::ipc::writer::StreamWriter;
        use arrow::record_batch::RecordBatch;

        // Craft an archive whose footer + manifest are perfectly well-formed and
        // whose sub-index is *valid Arrow IPC* — but whose schema is missing the
        // columns the decompressor requires (only `relative_path` is present).
        // Reading it succeeds; the per-column downcast must then return a clean
        // Err instead of panicking on `.unwrap()` (attacker-controlled schema).
        let schema = Arc::new(Schema::new(vec![Field::new(
            "relative_path",
            DataType::Utf8,
            false,
        )]));
        let mut b = StringBuilder::new();
        b.append_value("a.txt");
        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(b.finish())]).unwrap();
        let mut sub = Vec::new();
        {
            let mut w = StreamWriter::try_new(&mut sub, &schema).unwrap();
            w.write(&batch).unwrap();
            w.finish().unwrap();
        }

        let manifest = crate::index::write_manifest_bytes(&[crate::index::ManifestEntry {
            pkg_type: 0,
            repo: String::new(),
            module_name: "data".into(),
            index_offset: 0,
            index_len: sub.len() as u64,
            row_count: 1,
        }])
        .unwrap();

        let mut bytes = Vec::new();
        bytes.extend_from_slice(&sub);
        let manifest_offset = bytes.len() as u64;
        bytes.extend_from_slice(&manifest);
        bytes.extend_from_slice(&crate::index::MULTI_INDEX_MAGIC);
        bytes.extend_from_slice(&manifest_offset.to_le_bytes());

        let dir = tempfile::tempdir().unwrap();
        let archive = dir.path().join("wrong_schema.znippy");
        std::fs::write(&archive, &bytes).unwrap();
        let out_dir = dir.path().join("out");
        let res = decompress_archive(&archive, true, &out_dir);
        assert!(
            res.is_err(),
            "archive with a wrong/missing-column index schema must Err, not panic"
        );
    }
}