armdb 0.7.0

use std::fs;
use std::path::{Path, PathBuf};

#[cfg(test)]
use std::io;

use zerocopy::FromBytes;

use crate::Key;
use crate::entry::{EntryHeader, SEQUENCE_MASK, TOMBSTONE_BIT, entry_size};
use crate::error::{DbError, DbResult};
use crate::io::direct;

type ReadFn<'a> = dyn Fn(&std::fs::File, u64, usize) -> DbResult<Vec<u8>> + 'a;

#[cfg(test)]
static HINT_WRITE_FAULT: std::sync::Mutex<Option<(PathBuf, io::ErrorKind, String)>> =
    std::sync::Mutex::new(None);

#[cfg(test)]
pub(crate) struct HintWriteFaultGuard {
    path: PathBuf,
}

#[cfg(test)]
impl Drop for HintWriteFaultGuard {
    fn drop(&mut self) {
        let mut fault = HINT_WRITE_FAULT.lock().expect("hint fault mutex poisoned");
        if fault
            .as_ref()
            .is_some_and(|(path, _, _)| path == &self.path)
        {
            *fault = None;
        }
    }
}

#[cfg(test)]
pub(crate) fn fail_hint_write_for(path: PathBuf, error: io::Error) -> HintWriteFaultGuard {
    let mut fault = HINT_WRITE_FAULT.lock().expect("hint fault mutex poisoned");
    assert!(fault.is_none(), "a hint write fault is already registered");
    *fault = Some((path.clone(), error.kind(), error.to_string()));
    HintWriteFaultGuard { path }
}

#[cfg(test)]
fn injected_hint_write_error(path: &Path) -> Option<io::Error> {
    let fault = HINT_WRITE_FAULT.lock().expect("hint fault mutex poisoned");
    let (target, kind, message) = fault.as_ref()?;
    (target == path).then(|| io::Error::new(*kind, message.clone()))
}

/// Size of a single hint entry: GSN(8) + Key(key_len) + Offset(8) + Len(4) + CRC(4).
#[inline]
pub const fn hint_entry_size(key_len: usize) -> usize {
    8 + key_len + 8 + 4 + 4 // GSN(8) + Key(key_len) + Offset(8) + Len(4) + CRC(4)
}

/// Hint file format v3 magic. There is no compatibility decoder: databases
/// created during development must regenerate their hint files.
pub(crate) const HINT_MAGIC: [u8; 4] = *b"AHF3";
pub(crate) const HINT_VERSION: u8 = 3;
pub(crate) const HINT_HEADER_SIZE: usize = 36;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct HintHeader {
    pub logical_end: u64,
    pub max_gsn: u64,
    pub entry_count: u64,
}

fn entry_metadata(entries: &[u8], key_len: usize) -> Option<(u64, u64)> {
    let entry_size = hint_entry_size(key_len);
    if !entries.len().is_multiple_of(entry_size) {
        return None;
    }

    let mut max_gsn = 0;
    let mut entry_count = 0;
    for entry in entries.chunks_exact(entry_size) {
        let gsn = u64::from_ne_bytes(entry[..8].try_into().ok()?) & SEQUENCE_MASK;
        max_gsn = max_gsn.max(gsn);
        entry_count += 1;
    }
    Some((max_gsn, entry_count))
}

/// Wrap a raw entries buffer in a checksummed v3 header.
pub(crate) fn hint_file_bytes(
    entries: &[u8],
    logical_end: u64,
    key_len: usize,
) -> DbResult<Vec<u8>> {
    let Some((max_gsn, entry_count)) = entry_metadata(entries, key_len) else {
        return Err(DbError::FormatMismatch(
            "hint entries are not aligned to the key size".into(),
        ));
    };

    let mut buf = Vec::with_capacity(HINT_HEADER_SIZE + entries.len());
    buf.extend_from_slice(&HINT_MAGIC);
    buf.push(HINT_VERSION);
    buf.push(0); // flags
    buf.extend_from_slice(&(HINT_HEADER_SIZE as u16).to_le_bytes());
    buf.extend_from_slice(&logical_end.to_le_bytes());
    buf.extend_from_slice(&max_gsn.to_le_bytes());
    buf.extend_from_slice(&entry_count.to_le_bytes());
    let header_crc = crc32fast::hash(&buf[..32]);
    buf.extend_from_slice(&header_crc.to_le_bytes());
    buf.extend_from_slice(entries);
    Ok(buf)
}

