//! C3 — the cell's address space captured as a recorded-VA memory dump for the
//! native-execution VM (Stage 1; see the Stage-1 spec doc).
//!
//! The warm-restore snapshot must reproduce a guest cell's memory **at the same
//! virtual addresses** in a fresh process (pointers, RIP, RSP, the rewrite-patched
//! code — all assume their original VAs). The adversarial review's binding
//! correction: the cell's address space is **not** ASLR-free by construction (the
//! stack, brk/heap arena, and ring table are `mmap(NULL)`, host-ASLR-placed — the
//! cell's own `/proc/self/maps` infoleak filter is proof ASLR is live), so we must
//! **record each VMA's actual VA** from `/proc/<pid>/maps` and `MAP_FIXED` it back
//! (CRIU-style), under `personality(ADDR_NO_RANDOMIZE)` for the restorer. This is
//! NOT a free fall-out of a single managed mmap.
//!
//! This module is **pure** (no supervisor hot-path edits — the capture/restore
//! orchestration is C4): it parses maps, decodes prot, copies VMA contents into a
//! page-aligned RAM blob through a caller-supplied `read_at` closure (in-process
//! `memcpy` in tests; cross-process `process_vm_readv` / `pread(/proc/<pid>/mem)`
//! at runtime), and `MAP_FIXED`-restores a region from that blob. Thread register
//! files (`CAPTURED[18]` + `%fs` at the SENTINEL checkpoint) are gathered by C4
//! and stored as `ThreadRegs`; this module owns only the *memory*.

use std::io::{self, Write};

use super::state_snap::VmaRecord;

/// The guest arena ceiling: only VMAs strictly below this are the guest's own
/// (the loader/supervisor/vdso live at/above it — never dumped). Mirrors
/// `WINDOW_FLOOR` in `sentry::mod`; kept as its own const so this pure module does
/// not depend on the (cfg-heavy) parent for a single number.
pub const ARENA_FLOOR: u64 = 0x4000_0000;

/// Page size assumed by the dump layout (snapshot writers page-align every VMA's
/// offset in the RAM blob so `mmap` can map it — `mmap` requires a page-aligned
/// file offset). x86_64 base pages.
pub const PAGE: u64 = 4096;

/// One parsed `/proc/<pid>/maps` entry kept for the dump (start below the arena
/// floor). `prot` is the `PROT_*` bitmask decoded from the `rwx` perms.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct MapEntry {
    pub va: u64,
    pub len: u64,
    pub prot: u32,
}

/// Decode a `/proc/<pid>/maps` perms field (`"rwxp"` / `"r-xp"` / …) into a
/// `PROT_*` bitmask. The 4th char (`p`/`s`) is the share mode, handled at restore
/// (guest segments are all private), not part of prot.
pub fn perms_to_prot(perms: &[u8]) -> u32 {
    let mut prot = 0u32;
    if perms.first() == Some(&b'r') {
        prot |= libc::PROT_READ as u32;
    }
    if perms.get(1) == Some(&b'w') {
        prot |= libc::PROT_WRITE as u32;
    }
    if perms.get(2) == Some(&b'x') {
        prot |= libc::PROT_EXEC as u32;
    }
    prot
}

/// Parse `/proc/<pid>/maps` text, keeping only the guest's own mappings: those
/// whose start VA is strictly below `floor` (exactly the sentry's existing maps
/// filter). Each kept line `START-END PERMS OFFSET DEV INODE [PATH]` becomes a
/// [`MapEntry`]. Malformed lines are skipped (never panic on guest-influenced
/// text). Results are returned in file order (ascending VA, as the kernel emits).
pub fn parse_maps(raw: &[u8], floor: u64) -> Vec<MapEntry> {
    let mut out = Vec::new();
    for line in raw.split(|&b| b == b'\n') {
        if line.is_empty() {
            continue;
        }
        // START-END PERMS ...
        let dash = match line.iter().position(|&b| b == b'-') {
            Some(d) => d,
            None => continue,
        };
        let sp = match line[dash + 1..].iter().position(|&b| b == b' ') {
            Some(s) => dash + 1 + s,
            None => continue,
        };
        let start = match std::str::from_utf8(&line[..dash])
            .ok()
            .and_then(|s| u64::from_str_radix(s, 16).ok())
        {
            Some(v) => v,
            None => continue,
        };
        let end = match std::str::from_utf8(&line[dash + 1..sp])
            .ok()
            .and_then(|s| u64::from_str_radix(s, 16).ok())
        {
            Some(v) => v,
            None => continue,
        };
        if start >= floor || end <= start {
            continue;
        }
        // perms is the next whitespace-delimited field after END.
        let after = &line[sp + 1..];
        let pe = after.iter().position(|&b| b == b' ').unwrap_or(after.len());
        let prot = perms_to_prot(&after[..pe]);
        out.push(MapEntry {
            va: start,
            len: end - start,
            prot,
        });
    }
    out
}

/// Round `n` up to the next multiple of [`PAGE`].
#[inline]
fn page_up(n: u64) -> u64 {
    (n + PAGE - 1) & !(PAGE - 1)
}

