armdb 0.11.2

use std::fs::{File, OpenOptions};
use std::path::Path;

use crate::error::{DbError, DbResult};
use crate::io::aligned_buf::AlignedBuf;

/// Fill `buf` (length is a 4096 multiple) starting at `aligned_offset` (4096
/// aligned) by reading from `file`, requiring at least `needed` bytes. Each
/// retry after a short read re-aligns the read offset and destination slice
/// **down** to the previous sector boundary, so every issued read stays
/// 4096-aligned (required under `O_DIRECT`). Returns `UnexpectedEof` if fewer
/// than `needed` bytes are available.
fn aligned_read_exact(
    file: &File,
    buf: &mut [u8],
    aligned_offset: u64,
    needed: usize,
) -> DbResult<()> {
    use std::os::unix::fs::FileExt;
    const SECTOR: usize = 4096;
    let mut total_read = 0usize;
    while total_read < needed {
        // Re-align the cursor down to the previous sector boundary so the read
        // offset and the destination slice start are both 4096-aligned.
        let aligned_cursor = total_read & !(SECTOR - 1);
        let r = file.read_at(
            &mut buf[aligned_cursor..],
            aligned_offset + aligned_cursor as u64,
        )?;
        if r == 0 {
            break; // EOF
        }
        // The retry re-aligns `aligned_cursor` *down* to a sector boundary, so a
        // short read whose length isn't sector-aligned (truncated/corrupted file)
        // can re-read the same tail without advancing. Break on no progress so we
        // fall through to the `UnexpectedEof` error instead of spinning forever.
        let new_total = aligned_cursor + r;
        if new_total <= total_read {
            break;
        }
        total_read = new_total;
    }
    if total_read < needed {
        return Err(DbError::Io(std::io::Error::new(
            std::io::ErrorKind::UnexpectedEof,
            "failed to read full value",
        )));
    }
    Ok(())
}

#[cfg(target_os = "macos")]
unsafe extern "C" {
    fn fcntl(fd: std::os::raw::c_int, cmd: std::os::raw::c_int, ...) -> std::os::raw::c_int;
}

/// Read `len` bytes from `file` at `offset`.
///
/// On Linux with O_DIRECT this uses aligned buffers.
/// On macOS/other platforms this uses standard pread.
pub fn pread_value(file: &File, offset: u64, len: usize) -> DbResult<Vec<u8>> {
    // For O_DIRECT we need aligned buffers and aligned reads.
    let sector_size = 4096;
    let aligned_offset = offset & !(sector_size - 1);
    let diff = (offset - aligned_offset) as usize;
    let aligned_len = (diff + len + sector_size as usize - 1) & !(sector_size as usize - 1);

    let mut buf = AlignedBuf::zeroed(aligned_len);

    aligned_read_exact(file, &mut buf, aligned_offset, diff + len)?;

    let mut result = vec![0u8; len];
    result.copy_from_slice(&buf[diff..diff + len]);
    Ok(result)
}

#[cfg(feature = "var-collections")]
/// Read a single 4096-byte aligned block from `file`.
/// Short reads (end of file) leave remaining bytes zeroed.
/// Read a 4096-byte aligned block. Returns `(block, bytes_read)`.
/// Short reads (at end of file) are zero-padded; `bytes_read` reflects
/// actual data so callers can avoid caching partial blocks.
pub fn pread_block(file: &File, block_offset: u64) -> DbResult<(AlignedBuf, usize)> {
    use std::os::unix::fs::FileExt;
    debug_assert!(
        block_offset & 4095 == 0,
        "block_offset must be 4096-aligned"
    );

    let mut buf = AlignedBuf::zeroed(4096);
    let n = file.read_at(&mut buf, block_offset)?;
    Ok((buf, n))
}

/// Append data to a file at the given offset. Returns number of bytes written.
pub fn pwrite_at(file: &File, data: &[u8], offset: u64) -> DbResult<()> {
    use std::os::unix::fs::FileExt;
    file.write_all_at(data, offset)?;
    Ok(())
}

/// Sync file data to disk.
pub fn fsync(file: &File) -> DbResult<()> {
    file.sync_data()?;
    Ok(())
}

/// Test-only counter of [`sync_dir`] invocations. Lets rotation/open tests
/// assert that the shard directory is fsynced after a new data file is created
/// (F5) without needing to observe an actual power-loss.
#[cfg(test)]
pub(crate) static SYNC_DIR_CALLS: std::sync::atomic::AtomicU64 =
    std::sync::atomic::AtomicU64::new(0);