/// Validate only the fixed v3 header and split it from the payload in O(1).
pub(crate) fn parse_hint_header(data: &[u8]) -> Option<(HintHeader, &[u8])> {
    if data.len() < HINT_HEADER_SIZE {
        return None;
    }
    if data[..4] != HINT_MAGIC {
        return None;
    }
    if data[4] != HINT_VERSION {
        return None;
    }
    if data[5] != 0 {
        return None;
    }
    if u16::from_le_bytes(data[6..8].try_into().ok()?) as usize != HINT_HEADER_SIZE {
        return None;
    }
    let expected_crc = u32::from_le_bytes(data[32..36].try_into().ok()?);
    if crc32fast::hash(&data[..32]) != expected_crc {
        return None;
    }

    let header = HintHeader {
        logical_end: u64::from_le_bytes(data[8..16].try_into().ok()?),
        max_gsn: u64::from_le_bytes(data[16..24].try_into().ok()?),
        entry_count: u64::from_le_bytes(data[24..32].try_into().ok()?),
    };
    Some((header, &data[HINT_HEADER_SIZE..]))
}

fn read_hint_header_from(reader: &mut impl std::io::Read) -> std::io::Result<Option<HintHeader>> {
    let mut data = [0u8; HINT_HEADER_SIZE];
    match reader.read_exact(&mut data) {
        Ok(()) => Ok(parse_hint_header(&data).map(|(header, _)| header)),
        Err(error) if error.kind() == std::io::ErrorKind::UnexpectedEof => Ok(None),
        Err(error) => Err(error),
    }
}

/// Read and validate exactly the fixed v3 header without touching the payload.
pub(crate) fn read_hint_header(path: &Path) -> DbResult<Option<HintHeader>> {
    let mut file = match fs::File::open(path) {
        Ok(file) => file,
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
        Err(error) => return Err(DbError::Io(error)),
    };
    read_hint_header_from(&mut file).map_err(DbError::Io)
}

/// Validate the v3 header and all payload metadata in O(number of entries).
pub(crate) fn validate_hint_file(data: &[u8], key_len: usize) -> Option<(HintHeader, &[u8])> {
    let (header, entries) = parse_hint_header(data)?;
    let (max_gsn, entry_count) = entry_metadata(entries, key_len)?;
    if header.entry_count != entry_count || header.max_gsn != max_gsn {
        return None;
    }
    Some((header, entries))
}

/// A parsed hint entry.
#[derive(Debug, Clone, Copy)]
pub struct HintEntry<K: Key> {
    pub gsn: u64,
    pub key: K,
    pub value_offset: u64,
    pub value_len: u32,
    pub crc32: u32,
}

impl<K: Key> HintEntry<K> {
    #[inline]
    pub fn is_tombstone(&self) -> bool {
        self.gsn & TOMBSTONE_BIT != 0
    }

    #[inline]
    pub fn sequence(&self) -> u64 {
        self.gsn & SEQUENCE_MASK
    }
}

/// Generate hint data by scanning a data file.
///
/// Reads only headers + keys (skips values), making this much cheaper
/// than a full recovery scan. Takes `key_len` at runtime so it can be
/// called from non-generic `ShardInner`.
pub fn generate_hint_data_dyn(
    data_file: &std::fs::File,
    file_len: u64,
    key_len: usize,
) -> DbResult<Vec<u8>> {
    generate_hint_data_with_reader(data_file, file_len, key_len, &read_plain)
}

/// Generate hint data from an encrypted data file.
#[cfg(feature = "encryption")]
pub fn generate_hint_data_dyn_encrypted(
    data_file: &std::fs::File,
    file_len: u64,
    key_len: usize,
    cipher: &crate::crypto::PageCipher,
    tag_file: &crate::io::tags::TagFile,
    file_id: u32,
) -> DbResult<Vec<u8>> {
    let reader = move |file: &std::fs::File, offset: u64, len: usize| {
        direct::pread_value_encrypted(file, tag_file, cipher, file_id, offset, len)
    };
    generate_hint_data_with_reader(data_file, file_len, key_len, &reader)
}