/// Capture the contents of `entries` into a page-aligned RAM blob written to
/// `ram`, returning the [`VmaRecord`]s that index into it. Each record's
/// `file_off` is the (page-aligned) offset of that VMA's bytes **within the RAM
/// blob** — the meaning of `file_off` for a snapshot dump (guest VMAs are
/// anonymous; the blob *is* their backing store on restore). `read_at(va, buf)`
/// fills `buf` with the bytes at guest VA `va`, returning `false` on a read
/// failure (which becomes an `Err`): in tests it is a `memcpy` from this process;
/// at runtime C4 passes a `process_vm_readv` / `pread(/proc/<pid>/mem)` reader for
/// the (quiesced) cell. The writer must already be at a page-aligned position
/// (callers start the RAM blob on a page boundary).
pub fn capture_regions(
    entries: &[MapEntry],
    read_at: impl Fn(u64, &mut [u8]) -> bool,
    ram: &mut impl Write,
) -> io::Result<Vec<VmaRecord>> {
    let mut recs = Vec::with_capacity(entries.len());
    let mut off = 0u64; // running, page-aligned offset into the blob
    let mut buf = vec![0u8; PAGE as usize];
    for e in entries {
        let rec = VmaRecord {
            va_start: e.va,
            len: e.len,
            prot: e.prot,
            file_off: off,
        };
        // Copy the region a page at a time (bounded scratch; large VMAs stream).
        let mut done = 0u64;
        while done < e.len {
            let chunk = ((e.len - done) as usize).min(PAGE as usize);
            let dst = &mut buf[..chunk];
            if !read_at(e.va + done, dst) {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    format!("memimage: read_at failed @ {:#x}+{:#x}", e.va, done),
                ));
            }
            ram.write_all(dst)?;
            done += chunk as u64;
        }
        // Pad the blob to a page boundary so the NEXT VMA's file_off is
        // page-aligned (mmap requires a page-aligned file offset).
        let padded = page_up(e.len);
        if padded > e.len {
            let pad = vec![0u8; (padded - e.len) as usize];
            ram.write_all(&pad)?;
        }
        off += padded;
        recs.push(rec);
    }
    Ok(recs)
}

