grit-lib 0.5.21

//! Memory-mapped pack (and MIDX) file bytes with an owned-buffer fallback.
//!
//! This module is the only place that uses `unsafe` for pack backing storage.

use crate::error::{Error, Result};
use crate::objects::HashAlgo;
use memmap2::{Mmap, MmapOptions};
use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::ops::Deref;
use std::path::Path;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

const PACK_SIGNATURE: &[u8; 4] = b"PACK";
const HEADER_LEN: usize = 12;

/// Files at or below this size are read into an owned `Vec` (tiny packs, tests,
/// platforms or filesystems where mapping fails).
const OWNED_READ_THRESHOLD: u64 = 8192;

/// Header plus trailing pack checksum — enough to detect on-disk replacement when
/// mtime/size are unchanged (same-second repack) without re-reading the whole file.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub(crate) struct PackFingerprint {
    /// `PACK`, version, and object count (`12` bytes).
    pub header: [u8; HEADER_LEN],
    /// Trailer width in bytes (20 for SHA-1, 32 for SHA-256).
    pub hash_bytes: u8,
    /// Pack trailer hash (only the first `hash_bytes` entries are meaningful).
    pub trailer: [u8; 32],
}

enum Storage {
    Mapped(Mmap),
    Owned(Vec<u8>),
}

/// Immutable bytes of a `.pack`, pack index, commit-graph, MIDX, or `packed-refs` file.
///
/// # Safety (memory-mapped storage)
///
/// Pack, index, MIDX, and `packed-refs` files are immutable at a stable path: grit
/// and Git publish updates by writing a temporary file and renaming it into place,
/// so active readers never map a file that is being truncated in place. If a pack
/// were truncated while mapped, reads could fault or return stale bytes — the same
/// risk Git accepts for memory-mapped packfiles. `open` uses a single file
/// descriptor and `fstat`s that handle before mapping so length and inode stay paired.
pub struct PackData {
    storage: Storage,
}

impl std::fmt::Debug for PackData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PackData")
            .field("len", &self.len())
            .field(
                "backing",
                &match &self.storage {
                    Storage::Mapped(_) => "mapped",
                    Storage::Owned(_) => "owned",
                },
            )
            .finish()
    }
}

impl Deref for PackData {
    type Target = [u8];

    fn deref(&self) -> &[u8] {
        match &self.storage {
            Storage::Mapped(m) => m.deref(),
            Storage::Owned(v) => v.as_slice(),
        }
    }
}

impl PackData {
    /// Map `path` read-only, or read it into an owned buffer when small or mapping fails.
    ///
    /// # Errors
    ///
    /// Returns [`Error::Io`] when the file cannot be opened or read, and
    /// [`Error::CorruptObject`] when the file is shorter than a pack header.
    pub(crate) fn open(path: &Path) -> Result<Arc<Self>> {
        Self::open_with_min_len(path, HEADER_LEN as u64)
    }

    /// Map or read `path`, requiring at least `min_len` bytes (pack header vs index fanout).
    pub(crate) fn open_with_min_len(path: &Path, min_len: u64) -> Result<Arc<Self>> {
        let file = File::open(path).map_err(Error::Io)?;
        Self::open_from_file(file, min_len)
    }

    fn open_from_file(file: File, min_len: u64) -> Result<Arc<Self>> {
        let len = file.metadata().map_err(Error::Io)?.len();
        if len < min_len {
            return Err(Error::CorruptObject(format!(
                "file is too small ({len} bytes, need at least {min_len})"
            )));
        }
        let len_usize = usize::try_from(len).map_err(|_| {
            Error::CorruptObject(format!(
                "pack file length {len} does not fit in address space"
            ))
        })?;
        if len <= OWNED_READ_THRESHOLD {
            let mut data = Vec::with_capacity(len_usize);
            let mut reader = file;
            reader.read_to_end(&mut data).map_err(Error::Io)?;
            return Ok(Arc::new(Self {
                storage: Storage::Owned(data),
            }));
        }
        match map_read_only(&file, len_usize) {
            Ok(mmap) => Ok(Arc::new(Self {
                storage: Storage::Mapped(mmap),
            })),
            Err(_) => {
                let mut data = Vec::with_capacity(len_usize);
                let mut reader = file;
                reader.seek(SeekFrom::Start(0)).map_err(Error::Io)?;
                reader.read_to_end(&mut data).map_err(Error::Io)?;
                Ok(Arc::new(Self {
                    storage: Storage::Owned(data),
                }))
            }
        }
    }