pub(crate) fn generate_hint_data_with_reader(
    data_file: &std::fs::File,
    logical_end: u64,
    key_len: usize,
    read_fn: &ReadFn<'_>,
) -> DbResult<Vec<u8>> {
    let hint_sz = hint_entry_size(key_len);
    let min_data_entry = entry_size(key_len, 0);
    let estimated = logical_end.checked_div(min_data_entry).unwrap_or_default() as usize;
    let mut entries = Vec::with_capacity(estimated * hint_sz);

    let header_size = size_of::<EntryHeader>() as u64;
    let mut offset: u64 = 0;

    while offset + header_size <= logical_end {
        let header_bytes = match read_fn(data_file, offset, size_of::<EntryHeader>()) {
            Ok(b) => b,
            Err(_) => break,
        };
        let header: EntryHeader = match EntryHeader::read_from_bytes(&header_bytes) {
            Ok(h) => h,
            Err(_) => break,
        };

        // Page padding sentinel: force-flush pads the partial page with zeros.
        if header.gsn == 0
            && header.value_len == 0
            && header.crc32 == 0
            && crate::entry::check_page_padding_at(
                data_file,
                offset + size_of::<EntryHeader>() as u64,
                read_fn,
            )?
        {
            const PAGE_SIZE: u64 = 4096;
            let next_page =
                (offset + size_of::<EntryHeader>() as u64).div_ceil(PAGE_SIZE) * PAGE_SIZE;
            if next_page >= logical_end {
                break;
            }
            offset = next_page;
            continue;
        }

        let total = entry_size(key_len, header.value_len);
        if offset + total > logical_end {
            break; // partial entry at end
        }

        // Read only the key (K bytes after header)
        let key_bytes = read_fn(data_file, offset + size_of::<EntryHeader>() as u64, key_len)?;

        let value_offset = offset + size_of::<EntryHeader>() as u64 + key_len as u64;

        // Serialize: GSN(8) | Key(key_len) | Offset(8) | Len(4) | CRC(4)
        entries.extend_from_slice(&header.gsn.to_ne_bytes());
        entries.extend_from_slice(&key_bytes);
        entries.extend_from_slice(&value_offset.to_ne_bytes());
        entries.extend_from_slice(&header.value_len.to_ne_bytes());
        entries.extend_from_slice(&header.crc32.to_ne_bytes());

        offset += total;
    }

    if offset < logical_end {
        return Err(DbError::CorruptedEntry { offset });
    }

    hint_file_bytes(&entries, logical_end, key_len)
}

fn read_plain(file: &std::fs::File, offset: u64, len: usize) -> DbResult<Vec<u8>> {
    direct::pread_value(file, offset, len)
}

/// Parse hint entries from a complete v3 hint file image.
///
/// The fixed header is validated and skipped; a file with a missing/unknown
/// header yields an empty iterator (callers that need to distinguish that from
/// a genuinely empty entry list call [`parse_hint_header`] first).
pub fn parse_hint_entries<K: Key>(data: &[u8]) -> impl Iterator<Item = HintEntry<K>> + '_ {
    let entry_sz = hint_entry_size(size_of::<K>());
    let entries = match parse_hint_header(data) {
        Some((_, entries)) => entries,
        None => &[][..],
    };
    entries.chunks_exact(entry_sz).filter_map(|chunk| {
        let gsn = u64::from_ne_bytes(chunk[..8].try_into().ok()?);
        let key: K = K::from_bytes(&chunk[8..8 + size_of::<K>()]);
        let value_offset = u64::from_ne_bytes(
            chunk[8 + size_of::<K>()..8 + size_of::<K>() + 8]
                .try_into()
                .ok()?,
        );
        let value_len = u32::from_ne_bytes(
            chunk[8 + size_of::<K>() + 8..8 + size_of::<K>() + 12]
                .try_into()
                .ok()?,
        );
        let crc32 = u32::from_ne_bytes(
            chunk[8 + size_of::<K>() + 12..8 + size_of::<K>() + 16]
                .try_into()
                .ok()?,
        );
        Some(HintEntry {
            gsn,
            key,
            value_offset,
            value_len,
            crc32,
        })
    })
}