/// `MAP_PRIVATE | MAP_FIXED`-map one VMA's bytes from the dump `file` at its
/// recorded VA, so the restored process sees identical-VA memory (CoW: guest
/// writes never touch the dump). `ram_base` is the byte offset of the RAM blob
/// within `file` (0 when the dump file IS the blob); the mapped file offset is
/// `ram_base + rec.file_off` and must be page-aligned (it is, by construction).
/// The region is bounds-checked against the file length first (a corrupt
/// offset/len would otherwise SIGBUS the restorer on first touch).
///
/// # Safety
/// `MAP_FIXED` silently unmaps anything already at `[rec.va_start, +len)`. The
/// caller (C4) must run in a fresh restorer whose own image lives ABOVE the arena
/// floor and under `personality(ADDR_NO_RANDOMIZE)`, so the guest arena is free.
pub unsafe fn restore_region(
    rec: &VmaRecord,
    file: &std::fs::File,
    ram_base: u64,
) -> io::Result<*mut u8> {
    use std::os::fd::AsRawFd;
    let file_off = ram_base
        .checked_add(rec.file_off)
        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "memimage: file_off overflow"))?;
    let file_len = file.metadata()?.len();
    if !crate::snapshot_frame::ram_region_within(file_len, file_off, rec.len) {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "memimage: VMA [{:#x},+{:#x}) at blob off {file_off} exceeds dump len {file_len}",
                rec.va_start, rec.len
            ),
        ));
    }
    if file_off % PAGE != 0 || rec.va_start % PAGE != 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "memimage: VA and blob offset must be page-aligned for MAP_FIXED",
        ));
    }
    let ptr = unsafe {
        libc::mmap(
            rec.va_start as *mut libc::c_void,
            rec.len as usize,
            rec.prot as libc::c_int,
            libc::MAP_PRIVATE | libc::MAP_FIXED,
            file.as_raw_fd(),
            file_off as libc::off_t,
        )
    };
    if ptr == libc::MAP_FAILED {
        return Err(io::Error::last_os_error());
    }
    Ok(ptr as *mut u8)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::os::fd::FromRawFd;

    /// A RAM-backed `File` (no `tempfile` dep, no disk): `memfd_create` supports
    /// write, `metadata().len()`, and `mmap` — everything restore_region needs.
    fn memfd() -> std::fs::File {
        let fd = unsafe { libc::memfd_create(b"nev-memimg\0".as_ptr() as *const libc::c_char, 0) };
        assert!(fd >= 0, "memfd_create: {}", std::io::Error::last_os_error());
        unsafe { std::fs::File::from_raw_fd(fd) }
    }

    #[test]
    fn parse_maps_filters_to_arena_and_decodes_prot() {
        // A realistic mix: guest anon segments below the floor + the supervisor's
        // own code/stack/vdso at/above it (which must be dropped).
        let raw = b"\
00400000-00402000 r-xp 00000000 00:00 0                          \n\
00402000-00403000 rw-p 00000000 00:00 0                          \n\
3fffe000-40000000 rw-p 00000000 00:00 0                          [stack-ish]\n\
40000000-40010000 r-xp 00000000 00:00 0                          \n\
7f0000000000-7f0000001000 r--p 00000000 00:00 0                  [vvar]\n\
ffffffffff600000-ffffffffff601000 --xp 00000000 00:00 0          [vsyscall]\n";
        let m = parse_maps(raw, ARENA_FLOOR);
        assert_eq!(m.len(), 3, "only the 3 sub-floor VMAs are kept");
        assert_eq!(
            m[0],
            MapEntry {
                va: 0x400000,
                len: 0x2000,
                prot: (libc::PROT_READ | libc::PROT_EXEC) as u32
            }
        );
        assert_eq!(
            m[1],
            MapEntry {
                va: 0x402000,
                len: 0x1000,
                prot: (libc::PROT_READ | libc::PROT_WRITE) as u32
            }
        );
        assert_eq!(m[2].va, 0x3fffe000);
        // The VMA starting exactly at the floor (0x40000000) is excluded.
        assert!(m.iter().all(|e| e.va < ARENA_FLOOR));
    }

    #[test]
    fn parse_maps_skips_malformed_lines() {
        let raw = b"garbage\n00400000-00401000 rw-p 0 0 0\nnodash here\n";
        let m = parse_maps(raw, ARENA_FLOOR);
        assert_eq!(m.len(), 1);
        assert_eq!(m[0].va, 0x400000);
    }

    // The headline test: capture a live anon region into a blob, drop it, and
    // MAP_FIXED-restore it at the IDENTICAL VA from the blob — proving the
    // recorded-VA round-trip (the heart of warm-restore).
    #[test]
    fn recorded_va_round_trip_map_fixed() {
        let len = 2 * PAGE as usize;
        // Get a real anon mapping; the kernel picks the address A.
        let a = unsafe {
            libc::mmap(
                std::ptr::null_mut(),
                len,
                libc::PROT_READ | libc::PROT_WRITE,
                libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
                -1,
                0,
            )
        };
        assert_ne!(a, libc::MAP_FAILED, "mmap anon");
        let a = a as u64;
        // Fill with a position-dependent pattern.
        let pat = |i: usize| (i as u8).wrapping_mul(31).wrapping_add(7);
        unsafe {
            for i in 0..len {
                *((a as *mut u8).add(i)) = pat(i);
            }
        }

        // Capture: read_at = memcpy from our own memory; blob → a memfd.
        let entries = [MapEntry {
            va: a,
            len: len as u64,
            prot: (libc::PROT_READ | libc::PROT_WRITE) as u32,
        }];
        let mut file = memfd();
        let recs = capture_regions(
            &entries,
            |va, buf| {
                unsafe {
                    std::ptr::copy_nonoverlapping(va as *const u8, buf.as_mut_ptr(), buf.len())
                };
                true
            },
            &mut file,
        )
        .expect("capture");
        assert_eq!(recs.len(), 1);
        assert_eq!(recs[0].va_start, a);
        assert_eq!(recs[0].file_off, 0);

        // Drop the original mapping, then restore from the blob at the SAME VA.
        unsafe { assert_eq!(libc::munmap(a as *mut libc::c_void, len), 0, "munmap") };
        let restored = unsafe { restore_region(&recs[0], &file, 0).expect("restore") };
        assert_eq!(restored as u64, a, "restored at the identical VA");
        // Contents survived byte-for-byte.
        unsafe {
            for i in 0..len {
                assert_eq!(
                    *((a as *const u8).add(i)),
                    pat(i),
                    "byte {i} mismatch after restore"
                );
            }
            libc::munmap(a as *mut libc::c_void, len);
        }
    }

    #[test]
    fn restore_rejects_unaligned_and_oob() {
        let mut file = memfd();
        file.write_all(&vec![0u8; PAGE as usize]).unwrap();
        // Unaligned VA.
        let bad_va = VmaRecord {
            va_start: 0x1234,
            len: PAGE,
            prot: libc::PROT_READ as u32,
            file_off: 0,
        };
        assert!(unsafe { restore_region(&bad_va, &file, 0) }.is_err());
        // Region past EOF (len exceeds the 1-page dump).
        let oob = VmaRecord {
            va_start: 0,
            len: 4 * PAGE,
            prot: libc::PROT_READ as u32,
            file_off: 0,
        };
        assert!(unsafe { restore_region(&oob, &file, 0) }.is_err());
    }

    #[test]
    fn capture_surfaces_read_failure() {
        let entries = [MapEntry {
            va: 0x1000,
            len: PAGE,
            prot: libc::PROT_READ as u32,
        }];
        let mut blob = Vec::new();
        let r = capture_regions(&entries, |_, _| false, &mut blob);
        assert!(
            r.is_err(),
            "a failing read_at must surface as Err, not silent zero-fill"
        );
    }
}