    /// Wrap an owned buffer (tests, small index files, and stale-byte injection).
    #[must_use]
    pub(crate) fn from_owned(data: Vec<u8>) -> Arc<Self> {
        Arc::new(Self {
            storage: Storage::Owned(data),
        })
    }

    /// Drop this mapping's pages from the process's resident set.
    ///
    /// The bytes stay valid: later reads fault the pages back in from the file (usually
    /// straight from the page cache). Owned buffers and non-Unix platforms are left as is.
    pub(crate) fn release_resident(&self) {
        #[cfg(unix)]
        if let Storage::Mapped(map) = &self.storage {
            // SAFETY: the mapping is a read-only shared file mapping and the file is never
            // modified while mapped (see the type docs), so dropping its pages only makes
            // the next read reload the same bytes from the file; no borrow observes a change.
            let _ = unsafe {
                map.unchecked_advise_range(memmap2::UncheckedAdvice::DontNeed, 0, map.len())
            };
        }
    }

    /// True when backing storage is a memory map (not an owned buffer).
    #[cfg(test)]
    pub(crate) fn is_mapped_backing(&self) -> bool {
        matches!(self.storage, Storage::Mapped(_))
    }

    /// Fingerprint of the bytes currently held in memory.
    ///
    /// `hash_bytes` is the repository object-id / pack-trailer width (from the `.idx` or
    /// repo format), not the pack header version field.
    ///
    /// # Errors
    ///
    /// Returns [`Error::CorruptObject`] when the buffer is not a well-formed pack prefix/trailer.
    pub(crate) fn fingerprint(&self, hash_bytes: usize) -> Result<PackFingerprint> {
        fingerprint_from_slice(self, hash_bytes)
    }
}

/// Caps how much of a memory-mapped pack a sequential pass over it keeps resident.
///
/// Readers [`charge`](Self::charge) the bytes they read; once the charges reach the
/// limit, the whole mapping's pages are dropped and counting starts over. Charging is
/// lock-free and shared by all worker threads. Meant for passes that move forward
/// through the pack: scattered small reads map whole fault-around windows each, so
/// they are better served by positioned file reads than by a budget.
pub(crate) struct ResidentBudget<'a> {
    pack: Option<&'a PackData>,
    limit: usize,
    charged: AtomicUsize,
}

impl<'a> ResidentBudget<'a> {
    /// A budget that never releases (the caller owns the bytes, or they are not mapped).
    #[must_use]
    pub(crate) fn unbounded() -> Self {
        Self {
            pack: None,
            limit: usize::MAX,
            charged: AtomicUsize::new(0),
        }
    }

    /// Keep at most about `limit` bytes of `pack` resident.
    #[must_use]
    pub(crate) fn new(pack: &'a PackData, limit: usize) -> Self {
        Self {
            pack: Some(pack),
            limit,
            charged: AtomicUsize::new(0),
        }
    }

    /// Record a read of `bytes` from the pack, releasing its pages when the limit is reached.
    pub(crate) fn charge(&self, bytes: usize) {
        let Some(pack) = self.pack else {
            return;
        };
        let before = self.charged.fetch_add(bytes, Ordering::Relaxed);
        if before.saturating_add(bytes) < self.limit {
            return;
        }
        // Of several threads crossing the limit together, only the one that resets wins.
        if self.charged.swap(0, Ordering::Relaxed) >= self.limit {
            pack.release_resident();
        }
    }

    /// Checksum `bytes` with `algo`, charging as it goes so a full pass stays within budget.
    pub(crate) fn checksum(&self, algo: HashAlgo, bytes: &[u8]) -> Vec<u8> {
        const CHUNK: usize = 1024 * 1024;
        let mut hasher = algo.checksum_hasher();
        for chunk in bytes.chunks(CHUNK) {
            hasher.update(chunk);
            self.charge(chunk.len());
        }
        hasher.finalize()
    }
}

fn map_read_only(file: &File, len: usize) -> std::io::Result<Mmap> {
    // SAFETY: the mapped file is only read through `&[u8]`; pack/MIDX paths are
    // replaced atomically via rename, not truncated in place (see module docs).
    // Length comes from fstat on this same `File` (see `open_from_file`).
    unsafe { MmapOptions::new().len(len).map(file) }
}