/// fsync a directory so that a preceding metadata operation (a rename, create,
/// or unlink of a file inside it) is durable. POSIX does not guarantee that an
/// unlink/rename survives a crash until the containing directory is fsynced.
pub fn sync_dir(dir: &Path) -> DbResult<()> {
    #[cfg(test)]
    SYNC_DIR_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    let d = File::open(dir)?;
    d.sync_all()?;
    Ok(())
}

/// Open a file for reading on the hot serving path. On Linux, sets `O_DIRECT`
/// when `direct` is true and the open succeeds; on `EINVAL` (filesystem does
/// not support direct I/O) it retries buffered and logs a warning. On macOS,
/// sets `F_NOCACHE`. Otherwise a plain buffered read handle.
pub fn open_serving_read(path: &Path, #[allow(unused_variables)] direct: bool) -> DbResult<File> {
    #[cfg(target_os = "linux")]
    if direct {
        match open_direct_linux(path, /*write=*/ false) {
            Ok(f) => return Ok(f),
            Err(e) if is_einval(&e) => {
                tracing::warn!(
                    ?path,
                    "O_DIRECT read open rejected by filesystem; using buffered"
                );
            }
            Err(e) => return Err(e),
        }
    }
    let file = OpenOptions::new().read(true).open(path)?;
    #[cfg(target_os = "macos")]
    {
        use std::os::unix::io::AsRawFd;
        unsafe {
            fcntl(file.as_raw_fd(), 48 /* F_NOCACHE */, 1);
        }
    }
    Ok(file)
}

/// Open a file for writing (append) on the hot serving path. Same `O_DIRECT`
/// semantics as [`open_serving_read`].
pub fn open_serving_write(path: &Path, #[allow(unused_variables)] direct: bool) -> DbResult<File> {
    #[cfg(target_os = "linux")]
    if direct {
        match open_direct_linux(path, /*write=*/ true) {
            Ok(f) => return Ok(f),
            Err(e) if is_einval(&e) => {
                tracing::warn!(
                    ?path,
                    "O_DIRECT write open rejected by filesystem; using buffered"
                );
            }
            Err(e) => return Err(e),
        }
    }
    let file = OpenOptions::new()
        .create(true)
        .read(true)
        .write(true)
        .truncate(false)
        .open(path)?;
    #[cfg(target_os = "macos")]
    {
        use std::os::unix::io::AsRawFd;
        unsafe {
            fcntl(file.as_raw_fd(), 48 /* F_NOCACHE */, 1);
        }
    }
    Ok(file)
}

/// Open a file for reading on the cold/bulk path (recovery, compaction inputs).
/// Always buffered; advises the kernel of sequential access on Linux.
pub fn open_bulk_read(path: &Path) -> DbResult<File> {
    let file = OpenOptions::new().read(true).open(path)?;
    #[cfg(target_os = "linux")]
    fadvise_sequential(&file);
    Ok(file)
}

/// Open a file for writing on the cold/bulk path (compaction tmp output).
/// Always buffered.
pub fn open_bulk_write(path: &Path) -> DbResult<File> {
    let file = OpenOptions::new()
        .create(true)
        .read(true)
        .write(true)
        .truncate(false)
        .open(path)?;
    #[cfg(target_os = "linux")]
    fadvise_sequential(&file);
    Ok(file)
}

/// Probe whether the directory's filesystem supports `O_DIRECT`. Creates a
/// temp file, opens it with `O_DIRECT`, does one aligned 4096 write + read,
/// removes it. Returns false on any failure (e.g. tmpfs `EINVAL`).
#[allow(dead_code)]
pub fn probe_direct_io(dir: &Path) -> bool {
    #[cfg(target_os = "linux")]
    {
        let path = dir.join(format!(".armdb_odirect_probe_{}", std::process::id()));
        let result = (|| -> DbResult<()> {
            let file = open_direct_linux(&path, /*write=*/ true)?;
            let buf = AlignedBuf::zeroed(4096);
            pwrite_at(&file, &buf, 0)?;
            let mut rbuf = AlignedBuf::zeroed(4096);
            use std::os::unix::fs::FileExt;
            file.read_at(&mut rbuf, 0)?;
            Ok(())
        })();
        let _ = std::fs::remove_file(&path);
        result.is_ok()
    }
    #[cfg(not(target_os = "linux"))]
    {
        let _ = dir;
        false
    }
}