/// Write hint data to a file. Uses standard buffered I/O (not O_DIRECT) —
/// hint files are small and written once.
pub fn write_hint_file(path: &Path, data: &[u8]) -> DbResult<()> {
    use std::io::Write;

    #[cfg(test)]
    if let Some(error) = injected_hint_write_error(path) {
        return Err(DbError::Io(error));
    }

    let tmp_path = path.with_extension("hint.tmp");
    let mut file = fs::OpenOptions::new()
        .create(true)
        .truncate(true)
        .write(true)
        .open(&tmp_path)?;
    file.write_all(data)?;
    file.sync_data()?;
    drop(file);
    fs::rename(&tmp_path, path)?;
    let parent = path.parent().ok_or_else(|| {
        DbError::Io(std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            "hint path has no parent directory",
        ))
    })?;
    direct::sync_dir(parent)?;
    Ok(())
}

/// Read a hint file into memory. Returns `None` if the file doesn't exist.
pub fn read_hint_file(path: &Path) -> DbResult<Option<Vec<u8>>> {
    match fs::read(path) {
        Ok(data) => Ok(Some(data)),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
        Err(e) => Err(DbError::Io(e)),
    }
}

/// Compute the hint file path for a data file path.
/// e.g. "000001.data" → "000001.hint"
#[inline]
pub fn hint_path_for_data(data_path: &Path) -> PathBuf {
    data_path.with_extension("hint")
}