/// Read only the pack header and trailer from disk for comparison with cached bytes.
///
/// # Errors
///
/// Returns [`Error::Io`] on read failure and [`Error::CorruptObject`] for malformed packs.
pub(crate) fn fingerprint_from_file(path: &Path, hash_bytes: usize) -> Result<PackFingerprint> {
    let mut file = File::open(path).map_err(Error::Io)?;
    let file_len = file.metadata().map_err(Error::Io)?.len();
    read_fingerprint_from_open_file(&mut file, file_len, hash_bytes)
}

fn read_fingerprint_from_open_file(
    file: &mut File,
    file_len: u64,
    hash_bytes: usize,
) -> Result<PackFingerprint> {
    if hash_bytes != 20 && hash_bytes != 32 {
        return Err(Error::CorruptObject(format!(
            "unsupported pack trailer width {hash_bytes}"
        )));
    }
    let mut header = [0u8; HEADER_LEN];
    file.seek(SeekFrom::Start(0)).map_err(Error::Io)?;
    file.read_exact(&mut header).map_err(Error::Io)?;
    validate_pack_header(&header)?;
    if file_len < u64::try_from(HEADER_LEN + hash_bytes).unwrap_or(u64::MAX) {
        return Err(Error::CorruptObject(
            "pack too small for trailer".to_owned(),
        ));
    }
    let mut trailer = [0u8; 32];
    file.seek(SeekFrom::Start(file_len - hash_bytes as u64))
        .map_err(Error::Io)?;
    file.read_exact(&mut trailer[..hash_bytes])
        .map_err(Error::Io)?;
    Ok(PackFingerprint {
        header,
        hash_bytes: hash_bytes as u8,
        trailer,
    })
}

pub(crate) fn fingerprint_from_slice(bytes: &[u8], hash_bytes: usize) -> Result<PackFingerprint> {
    if hash_bytes != 20 && hash_bytes != 32 {
        return Err(Error::CorruptObject(format!(
            "unsupported pack trailer width {hash_bytes}"
        )));
    }
    if bytes.len() < HEADER_LEN {
        return Err(Error::CorruptObject("pack too small".to_owned()));
    }
    let mut header = [0u8; HEADER_LEN];
    header.copy_from_slice(&bytes[..HEADER_LEN]);
    validate_pack_header(&header)?;
    if bytes.len() < HEADER_LEN + hash_bytes {
        return Err(Error::CorruptObject(
            "pack too small for trailer".to_owned(),
        ));
    }
    let mut trailer = [0u8; 32];
    trailer[..hash_bytes].copy_from_slice(&bytes[bytes.len() - hash_bytes..]);
    Ok(PackFingerprint {
        header,
        hash_bytes: hash_bytes as u8,
        trailer,
    })
}