/// Advise the kernel to drop cached pages for `[offset, offset+len)` after a
/// bulk pass. No-op off Linux.
pub fn fadvise_dontneed(file: &File, offset: u64, len: u64) {
    #[cfg(target_os = "linux")]
    {
        use std::num::NonZeroU64;
        use std::os::unix::io::AsFd;
        let len = NonZeroU64::new(len);
        let _ = rustix::fs::fadvise(file.as_fd(), offset, len, rustix::fs::Advice::DontNeed);
    }
    #[cfg(not(target_os = "linux"))]
    {
        let _ = (file, offset, len);
    }
}

#[cfg(target_os = "linux")]
fn fadvise_sequential(file: &File) {
    use std::os::unix::io::AsFd;
    // `len: None` — advise through end of file (POSIX len=0 semantics).
    let _ = rustix::fs::fadvise(file.as_fd(), 0, None, rustix::fs::Advice::Sequential);
}

#[cfg(target_os = "linux")]
fn open_direct_linux(path: &Path, write: bool) -> DbResult<File> {
    use std::os::unix::fs::OpenOptionsExt;
    let flag = rustix::fs::OFlags::DIRECT.bits() as i32;
    let mut opts = OpenOptions::new();
    opts.read(true).custom_flags(flag);
    if write {
        opts.create(true).write(true).truncate(false);
    }
    Ok(opts.open(path)?)
}

#[cfg(target_os = "linux")]
fn is_einval(e: &DbError) -> bool {
    matches!(
        e,
        DbError::Io(io)
            if io.raw_os_error() == Some(rustix::io::Errno::INVAL.raw_os_error())
    )
}