/// Scan a directory for `.hint` files and return sorted file IDs.
pub fn scan_hint_files(dir: &Path) -> DbResult<Vec<u32>> {
    let mut ids = Vec::new();
    if !dir.exists() {
        return Ok(ids);
    }
    for entry in fs::read_dir(dir)? {
        let entry = entry?;
        let name = entry.file_name();
        let name = name.to_string_lossy();
        if name.ends_with(".hint")
            && let Ok(id) = name.trim_end_matches(".hint").parse::<u32>()
        {
            ids.push(id);
        }
    }
    ids.sort();
    Ok(ids)
}

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

    struct ErrorAfterHeader {
        data: Vec<u8>,
        offset: usize,
    }

    impl std::io::Read for ErrorAfterHeader {
        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
            if self.offset >= HINT_HEADER_SIZE {
                return Err(std::io::Error::other("payload read attempted"));
            }
            let count = buf
                .len()
                .min(HINT_HEADER_SIZE - self.offset)
                .min(self.data.len() - self.offset);
            buf[..count].copy_from_slice(&self.data[self.offset..self.offset + count]);
            self.offset += count;
            Ok(count)
        }
    }

    fn raw_hint_entry(gsn: u64, key: [u8; 8], value_offset: u64, value_len: u32) -> Vec<u8> {
        let mut entry = Vec::new();
        entry.extend_from_slice(&gsn.to_ne_bytes());
        entry.extend_from_slice(&key);
        entry.extend_from_slice(&value_offset.to_ne_bytes());
        entry.extend_from_slice(&value_len.to_ne_bytes());
        entry.extend_from_slice(&0u32.to_ne_bytes());
        entry
    }

    fn refresh_header_crc(bytes: &mut [u8]) {
        let crc = crc32fast::hash(&bytes[..32]);
        bytes[32..36].copy_from_slice(&crc.to_le_bytes());
    }

    #[test]
    fn hint_v3_round_trip_has_exact_header_and_metadata() {
        let mut entries = raw_hint_entry(17, [1; 8], 16, 4);
        entries.extend_from_slice(&raw_hint_entry(3, [2; 8], 48, 4));
        let bytes = hint_file_bytes(&entries, 4096, 8).unwrap();
        assert_eq!(&bytes[..4], b"AHF3");
        assert_eq!(HINT_HEADER_SIZE, 36);
        assert_eq!(&bytes[6..8], &36u16.to_le_bytes());
        assert_eq!(&bytes[8..16], &4096u64.to_le_bytes());
        assert_eq!(&bytes[16..24], &17u64.to_le_bytes());
        assert_eq!(&bytes[24..32], &2u64.to_le_bytes());
        let (header, payload) = validate_hint_file(&bytes, 8).unwrap();
        assert_eq!(header.logical_end, 4096);
        assert_eq!(header.max_gsn, 17);
        assert_eq!(header.entry_count, 2);
        assert_eq!(payload.len(), 2 * hint_entry_size(8));
    }

    #[test]
    fn hint_v3_empty_hint_has_zero_metadata() {
        let bytes = hint_file_bytes(&[], 0, 8).unwrap();
        assert_eq!(bytes.len(), HINT_HEADER_SIZE);
        let (header, payload) = validate_hint_file(&bytes, 8).unwrap();
        assert_eq!(header.logical_end, 0);
        assert_eq!(header.max_gsn, 0);
        assert_eq!(header.entry_count, 0);
        assert!(payload.is_empty());
    }

    #[test]
    fn hint_v3_header_read_is_capped_at_fixed_size() {
        let entries = raw_hint_entry(17, [1; 8], 16, 4);
        let bytes = hint_file_bytes(&entries, 4096, 8).unwrap();
        let mut reader = ErrorAfterHeader {
            data: bytes,
            offset: 0,
        };

        let header = read_hint_header_from(&mut reader).unwrap().unwrap();

        assert_eq!(header.max_gsn, 17);
        assert_eq!(reader.offset, HINT_HEADER_SIZE);
    }

    #[test]
    fn hint_v3_rejects_unknown_flags() {
        let mut bytes = hint_file_bytes(&[], 0, 8).unwrap();
        bytes[5] = 1;
        refresh_header_crc(&mut bytes);
        assert!(parse_hint_header(&bytes).is_none());
    }

    #[test]
    fn hint_v3_rejects_wrong_header_len() {
        let mut bytes = hint_file_bytes(&[], 0, 8).unwrap();
        bytes[6..8].copy_from_slice(&35u16.to_le_bytes());
        refresh_header_crc(&mut bytes);
        assert!(parse_hint_header(&bytes).is_none());
    }

    #[test]
    fn hint_v3_rejects_corrupt_header_crc() {
        let mut bytes = hint_file_bytes(&[], 0, 8).unwrap();
        bytes[32] ^= 0xff;
        assert!(parse_hint_header(&bytes).is_none());
    }

    #[test]
    fn hint_v3_rejects_entry_count_mismatch() {
        let entries = raw_hint_entry(17, [1; 8], 16, 4);
        let mut bytes = hint_file_bytes(&entries, 4096, 8).unwrap();
        bytes[24..32].copy_from_slice(&2u64.to_le_bytes());
        refresh_header_crc(&mut bytes);
        assert!(validate_hint_file(&bytes, 8).is_none());
    }

    #[test]
    fn hint_v3_rejects_max_gsn_mismatch() {
        let entries = raw_hint_entry(17, [1; 8], 16, 4);
        let mut bytes = hint_file_bytes(&entries, 4096, 8).unwrap();
        bytes[16..24].copy_from_slice(&18u64.to_le_bytes());
        refresh_header_crc(&mut bytes);
        assert!(validate_hint_file(&bytes, 8).is_none());
    }

    #[test]
    fn hint_sequence_strips_flag_region() {
        use crate::entry::SOFT_DELETE_BIT;
        let e = HintEntry::<[u8; 8]> {
            gsn: 7 | SOFT_DELETE_BIT,
            key: [0u8; 8],
            value_offset: 0,
            value_len: 0,
            crc32: 0,
        };
        assert_eq!(e.sequence(), 7);
        assert!(!e.is_tombstone());
    }

    #[test]
    fn hint_generation_skips_mid_file_padding() {
        use crate::entry::serialize_entry;

        let dir = std::env::temp_dir().join(format!("armdb_hintpad_{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("000001.data");

        let key_len = 8usize;
        let e1 = serialize_entry(1, &[1u8; 8], b"v1", false);
        let e2 = serialize_entry(2, &[2u8; 8], b"v2", false);

        // Layout: e1, zero-pad to next page boundary, e2.
        let mut file_bytes = Vec::new();
        file_bytes.extend_from_slice(&e1);
        let pad = 4096 - (file_bytes.len() % 4096);
        file_bytes.extend(std::iter::repeat_n(0u8, pad));
        file_bytes.extend_from_slice(&e2);

        let f = crate::io::direct::open_bulk_write(&path).unwrap();
        crate::io::direct::pwrite_at(&f, &file_bytes, 0).unwrap();
        crate::io::direct::fsync(&f).unwrap();
        drop(f);

        let rf = crate::io::direct::open_bulk_read(&path).unwrap();
        let data =
            generate_hint_data_with_reader(&rf, file_bytes.len() as u64, key_len, &read_plain)
                .unwrap();
        // v3 header + 2 entries.
        assert_eq!(
            parse_hint_header(&data).unwrap().0.logical_end,
            file_bytes.len() as u64
        );
        let count = (data.len() - HINT_HEADER_SIZE) / hint_entry_size(key_len);
        assert_eq!(
            count, 2,
            "both real entries must produce hints, padding none"
        );

        std::fs::remove_dir_all(&dir).ok();
    }

    /// Encrypted variant: hint generation must skip mid-file page padding when
    /// scanning an *encrypted* data file. A page-aligned encrypted tree pads the
    /// partial page on `flush()`, so the second put lands at the next page
    /// boundary, leaving padding (ciphertext of zeros) between the two entries.
    /// `generate_hint_data_dyn_encrypted` reads through the decrypting reader, so
    /// it must recognise the decrypted-zero padding and emit hints only for the
    /// two real entries.
    #[cfg(feature = "encryption")]
    #[test]
    fn hint_generation_skips_mid_file_padding_encrypted() {
        use crate::config::Config;
        use crate::const_tree::ConstTree;
        use crate::crypto::PageCipher;
        use crate::io::tags::{TagFile, tags_path_for_data};

        let dir = tempfile::tempdir().unwrap();
        let key = [0x42u8; 32];
        let mut config = Config::test();
        config.shard_count = 1;
        config.encryption_key = Some(key);

        {
            let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config).unwrap();
            tree.put(&1u64.to_be_bytes(), &10u64.to_be_bytes()).unwrap();
            // Final flush pads the partial page; the next put starts at the next
            // page boundary, leaving mid-file padding in between.
            tree.flush().unwrap();
            tree.put(&2u64.to_be_bytes(), &20u64.to_be_bytes()).unwrap();
            tree.flush().unwrap();
            tree.close().unwrap();
        }

        // The shard's first data file is id 1 (`000001.data`). The physical file
        // is two 4096-byte pages, but production generates hints over the
        // *logical* data end (`write_offset`), not `metadata().len()`, so the
        // trailing final-page padding is excluded. Mirror that here: the two
        // real entries sit at offset 0 and at the 4096 boundary (after the
        // mid-file padding the first flush left), so the logical end is
        // `4096 + entry_size`.
        let data_path = dir.path().join("shard_000").join("000001.data");
        let data_file = crate::io::direct::open_bulk_read(&data_path).unwrap();
        let physical_len = data_file.metadata().unwrap().len();
        assert!(
            physical_len >= 8192,
            "expected at least two flushed pages, got {physical_len}"
        );
        let logical_len = 4096 + entry_size(8, 8);

        // The DB derives one HKDF subkey per shard from the master key +
        // db.salt (D1); reconstruct the cipher the same way production does
        // rather than deriving straight from the master key.
        let salt = crate::salt::read(dir.path()).expect("read db.salt");
        let cipher = PageCipher::derive(&key, &salt, 0).expect("derive cipher");
        let tag_file = TagFile::open_read(&tags_path_for_data(&data_path)).unwrap();

        let data =
            generate_hint_data_dyn_encrypted(&data_file, logical_len, 8, &cipher, &tag_file, 1)
                .unwrap();
        // v3 header records the logical end; two entries follow it.
        assert_eq!(parse_hint_header(&data).unwrap().0.logical_end, logical_len);
        let count = (data.len() - HINT_HEADER_SIZE) / hint_entry_size(8);
        assert_eq!(
            count, 2,
            "both real entries must produce hints across encrypted padding"
        );
    }
}