fn validate_pack_header(header: &[u8; HEADER_LEN]) -> Result<()> {
    if header[..4] != *PACK_SIGNATURE {
        return Err(Error::CorruptObject("bad pack signature".to_owned()));
    }
    let version = u32::from_be_bytes(header[4..8].try_into().unwrap_or([0; 4]));
    if version != 2 && version != 3 {
        return Err(Error::CorruptObject(format!(
            "unsupported pack version {version}"
        )));
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use std::io::Write;

    #[test]
    fn owned_open_roundtrip() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("tiny.pack");
        fs::write(&path, b"PACK\x00\x00\x00\x02\x00\x00\x00\x00").expect("write");
        let data = PackData::open(&path).expect("open");
        assert_eq!(&data[..], &b"PACK\x00\x00\x00\x02\x00\x00\x00\x00"[..]);
    }

    #[test]
    fn fingerprint_detects_trailer_change() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("p.pack");
        let mut pack = Vec::from(b"PACK\x00\x00\x00\x02\x00\x00\x00\x01" as &[u8]);
        pack.extend_from_slice(&[0u8; 20]);
        pack.extend_from_slice(b"aaaaaaaaaaaaaaaaaaaa");
        fs::write(&path, &pack).expect("write");
        let fp1 = fingerprint_from_file(&path, 20).expect("fp1");
        let last = pack.len() - 1;
        pack[last] = b'b';
        fs::write(&path, &pack).expect("rewrite");
        let fp2 = fingerprint_from_file(&path, 20).expect("fp2");
        assert_ne!(fp1, fp2);
    }

    #[test]
    fn fingerprint_sha256_trailer_on_pack_version_two() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("sha256.pack");
        let mut pack = Vec::from(b"PACK\x00\x00\x00\x02\x00\x00\x00\x00" as &[u8]);
        let trailer = [0xAB_u8; 32];
        pack.extend_from_slice(&trailer);
        fs::write(&path, &pack).expect("write");
        let fp = fingerprint_from_file(&path, 32).expect("fp");
        assert_eq!(&fp.trailer[..fp.hash_bytes as usize], &trailer[..]);
        assert_eq!(fp.header[4..8], 2u32.to_be_bytes());
    }

    #[test]
    fn fingerprint_rejects_invalid_trailer_width_and_malformed_header() {
        let tiny = b"PACK\x00\x00\x00\x02\x00\x00\x00\x00";
        let owned = PackData::from_owned(tiny.to_vec());
        assert!(owned.fingerprint(19).is_err());
        assert!(owned.fingerprint(33).is_err());
        assert!(fingerprint_from_file(
            &tempfile::tempdir().expect("d").path().join("missing.pack"),
            20
        )
        .is_err());

        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("bad.pack");
        fs::write(&path, b"NOTPACK\x00\x00\x00\x02\x00\x00\x00\x00").expect("write");
        let err = fingerprint_from_file(&path, 20).expect_err("bad sig");
        assert!(matches!(err, Error::CorruptObject(_)));

        fs::write(&path, b"PACK\x00\x00\x00\x09\x00\x00\x00\x00").expect("bad ver");
        let err = fingerprint_from_file(&path, 20).expect_err("bad ver");
        assert!(matches!(err, Error::CorruptObject(_)));

        fs::write(&path, b"PACK\x00\x00\x00\x02").expect("trunc");
        assert!(matches!(
            PackData::open(&path),
            Err(Error::CorruptObject(_))
        ));

        let mut short_trailer = Vec::from(b"PACK\x00\x00\x00\x02\x00\x00\x00\x00" as &[u8]);
        short_trailer.push(0);
        fs::write(&path, &short_trailer).expect("short trailer");
        let err = fingerprint_from_file(&path, 20).expect_err("short trailer");
        assert!(matches!(err, Error::CorruptObject(_)));
    }

    #[test]
    fn owned_buffer_when_pack_at_or_below_threshold() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("medium.pack");
        let mut pack = Vec::from(b"PACK\x00\x00\x00\x02\x00\x00\x00\x00" as &[u8]);
        pack.resize(8192, 0);
        pack.extend_from_slice(b"aaaaaaaaaaaaaaaaaaaa");
        fs::write(&path, &pack).expect("write");
        let data = PackData::open(&path).expect("open owned");
        assert_eq!(data.len(), pack.len());
        let fp = data.fingerprint(20).expect("fp");
        assert_eq!(&fp.trailer[..20], b"aaaaaaaaaaaaaaaaaaaa");
    }

    #[test]
    fn open_rejects_shorter_than_header() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("short.pack");
        fs::write(&path, b"PAC").expect("write");
        assert!(matches!(
            PackData::open(&path),
            Err(Error::CorruptObject(_))
        ));
    }

    #[test]
    fn fingerprint_rejects_unsupported_trailer_width() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("p.pack");
        fs::write(&path, b"PACK\x00\x00\x00\x02\x00\x00\x00\x00").expect("write");
        assert!(fingerprint_from_file(&path, 16).is_err());
        let data = PackData::open(&path).expect("open");
        assert!(data.fingerprint(16).is_err());
    }

    #[test]
    fn fingerprint_accepts_pack_version_three_header() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("v3.pack");
        let mut pack = Vec::from(b"PACK\x00\x00\x00\x03\x00\x00\x00\x00" as &[u8]);
        pack.extend_from_slice(&[0u8; 20]);
        fs::write(&path, &pack).expect("write");
        let fp = fingerprint_from_file(&path, 20).expect("fp");
        assert_eq!(u32::from_be_bytes(fp.header[4..8].try_into().unwrap()), 3);
    }

    #[test]
    fn fingerprint_from_slice_rejects_bad_signature() {
        let bytes = b"NOPE\x00\x00\x00\x02\x00\x00\x00\x00";
        assert!(matches!(
            fingerprint_from_slice(bytes, 20),
            Err(Error::CorruptObject(_))
        ));
    }

    /// Resident KiB of the mapping that contains `addr`, from `/proc/self/smaps`.
    #[cfg(target_os = "linux")]
    fn mapping_rss_kib(addr: *const u8) -> u64 {
        let smaps = fs::read_to_string("/proc/self/smaps").expect("smaps");
        let addr = addr as usize;
        let mut inside = false;
        for line in smaps.lines() {
            let range = line.split_whitespace().next().unwrap_or_default();
            if let Some((start, end)) = range.split_once('-') {
                if let (Ok(start), Ok(end)) = (
                    usize::from_str_radix(start, 16),
                    usize::from_str_radix(end, 16),
                ) {
                    inside = (start..end).contains(&addr);
                    continue;
                }
            }
            if inside {
                if let Some(kib) = line.strip_prefix("Rss:") {
                    return kib
                        .trim()
                        .trim_end_matches("kB")
                        .trim()
                        .parse()
                        .expect("kB");
                }
            }
        }
        panic!("no mapping contains {addr:#x}");
    }

    #[cfg(target_os = "linux")]
    fn mapped_test_pack(dir: &Path, len: usize) -> (Arc<PackData>, Vec<u8>) {
        let path = dir.join("resident.pack");
        let mut bytes = Vec::from(&b"PACK\x00\x00\x00\x02\x00\x00\x00\x00"[..]);
        bytes.extend((0..len).map(|i| u8::try_from(i % 251).unwrap_or(0)));
        fs::write(&path, &bytes).expect("write");
        let data = PackData::open(&path).expect("open");
        assert!(data.is_mapped_backing());
        (data, bytes)
    }

    #[cfg(target_os = "linux")]
    #[test]
    fn release_resident_drops_pages_but_keeps_the_bytes() {
        let dir = tempfile::tempdir().expect("tempdir");
        let (data, bytes) = mapped_test_pack(dir.path(), 24 << 20);
        let digest = HashAlgo::Sha1.checksum_digest(&data);
        assert!(
            mapping_rss_kib(data.as_ptr()) > 20 << 10,
            "a full read maps the file"
        );
        data.release_resident();
        assert!(
            mapping_rss_kib(data.as_ptr()) < 1 << 10,
            "release drops the pages"
        );
        assert_eq!(&data[..], &bytes[..]);
        assert_eq!(HashAlgo::Sha1.checksum_digest(&data), digest);
    }

    #[cfg(target_os = "linux")]
    #[test]
    fn resident_budget_bounds_a_full_pass() {
        let dir = tempfile::tempdir().expect("tempdir");
        let (data, bytes) = mapped_test_pack(dir.path(), 24 << 20);
        data.release_resident();
        let budget = ResidentBudget::new(&data, 4 << 20);
        assert_eq!(
            budget.checksum(HashAlgo::Sha1, &data),
            HashAlgo::Sha1.checksum_digest(&bytes)
        );
        let rss = mapping_rss_kib(data.as_ptr());
        assert!(rss <= 6 << 10, "budgeted pass left {rss} KiB resident");

        data.release_resident();
        let unbounded = ResidentBudget::unbounded();
        let _ = unbounded.checksum(HashAlgo::Sha1, &data);
        assert!(
            mapping_rss_kib(data.as_ptr()) > 20 << 10,
            "unbounded never releases"
        );
    }

    #[test]
    fn release_resident_on_owned_bytes_is_a_no_op() {
        let data = PackData::from_owned(b"PACK\x00\x00\x00\x02\x00\x00\x00\x00".to_vec());
        data.release_resident();
        let budget = ResidentBudget::new(&data, 0);
        budget.charge(1);
        assert_eq!(&data[..4], b"PACK");
    }

    #[test]
    fn mmap_large_sparse_pack() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("sparse.pack");
        let mut file = File::create(&path).expect("create");
        file.write_all(b"PACK\x00\x00\x00\x02\x00\x00\x00\x01")
            .expect("header");
        let huge: u64 = 1 << 32;
        file.seek(SeekFrom::Start(huge)).expect("seek");
        file.write_all(b"blob payload padding").expect("payload");
        file.write_all(b"bbbbbbbbbbbbbbbbbbbb").expect("trailer");
        drop(file);
        let len = fs::metadata(&path).expect("meta").len();
        assert!(len > huge);
        let data = PackData::open(&path).expect("mmap sparse");
        assert_eq!(&data[..4], b"PACK");
        assert_eq!(data.len() as u64, len);
    }
}