/// Read `len` bytes from an encrypted data file at `offset`, decrypting pages.
///
/// Same alignment logic as `pread_value`, plus per-page AES-256-GCM decryption
/// using tags from the corresponding `.tags` file.
#[cfg(feature = "encryption")]
pub fn pread_value_encrypted(
    file: &File,
    tag_file: &crate::io::tags::TagFile,
    cipher: &crate::crypto::PageCipher,
    file_id: u32,
    offset: u64,
    len: usize,
) -> DbResult<Vec<u8>> {
    let sector_size: u64 = 4096;
    let aligned_offset = offset & !(sector_size - 1);
    let diff = (offset - aligned_offset) as usize;
    let aligned_len = (diff + len + sector_size as usize - 1) & !(sector_size as usize - 1);

    let mut buf = AlignedBuf::zeroed(aligned_len);

    // F9: require the *full* pages, not just `diff + len`. Every page below is
    // decrypted whole, so a read that stops mid-page leaves the tail zeroed and
    // fails GCM authentication — reporting a truncated file as corruption.
    // Encrypted files are always page-padded, so these bytes exist on disk.
    aligned_read_exact(file, &mut buf, aligned_offset, aligned_len)?;

    // Decrypt each 4096-byte page
    let start_page = aligned_offset / sector_size;
    let num_pages = aligned_len / sector_size as usize;
    let tags = tag_file.read_tags(start_page, num_pages)?;
    #[allow(clippy::needless_range_loop)]
    for i in 0..num_pages {
        let page_start = i * sector_size as usize;
        let page = &mut buf[page_start..page_start + sector_size as usize];
        cipher.decrypt_page(file_id, start_page + i as u64, page, &tags[i])?;
    }

    let mut result = vec![0u8; len];
    result.copy_from_slice(&buf[diff..diff + len]);
    Ok(result)
}

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

    #[test]
    fn aligned_read_fills_full_range_across_short_reads() {
        let dir = std::env::temp_dir().join(format!("armdb_aligned_{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("a.data");

        // Write 3 pages of known bytes via buffered handle.
        let f = open_bulk_write(&path).unwrap();
        let mut data = vec![0u8; 4096 * 3];
        for (i, b) in data.iter_mut().enumerate() {
            *b = (i % 251) as u8;
        }
        pwrite_at(&f, &data, 0).unwrap();
        fsync(&f).unwrap();
        drop(f);

        // Read an unaligned range that spans page boundaries.
        let r = open_bulk_read(&path).unwrap();
        let got = pread_value(&r, 100, 4096 * 2 + 7).unwrap();
        assert_eq!(&got[..], &data[100..100 + 4096 * 2 + 7]);

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

    /// Regression: a read that needs more bytes than the file holds, where the
    /// available tail is not sector-aligned, must return `UnexpectedEof` rather
    /// than spinning forever. The retry re-aligns the cursor down to a sector
    /// boundary, so without a no-progress guard `aligned_read_exact` would
    /// re-read the same partial tail indefinitely.
    #[test]
    fn aligned_read_terminates_on_short_file() {
        let dir = std::env::temp_dir().join(format!("armdb_aligned_eof_{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("short.data");

        // 5000 bytes: one full sector (4096) plus a non-aligned 904-byte tail.
        let f = open_bulk_write(&path).unwrap();
        pwrite_at(&f, &vec![0xABu8; 5000], 0).unwrap();
        fsync(&f).unwrap();
        drop(f);

        // Ask for 8000 bytes from offset 0 — more than the file holds. Must error,
        // not hang.
        let r = open_bulk_read(&path).unwrap();
        let res = pread_value(&r, 0, 8000);
        assert!(
            matches!(&res, Err(DbError::Io(e)) if e.kind() == std::io::ErrorKind::UnexpectedEof),
            "expected UnexpectedEof on short file, got {res:?}"
        );

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

    #[test]
    fn probe_and_bulk_open_roundtrip() {
        let dir = std::env::temp_dir().join(format!("armdb_direct_probe_{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();

        // Probe never panics and returns a bool; on tmpfs/sandbox it is false.
        let _supported = probe_direct_io(&dir);

        // Bulk open + write + read works regardless of O_DIRECT support.
        let path = dir.join("t.data");
        let f = open_bulk_write(&path).unwrap();
        pwrite_at(&f, b"hello", 0).unwrap();
        fsync(&f).unwrap();
        drop(f);

        let r = open_bulk_read(&path).unwrap();
        let got = pread_value(&r, 0, 5).unwrap();
        assert_eq!(&got, b"hello");

        // Serving open with direct=false behaves like buffered.
        let _s = open_serving_read(&path, false).unwrap();
        fadvise_dontneed(&r, 0, 5); // no-op off Linux; must not error/panic

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

    /// F9: `pread_value_encrypted` decrypts whole 4096-byte pages, so it must
    /// *read* whole pages. Requiring only `diff + len` lets a read stop mid-page
    /// and leaves the page tail zeroed, which then fails GCM authentication —
    /// reporting a truncated file as "decryption/authentication failed".
    #[cfg(feature = "encryption")]
    #[test]
    fn pread_value_encrypted_reports_short_file_as_eof_not_auth_failure() {
        use crate::crypto::PageCipher;
        use crate::io::tags::{TagFile, tags_path_for_data};

        let dir = std::env::temp_dir().join(format!("armdb_f9_{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("f9.data");
        let _ = std::fs::remove_file(&path);
        let _ = std::fs::remove_file(tags_path_for_data(&path));

        let cipher = PageCipher::new(&[7u8; 32]).unwrap();
        let file_id = 3u32;

        // Two encrypted pages, tags written alongside.
        let mut plain = vec![0u8; 4096 * 2];
        for (i, b) in plain.iter_mut().enumerate() {
            *b = (i % 251) as u8;
        }
        let mut tags = Vec::new();
        for page in 0..2u64 {
            let start = page as usize * 4096;
            let tag = cipher
                .encrypt_page(file_id, page, &mut plain[start..start + 4096])
                .unwrap();
            tags.push(tag);
        }
        let tag_file = TagFile::open_write(&tags_path_for_data(&path)).unwrap();
        tag_file.write_tags(0, &tags).unwrap();
        tag_file.sync().unwrap();

        let w = open_bulk_write(&path).unwrap();
        pwrite_at(&w, &plain, 0).unwrap();
        fsync(&w).unwrap();
        drop(w);

        // Truncate mid-way through the second page: exactly `diff + len` bytes
        // exist for the read below, but the page that covers them does not.
        let (offset, len) = (100u64, 5000usize);
        std::fs::File::options()
            .write(true)
            .open(&path)
            .unwrap()
            .set_len(offset + len as u64)
            .unwrap();

        let r = open_bulk_read(&path).unwrap();
        let tag_file = TagFile::open_read(&tags_path_for_data(&path)).unwrap();
        let err = pread_value_encrypted(&r, &tag_file, &cipher, file_id, offset, len)
            .expect_err("the page covering this range is not fully on disk");
        match err {
            DbError::Io(e) => assert_eq!(e.kind(), std::io::ErrorKind::UnexpectedEof),
            other => panic!("expected UnexpectedEof, got {other:?}"),
        }

        let _ = std::fs::remove_file(&path);
        let _ = std::fs::remove_file(tags_path_for_data(&path));
    }
}