supermachine 0.2.0

// Snapshot/restore for the VMM. Captures the VM state needed to
// resume in a fresh process: per-vCPU registers (GP, SIMD, sysregs),
// per-vCPU ICC + redistributor state, the opaque GIC blob, vtimer
// offset, and full guest RAM.
//
// V2 (this revision): captures virtio mirror state — per-MMIO-device
// negotiated features, status, and per-queue ring addresses + cursors,
// plus the vsock muxer's TSI listener registry. In-flight TCP/UDP
// connections and pending handshakes are intentionally dropped (peer
// TCPs see RST; clients retry). Listeners are re-bound on restore
// with fresh ephemeral host ports — the guest only knows its vm_port,
// so the change is transparent.
//
// Wire format ("SMSNAP\x06\x00", little-endian):
//   header 64 B:
//     [0..8]   magic
//     [8..16]  version (u64)
//     [16..24] captured_mach_time (u64)
//     [24..32] captured_cntvct (u64)
//     [32..40] memory_bytes (u64)
//     [40..48] gic_blob_len (u64)
//     [48..52] n_vcpus (u32)
//     [52..56] reserved (u32)
//     [56..64] ram_gpa (u64)
//   per_vcpu × n_vcpus:
//     vtimer_offset (u64),
//     gp_n simd_n sys_n icc_n redist_n  (u32×5)
//     gp_regs (reg_id u32, val u64)        × gp_n
//     simd_regs (reg_id u32, val u128)     × simd_n
//     sys_regs (reg_id u32, val u64)       × sys_n
//     icc_regs (reg_id u32, val u64)       × icc_n
//     redist (off u32, val u64)            × redist_n
//   then RAM: memory_bytes raw bytes
//
// Memory and metadata live in the same file for v1 simplicity.

#![cfg(all(target_os = "macos", target_arch = "aarch64"))]

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

use applevisor_sys as av;

use crate::devices::virtio::mmio::{MmioSnapshot, QueueSnapshot};
use crate::devices::virtio::vsock::muxer::TsiListenerSnapshot;
use crate::hvf::{self, Vcpu};
use crate::vmm::vstate::MicroVm;

const SNAPSHOT_MAGIC: [u8; 8] = *b"SMSNAP\x08\x00";
const SNAPSHOT_VERSION: u64 = 8;
/// Page size to align the RAM region within the snapshot file. macOS
/// on Apple Silicon uses 16 KiB pages; aligning here lets `--cow-restore`
/// mmap the region with `MAP_PRIVATE` directly out of the file.
const RAM_PAGE_ALIGN: u64 = 16384;
const SPARSE_RAM_CHUNK: usize = 64 * 1024;

#[derive(Default, Clone)]
pub struct PerVcpuState {
    pub gp_regs: Vec<(u32, u64)>,     // hv_reg_t as u32
    pub simd_regs: Vec<(u32, u128)>,  // hv_simd_fp_reg_t as u32
    pub sys_regs: Vec<(u32, u64)>,    // hv_sys_reg_t as u32
    pub icc_regs: Vec<(u32, u64)>,    // hv_gic_icc_reg_t as u32
    pub redist_regs: Vec<(u32, u64)>, // GICR offset
    pub vtimer_offset: u64,
}

/// Per-VM virtio device state. Length of `mmio` matches the bus
/// device order in the file (vsock first, blk after).
#[derive(Default)]
pub struct VirtioSnapshot {
    pub mmio: Vec<MmioSnapshot>,
    pub vsock_listeners: Vec<TsiListenerSnapshot>,
}

pub struct Snapshot {
    pub captured_mach_time: u64,
    pub captured_cntvct: u64,
    pub ram_gpa: u64,
    pub memory: Vec<u8>,
    pub gic_blob: Vec<u8>,
    pub per_vcpu: Vec<PerVcpuState>,
    pub virtio: VirtioSnapshot,
}

#[derive(Default, Clone, Copy)]
pub struct SnapshotWriteStats {
    pub ram_bytes: u64,
    pub ram_data_bytes: u64,
    pub ram_zero_bytes: u64,
}

#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
pub struct SnapshotRestoreTimings {
    pub ram_copy_us: u128,
    pub gic_restore_us: u128,
    pub vcpu_restore_us: u128,
    pub vtimer_offset_us: u128,
}

#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
pub struct SnapshotRestoreOptions {
    pub skip_gic_blob: bool,
}

// ----- Register enumerations -----

fn gp_reg_enum() -> Vec<av::hv_reg_t> {
    let mut out = Vec::with_capacity(37);
    let x0 = av::hv_reg_t::X0 as u32;
    for i in 0..=30u32 {
        // SAFETY: X0..X30 are 31 contiguous variants starting at X0.
        out.push(unsafe { std::mem::transmute::<u32, av::hv_reg_t>(x0 + i) });
    }
    out.push(av::hv_reg_t::FP);
    out.push(av::hv_reg_t::LR);
    out.push(av::hv_reg_t::PC);
    out.push(av::hv_reg_t::CPSR);
    out.push(av::hv_reg_t::FPCR);
    out.push(av::hv_reg_t::FPSR);
    out
}

fn simd_reg_enum() -> Vec<av::hv_simd_fp_reg_t> {
    let q0 = av::hv_simd_fp_reg_t::Q0 as u32;
    (0..32u32)
        // SAFETY: Q0..Q31 are 32 contiguous variants starting at Q0.
        .map(|i| unsafe { std::mem::transmute::<u32, av::hv_simd_fp_reg_t>(q0 + i) })
        .collect()
}

fn sys_reg_enum() -> Vec<av::hv_sys_reg_t> {
    use av::hv_sys_reg_t::*;
    vec![
        MPIDR_EL1,
        SCTLR_EL1,
        CPACR_EL1,
        TCR_EL1,
        TTBR0_EL1,
        TTBR1_EL1,
        MAIR_EL1,
        AMAIR_EL1,
        VBAR_EL1,
        CONTEXTIDR_EL1,
        TPIDR_EL1,
        SPSR_EL1,
        ELR_EL1,
        SP_EL0,
        SP_EL1,
        ESR_EL1,
        FAR_EL1,
        PAR_EL1,
        TPIDR_EL0,
        TPIDRRO_EL0,
        CNTKCTL_EL1,
        CSSELR_EL1,
        MDSCR_EL1,
        // Pointer-auth keys (CONFIG_ARM64_PTR_AUTH guests OOPS without these).
        APIAKEYLO_EL1,
        APIAKEYHI_EL1,
        APIBKEYLO_EL1,
        APIBKEYHI_EL1,
        APDAKEYLO_EL1,
        APDAKEYHI_EL1,
        APDBKEYLO_EL1,
        APDBKEYHI_EL1,
        APGAKEYLO_EL1,
        APGAKEYHI_EL1,
        // Vtimer CTL + deadline (HVF DOES accept these despite earlier rumours).
        CNTV_CTL_EL0,
        CNTV_CVAL_EL0,
        CNTP_CTL_EL0,
        CNTP_CVAL_EL0,
        // CNTVOFF_EL2 is captured via hv_vcpu_get_vtimer_offset, not as a sysreg.
    ]
}

fn icc_reg_enum() -> Vec<av::hv_gic_icc_reg_t> {
    use av::hv_gic_icc_reg_t::*;
    vec![
        PMR_EL1,
        BPR0_EL1,
        BPR1_EL1,
        AP0R0_EL1,
        AP1R0_EL1,
        RPR_EL1,
        CTLR_EL1,
        SRE_EL1,
        IGRPEN0_EL1,
        IGRPEN1_EL1,
    ]
}

/// Per-vCPU redistributor offsets QEMU v11 captures explicitly because
/// the opaque blob doesn't cover them. From hw/intc/arm_gicv3_hvf.c.
fn redist_reg_offsets() -> Vec<u32> {
    let mut v = Vec::with_capacity(11);
    v.push(0x10080); // GICR_IGROUPR0
    v.push(0x10100); // GICR_ISENABLER0
    v.push(0x10C04); // GICR_ICFGR1
    v.push(0x10200); // GICR_ISPENDR0  (vtimer PPI 27 lives here)
    v.push(0x10300); // GICR_ISACTIVER0
    for n in 0..8u32 {
        v.push(0x10400 + 4 * n); // GICR_IPRIORITYR0..7
    }
    v
}

#[link(name = "System", kind = "framework")]
extern "C" {
    fn mach_absolute_time() -> u64;
}

// ----- Capture -----

pub fn capture_vcpu_state(vcpu: &Vcpu) -> hvf::Result<PerVcpuState> {
    let mut gp_regs = Vec::with_capacity(37);
    for r in gp_reg_enum() {
        gp_regs.push((r as u32, vcpu.get_reg(r)?));
    }
    let mut simd_regs = Vec::with_capacity(32);
    for r in simd_reg_enum() {
        simd_regs.push((r as u32, vcpu.get_simd_fp_reg(r)?));
    }
    let mut sys_regs = Vec::new();
    for r in sys_reg_enum() {
        // Best-effort: some sysregs may be unreadable in certain
        // states; skip rather than fail the whole capture.
        if let Ok(v) = vcpu.get_sys_reg(r) {
            sys_regs.push((r as u32, v));
        }
    }
    let mut icc_regs = Vec::new();
    for r in icc_reg_enum() {
        if let Ok(v) = vcpu.get_icc_reg(r) {
            icc_regs.push((r as u32, v));
        }
    }
    let mut redist_regs = Vec::new();
    for off in redist_reg_offsets() {
        // SAFETY: offsets are valid GICR register enum variants.
        let reg: av::hv_gic_redistributor_reg_t = unsafe { std::mem::transmute(off) };
        if let Ok(v) = vcpu.get_redist_reg(reg) {
            redist_regs.push((off, v));
        }
    }
    let vtimer_offset = vcpu.get_vtimer_offset()?;
    Ok(PerVcpuState {
        gp_regs,
        simd_regs,
        sys_regs,
        icc_regs,
        redist_regs,
        vtimer_offset,
    })
}

pub fn capture_snapshot(vm: &MicroVm, virtio: VirtioSnapshot) -> hvf::Result<Snapshot> {
    let per0 = capture_vcpu_state(&vm.vcpu)?;
    let gic_blob = hvf::gic_state_capture()?;
    let mut memory = vec![0u8; vm.ram_size];
    // SAFETY: ram_host is vm.ram_size bytes.
    unsafe {
        std::ptr::copy_nonoverlapping(vm.ram_host, memory.as_mut_ptr(), vm.ram_size);
    }
    let captured_mach_time = unsafe { mach_absolute_time() };
    let captured_cntvct = captured_mach_time.wrapping_sub(per0.vtimer_offset);
    Ok(Snapshot {
        captured_mach_time,
        captured_cntvct,
        ram_gpa: vm.ram_gpa,
        memory,
        gic_blob,
        per_vcpu: vec![per0],
        virtio,
    })
}

// ----- Restore -----

/// Restore per-vCPU state. Vtimer offset is the CALLER's
/// responsibility (one coherent value across all vCPUs).
pub fn restore_vcpu_state(vcpu: &Vcpu, st: &PerVcpuState) -> hvf::Result<()> {
    // 1. sysregs first (MMU, exception state, pointer-auth, vtimer).
    //    Some may be RO; only fail loudly for MMU-critical writes.
    use av::hv_sys_reg_t as S;
    let critical = |id: u32| {
        let r: S = unsafe { std::mem::transmute(id) };
        matches!(
            r,
            S::SCTLR_EL1 | S::TCR_EL1 | S::TTBR0_EL1 | S::TTBR1_EL1 | S::MAIR_EL1 | S::VBAR_EL1
        )
    };
    for (id, v) in &st.sys_regs {
        // SAFETY: id originated from sys_reg_enum() variants.
        let r: S = unsafe { std::mem::transmute(*id) };
        if let Err(e) = vcpu.set_sys_reg(r, *v) {
            if critical(*id) {
                return Err(e);
            }
        }
    }

    // 2. ICC: SRE_EL1 first (gates subsequent ICC writes), then
    //    everything except IGRPEN0/1, then IGRPEN0/1 last (unmask
    //    delivery only after PMR/BPR/AP state is in place).
    use av::hv_gic_icc_reg_t as I;
    let icc_find = |want: I| -> Option<u64> {
        st.icc_regs.iter().find_map(|(id, v)| {
            // SAFETY: id from icc_reg_enum.
            let r: I = unsafe { std::mem::transmute(*id) };
            (r == want).then_some(*v)
        })
    };
    if let Some(v) = icc_find(I::SRE_EL1) {
        let _ = vcpu.set_icc_reg(I::SRE_EL1, v);
    }
    for (id, v) in &st.icc_regs {
        // SAFETY: id from icc_reg_enum.
        let r: I = unsafe { std::mem::transmute(*id) };
        match r {
            I::SRE_EL1 | I::IGRPEN0_EL1 | I::IGRPEN1_EL1 => continue,
            _ => {
                let _ = vcpu.set_icc_reg(r, *v);
            }
        }
    }
    if let Some(v) = icc_find(I::IGRPEN0_EL1) {
        let _ = vcpu.set_icc_reg(I::IGRPEN0_EL1, v);
    }
    if let Some(v) = icc_find(I::IGRPEN1_EL1) {
        let _ = vcpu.set_icc_reg(I::IGRPEN1_EL1, v);
    }

    // 3. SIMD/FP regs.
    use av::hv_simd_fp_reg_t as Q;
    for (id, v) in &st.simd_regs {
        // SAFETY: id from simd_reg_enum.
        let r: Q = unsafe { std::mem::transmute(*id) };
        vcpu.set_simd_fp_reg(r, *v)?;
    }

    // 4. Per-vCPU redistributor regs. Order from QEMU v11
    //    arm_gicv3_hvf.c: group/config/priority first, then CLEAR each
    //    enable/pending/active mask before SET (otherwise restore is
    //    OR of default + captured bits).
    let find_off = |off: u32| -> u64 {
        st.redist_regs
            .iter()
            .find_map(|(o, v)| (*o == off).then_some(*v))
            .unwrap_or(0)
    };
    let write_off = |off: u32, val: u64| -> hvf::Result<()> {
        // SAFETY: off comes from our own enumeration; transmute matches repr.
        let r: av::hv_gic_redistributor_reg_t = unsafe { std::mem::transmute(off) };
        vcpu.set_redist_reg(r, val)
    };
    write_off(0x10080, find_off(0x10080))?; // IGROUPR0
    write_off(0x10C04, find_off(0x10C04))?; // ICFGR1
    for n in 0..8u32 {
        write_off(0x10400 + 4 * n, find_off(0x10400 + 4 * n))?;
    }
    write_off(0x10180, 0xFFFF_FFFF)?; // ICENABLER0 clear
    write_off(0x10100, find_off(0x10100))?; // ISENABLER0 set
    write_off(0x10280, 0xFFFF_FFFF)?; // ICPENDR0 clear
    write_off(0x10200, find_off(0x10200))?; // ISPENDR0 set
    write_off(0x10380, 0xFFFF_FFFF)?; // ICACTIVER0 clear
    write_off(0x10300, find_off(0x10300))?; // ISACTIVER0 set

    // 5. Force vtimer mask off so HVF re-evaluates on next run.
    let _ = vcpu.set_vtimer_mask(false);

    // 6. Vtimer force-fire: if the captured vtimer was enabled and
    //    unmasked, set CVAL=0 and force-pend the PPI bit so the guest
    //    wakes immediately rather than waiting on a stale deadline.
    let cntv_ctl = st
        .sys_regs
        .iter()
        .find_map(|(id, v)| {
            let r: S = unsafe { std::mem::transmute(*id) };
            (r == S::CNTV_CTL_EL0).then_some(*v)
        })
        .unwrap_or(0);
    let enable = cntv_ctl & 1 != 0;
    let imask = cntv_ctl & 2 != 0;
    if enable && !imask {
        vcpu.set_sys_reg(S::CNTV_CVAL_EL0, 0)?;
        // GICR_ISPENDR0 bit 27 = vtimer PPI.
        write_off(0x10200, 1u64 << 27)?;
    }

    // 7. GP regs LAST (PC/CPSR finalize the vCPU).
    use av::hv_reg_t as R;
    for (id, v) in &st.gp_regs {
        // SAFETY: id from gp_reg_enum.
        let r: R = unsafe { std::mem::transmute(*id) };
        vcpu.set_reg(r, *v)?;
    }
    Ok(())
}

/// Restore the full snapshot into `vm`. RAM is memcpy'd in (for v1
/// we don't do CoW — that's a follow-up). Caller must have created
/// `vm` with the same RAM size as the snapshot (otherwise the memcpy
/// is silently truncated).
pub fn restore_snapshot(vm: &MicroVm, snap: &Snapshot) -> hvf::Result<()> {
    restore_snapshot_timed(vm, snap).map(|_| ())
}

pub fn restore_snapshot_timed(
    vm: &MicroVm,
    snap: &Snapshot,
) -> hvf::Result<SnapshotRestoreTimings> {
    restore_snapshot_timed_with_options(vm, snap, SnapshotRestoreOptions::default())
}

pub fn restore_snapshot_timed_with_options(
    vm: &MicroVm,
    snap: &Snapshot,
    options: SnapshotRestoreOptions,
) -> hvf::Result<SnapshotRestoreTimings> {
    let mut timings = SnapshotRestoreTimings::default();
    // 1. RAM. CoW-restored snapshots leave `memory` empty (pages are
    //    already mapped via mmap(MAP_PRIVATE) before MicroVm::new).
    if !snap.memory.is_empty() {
        let t0 = std::time::Instant::now();
        // SAFETY: ram_host is vm.ram_size bytes; we cap copy length.
        unsafe {
            std::ptr::copy_nonoverlapping(
                snap.memory.as_ptr(),
                vm.ram_host,
                vm.ram_size.min(snap.memory.len()),
            );
        }
        timings.ram_copy_us = t0.elapsed().as_micros();
    }
    // 2. GIC blob (covers distributor + per-PE pending/active).
    if !options.skip_gic_blob {
        let t0 = std::time::Instant::now();
        hvf::gic_state_restore(&snap.gic_blob)?;
        timings.gic_restore_us = t0.elapsed().as_micros();
    }
    // 3. Per-vCPU state.
    let boot_vcpu = snap.per_vcpu.first().ok_or(hvf::Error::Hv(-1))?;
    let t0 = std::time::Instant::now();
    restore_vcpu_state(&vm.vcpu, boot_vcpu)?;
    timings.vcpu_restore_us = t0.elapsed().as_micros();
    // 4. Coherent vtimer offset for all vCPUs (single value).
    let now = unsafe { mach_absolute_time() };
    let new_offset = now.wrapping_sub(snap.captured_cntvct);
    let t0 = std::time::Instant::now();
    vm.vcpu.set_vtimer_offset(new_offset)?;
    timings.vtimer_offset_us = t0.elapsed().as_micros();
    Ok(timings)
}

// ----- File serialization -----

#[derive(Debug)]
pub enum FileError {
    Io(std::io::Error),
    BadMagic,
    BadVersion(u64),
    Malformed(&'static str),
    Truncated,
}
impl From<std::io::Error> for FileError {
    fn from(e: std::io::Error) -> Self {
        Self::Io(e)
    }
}

fn le_u16(bytes: &[u8]) -> Result<u16, FileError> {
    let bytes = bytes.try_into().map_err(|_| FileError::Truncated)?;
    Ok(u16::from_le_bytes(bytes))
}

fn le_u32(bytes: &[u8]) -> Result<u32, FileError> {
    let bytes = bytes.try_into().map_err(|_| FileError::Truncated)?;
    Ok(u32::from_le_bytes(bytes))
}

fn le_u64(bytes: &[u8]) -> Result<u64, FileError> {
    let bytes = bytes.try_into().map_err(|_| FileError::Truncated)?;
    Ok(u64::from_le_bytes(bytes))
}

fn le_u128(bytes: &[u8]) -> Result<u128, FileError> {
    let bytes = bytes.try_into().map_err(|_| FileError::Truncated)?;
    Ok(u128::from_le_bytes(bytes))
}

pub fn save_to_file(path: &str, snap: &Snapshot) -> Result<(), FileError> {
    save_to_file_with_stats(path, snap).map(|_| ())
}

pub fn save_to_file_with_stats(
    path: &str,
    snap: &Snapshot,
) -> Result<SnapshotWriteStats, FileError> {
    use std::io::{Seek, SeekFrom};
    let mut f = std::fs::File::create(path)?;
    // Header (72 bytes). ram_offset is filled in after we know where
    // RAM ends up — see the seek-and-rewrite at the bottom.
    let mut hdr = Vec::with_capacity(72);
    hdr.extend_from_slice(&SNAPSHOT_MAGIC);
    hdr.extend_from_slice(&SNAPSHOT_VERSION.to_le_bytes());
    hdr.extend_from_slice(&snap.captured_mach_time.to_le_bytes());
    hdr.extend_from_slice(&snap.captured_cntvct.to_le_bytes());
    hdr.extend_from_slice(&(snap.memory.len() as u64).to_le_bytes());
    hdr.extend_from_slice(&(snap.gic_blob.len() as u64).to_le_bytes());
    hdr.extend_from_slice(&(snap.per_vcpu.len() as u32).to_le_bytes());
    hdr.extend_from_slice(&0u32.to_le_bytes()); // reserved
    hdr.extend_from_slice(&snap.ram_gpa.to_le_bytes());
    hdr.extend_from_slice(&0u64.to_le_bytes()); // ram_offset placeholder
    f.write_all(&hdr)?;
    f.write_all(&snap.gic_blob)?;
    for st in &snap.per_vcpu {
        let mut p = Vec::with_capacity(64 + st.gp_regs.len() * 12 + st.simd_regs.len() * 20);
        p.extend_from_slice(&st.vtimer_offset.to_le_bytes());
        p.extend_from_slice(&(st.gp_regs.len() as u32).to_le_bytes());
        p.extend_from_slice(&(st.simd_regs.len() as u32).to_le_bytes());
        p.extend_from_slice(&(st.sys_regs.len() as u32).to_le_bytes());
        p.extend_from_slice(&(st.icc_regs.len() as u32).to_le_bytes());
        p.extend_from_slice(&(st.redist_regs.len() as u32).to_le_bytes());
        for (id, v) in &st.gp_regs {
            p.extend_from_slice(&id.to_le_bytes());
            p.extend_from_slice(&v.to_le_bytes());
        }
        for (id, v) in &st.simd_regs {
            p.extend_from_slice(&id.to_le_bytes());
            p.extend_from_slice(&v.to_le_bytes());
        }
        for (id, v) in &st.sys_regs {
            p.extend_from_slice(&id.to_le_bytes());
            p.extend_from_slice(&v.to_le_bytes());
        }
        for (id, v) in &st.icc_regs {
            p.extend_from_slice(&id.to_le_bytes());
            p.extend_from_slice(&v.to_le_bytes());
        }
        for (off, v) in &st.redist_regs {
            p.extend_from_slice(&off.to_le_bytes());
            p.extend_from_slice(&v.to_le_bytes());
        }
        f.write_all(&p)?;
    }

    // Virtio section: u32 n_mmio, then per device:
    //   u32 driver_features[2], u32 status, u32 n_queues
    //   per queue: u16 size, u8 ready, _pad, u64 desc, u64 avail, u64 used,
    //              u16 last_avail_idx, u16 next_used_idx
    // then u32 n_listeners, per listener: u64 cid, u32 peer_port,
    //              u32 vm_port, u16 family, u16 socktype.
    let mut v = Vec::with_capacity(
        64 + snap.virtio.mmio.len() * 256 + snap.virtio.vsock_listeners.len() * 24,
    );
    v.extend_from_slice(&(snap.virtio.mmio.len() as u32).to_le_bytes());
    for m in &snap.virtio.mmio {
        v.extend_from_slice(&m.driver_features[0].to_le_bytes());
        v.extend_from_slice(&m.driver_features[1].to_le_bytes());
        v.extend_from_slice(&m.status.to_le_bytes());
        v.extend_from_slice(&m.interrupt_status.to_le_bytes());
        v.extend_from_slice(&(m.queues.len() as u32).to_le_bytes());
        for q in &m.queues {
            v.extend_from_slice(&q.size.to_le_bytes());
            v.push(if q.ready { 1 } else { 0 });
            v.push(0); // pad
            v.extend_from_slice(&q.desc_table.to_le_bytes());
            v.extend_from_slice(&q.avail_ring.to_le_bytes());
            v.extend_from_slice(&q.used_ring.to_le_bytes());
            v.extend_from_slice(&q.last_avail_idx.to_le_bytes());
            v.extend_from_slice(&q.next_used_idx.to_le_bytes());
        }
    }
    v.extend_from_slice(&(snap.virtio.vsock_listeners.len() as u32).to_le_bytes());
    for l in &snap.virtio.vsock_listeners {
        v.extend_from_slice(&l.cid.to_le_bytes());
        v.extend_from_slice(&l.peer_port.to_le_bytes());
        v.extend_from_slice(&l.vm_port.to_le_bytes());
        v.extend_from_slice(&l.family.to_le_bytes());
        v.extend_from_slice(&l.socktype.to_le_bytes());
    }
    f.write_all(&v)?;

    // Pad to RAM_PAGE_ALIGN, then write RAM. Stash the offset back
    // into the header so loaders (especially --cow-restore) know
    // where to mmap from.
    let cur = f.stream_position()? as u64;
    let pad = (RAM_PAGE_ALIGN - (cur % RAM_PAGE_ALIGN)) % RAM_PAGE_ALIGN;
    if pad > 0 {
        f.write_all(&vec![0u8; pad as usize])?;
    }
    let ram_offset = f.stream_position()? as u64;
    let stats = write_sparse_ram(&mut f, &snap.memory)?;
    f.seek(SeekFrom::Start(64))?;
    f.write_all(&ram_offset.to_le_bytes())?;
    Ok(stats)
}

fn write_sparse_ram(f: &mut std::fs::File, memory: &[u8]) -> Result<SnapshotWriteStats, FileError> {
    use std::io::{Seek, SeekFrom};

    let start = f.stream_position()? as u64;
    let mut pos = 0usize;
    let mut data_bytes = 0u64;
    let mut zero_bytes = 0u64;

    while pos < memory.len() {
        let run_is_zero =
            is_zero_chunk(&memory[pos..(pos + (memory.len() - pos).min(SPARSE_RAM_CHUNK))]);
        let run_start = pos;
        pos += (memory.len() - pos).min(SPARSE_RAM_CHUNK);
        while pos < memory.len() {
            let next_len = (memory.len() - pos).min(SPARSE_RAM_CHUNK);
            let next_is_zero = is_zero_chunk(&memory[pos..pos + next_len]);
            if next_is_zero != run_is_zero {
                break;
            }
            pos += next_len;
        }

        let run_len = pos - run_start;
        if run_is_zero {
            f.seek(SeekFrom::Current(run_len as i64))?;
            zero_bytes += run_len as u64;
        } else {
            f.write_all(&memory[run_start..pos])?;
            data_bytes += run_len as u64;
        }
    }

    let logical_end = start + memory.len() as u64;
    f.set_len(logical_end)?;
    f.seek(SeekFrom::Start(logical_end))?;
    Ok(SnapshotWriteStats {
        ram_bytes: memory.len() as u64,
        ram_data_bytes: data_bytes,
        ram_zero_bytes: zero_bytes,
    })
}

fn is_zero_chunk(chunk: &[u8]) -> bool {
    chunk.iter().all(|b| *b == 0)
}

pub fn load_from_file(path: &str) -> Result<Snapshot, FileError> {
    load_from_file_inner(path, /* skip_ram = */ false).map(|(snap, _, _)| snap)
}

/// Like `load_from_file` but skips RAM bytes (sets `memory = Vec::new()`).
/// Returns `(Snapshot, ram_offset, memory_bytes)` so callers can mmap
/// the RAM region directly for CoW restore.
pub fn load_meta(path: &str) -> Result<(Snapshot, u64, usize), FileError> {
    load_from_file_inner(path, true)
}

fn load_from_file_inner(path: &str, skip_ram: bool) -> Result<(Snapshot, u64, usize), FileError> {
    use std::io::{Seek, SeekFrom};
    let mut f = std::fs::File::open(path)?;
    let mut hdr = [0u8; 72];
    f.read_exact(&mut hdr).map_err(|_| FileError::Truncated)?;
    if hdr[0..8] != SNAPSHOT_MAGIC {
        return Err(FileError::BadMagic);
    }
    let version = le_u64(&hdr[8..16])?;
    if version != SNAPSHOT_VERSION {
        return Err(FileError::BadVersion(version));
    }
    let captured_mach_time = le_u64(&hdr[16..24])?;
    let captured_cntvct = le_u64(&hdr[24..32])?;
    let memory_bytes = le_u64(&hdr[32..40])? as usize;
    let gic_blob_len = le_u64(&hdr[40..48])? as usize;
    let n_vcpus = le_u32(&hdr[48..52])? as usize;
    if n_vcpus == 0 {
        return Err(FileError::Malformed("snapshot contains no vCPU state"));
    }
    let ram_gpa = le_u64(&hdr[56..64])?;
    let ram_offset = le_u64(&hdr[64..72])?;
    let mut gic_blob = vec![0u8; gic_blob_len];
    f.read_exact(&mut gic_blob)
        .map_err(|_| FileError::Truncated)?;
    let mut per_vcpu = Vec::with_capacity(n_vcpus);
    for _ in 0..n_vcpus {
        let mut hh = [0u8; 28];
        f.read_exact(&mut hh).map_err(|_| FileError::Truncated)?;
        let vtimer_offset = le_u64(&hh[0..8])?;
        let gp_n = le_u32(&hh[8..12])? as usize;
        let simd_n = le_u32(&hh[12..16])? as usize;
        let sys_n = le_u32(&hh[16..20])? as usize;
        let icc_n = le_u32(&hh[20..24])? as usize;
        let redist_n = le_u32(&hh[24..28])? as usize;
        let mut gp_regs = Vec::with_capacity(gp_n);
        for _ in 0..gp_n {
            let mut e = [0u8; 12];
            f.read_exact(&mut e).map_err(|_| FileError::Truncated)?;
            gp_regs.push((le_u32(&e[0..4])?, le_u64(&e[4..12])?));
        }
        let mut simd_regs = Vec::with_capacity(simd_n);
        for _ in 0..simd_n {
            let mut e = [0u8; 20];
            f.read_exact(&mut e).map_err(|_| FileError::Truncated)?;
            simd_regs.push((le_u32(&e[0..4])?, le_u128(&e[4..20])?));
        }
        let mut sys_regs = Vec::with_capacity(sys_n);
        for _ in 0..sys_n {
            let mut e = [0u8; 12];
            f.read_exact(&mut e).map_err(|_| FileError::Truncated)?;
            sys_regs.push((le_u32(&e[0..4])?, le_u64(&e[4..12])?));
        }
        let mut icc_regs = Vec::with_capacity(icc_n);
        for _ in 0..icc_n {
            let mut e = [0u8; 12];
            f.read_exact(&mut e).map_err(|_| FileError::Truncated)?;
            icc_regs.push((le_u32(&e[0..4])?, le_u64(&e[4..12])?));
        }
        let mut redist_regs = Vec::with_capacity(redist_n);
        for _ in 0..redist_n {
            let mut e = [0u8; 12];
            f.read_exact(&mut e).map_err(|_| FileError::Truncated)?;
            redist_regs.push((le_u32(&e[0..4])?, le_u64(&e[4..12])?));
        }
        per_vcpu.push(PerVcpuState {
            gp_regs,
            simd_regs,
            sys_regs,
            icc_regs,
            redist_regs,
            vtimer_offset,
        });
    }
    // Virtio section.
    let mut buf4 = [0u8; 4];
    f.read_exact(&mut buf4).map_err(|_| FileError::Truncated)?;
    let n_mmio = u32::from_le_bytes(buf4) as usize;
    let mut mmio = Vec::with_capacity(n_mmio);
    for _ in 0..n_mmio {
        let mut h = [0u8; 20];
        f.read_exact(&mut h).map_err(|_| FileError::Truncated)?;
        let driver_features = [le_u32(&h[0..4])?, le_u32(&h[4..8])?];
        let status = le_u32(&h[8..12])?;
        let interrupt_status = le_u32(&h[12..16])?;
        let n_q = le_u32(&h[16..20])? as usize;
        let mut queues = Vec::with_capacity(n_q);
        for _ in 0..n_q {
            let mut q = [0u8; 32];
            f.read_exact(&mut q).map_err(|_| FileError::Truncated)?;
            queues.push(QueueSnapshot {
                size: le_u16(&q[0..2])?,
                ready: q[2] != 0,
                desc_table: le_u64(&q[4..12])?,
                avail_ring: le_u64(&q[12..20])?,
                used_ring: le_u64(&q[20..28])?,
                last_avail_idx: le_u16(&q[28..30])?,
                next_used_idx: le_u16(&q[30..32])?,
            });
        }
        mmio.push(MmioSnapshot {
            driver_features,
            status,
            interrupt_status,
            queues,
        });
    }
    f.read_exact(&mut buf4).map_err(|_| FileError::Truncated)?;
    let n_lis = u32::from_le_bytes(buf4) as usize;
    let mut vsock_listeners = Vec::with_capacity(n_lis);
    for _ in 0..n_lis {
        let mut e = [0u8; 20];
        f.read_exact(&mut e).map_err(|_| FileError::Truncated)?;
        vsock_listeners.push(TsiListenerSnapshot {
            cid: le_u64(&e[0..8])?,
            peer_port: le_u32(&e[8..12])?,
            vm_port: le_u32(&e[12..16])?,
            family: le_u16(&e[16..18])?,
            socktype: le_u16(&e[18..20])?,
        });
    }

    let memory = if skip_ram {
        Vec::new()
    } else {
        f.seek(SeekFrom::Start(ram_offset))
            .map_err(|_| FileError::Truncated)?;
        let mut m = vec![0u8; memory_bytes];
        f.read_exact(&mut m).map_err(|_| FileError::Truncated)?;
        m
    };

    Ok((
        Snapshot {
            captured_mach_time,
            captured_cntvct,
            ram_gpa,
            memory,
            gic_blob,
            per_vcpu,
            virtio: VirtioSnapshot {
                mmio,
                vsock_listeners,
            },
        },
        ram_offset,
        memory_bytes,
    ))
}

/// mmap a snapshot file's RAM region into a new private mapping.
/// Returns the host-side pointer + length suitable for handing to
/// `MicroVm::new_with_ram` and `Vm::map`. The mapping is `MAP_PRIVATE`,
/// so guest writes go to anon pages backing the COW'd portion only —
/// the snapshot file on disk is never written.
pub fn mmap_ram_cow(path: &str) -> std::io::Result<(*mut u8, usize)> {
    let (_snap, ram_offset, memory_bytes) =
        load_meta(path).map_err(|e| std::io::Error::other(format!("load_meta: {e:?}")))?;
    mmap_ram_cow_at(path, ram_offset, memory_bytes)
}

pub fn mmap_ram_cow_at(
    path: &str,
    ram_offset: u64,
    memory_bytes: usize,
) -> std::io::Result<(*mut u8, usize)> {
    let f = std::fs::File::open(path)?;

    // SAFETY: kernel allocates pages.
    let fd = std::os::fd::AsRawFd::as_raw_fd(&f);
    let ptr = unsafe {
        libc::mmap(
            std::ptr::null_mut(),
            memory_bytes,
            libc::PROT_READ | libc::PROT_WRITE,
            libc::MAP_PRIVATE,
            fd,
            ram_offset as libc::off_t,
        )
    };
    if ptr == libc::MAP_FAILED {
        return Err(std::io::Error::last_os_error());
    }
    Ok((ptr as *mut u8, memory_bytes))
}

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

    fn temp_snapshot_path(name: &str) -> PathBuf {
        let mut path = std::env::temp_dir();
        path.push(format!(
            "snapshot-{name}-{}-{:?}.snap",
            std::process::id(),
            std::thread::current().id()
        ));
        path
    }

    fn write_bytes(path: &PathBuf, bytes: &[u8]) -> std::io::Result<()> {
        let mut file = std::fs::File::create(path)?;
        file.write_all(bytes)
    }

    fn minimal_header(n_vcpus: u32) -> [u8; 72] {
        let mut hdr = [0u8; 72];
        hdr[0..8].copy_from_slice(&SNAPSHOT_MAGIC);
        hdr[8..16].copy_from_slice(&SNAPSHOT_VERSION.to_le_bytes());
        hdr[32..40].copy_from_slice(&0u64.to_le_bytes());
        hdr[40..48].copy_from_slice(&0u64.to_le_bytes());
        hdr[48..52].copy_from_slice(&n_vcpus.to_le_bytes());
        hdr[64..72].copy_from_slice(&72u64.to_le_bytes());
        hdr
    }

    #[test]
    fn load_rejects_truncated_snapshot() -> std::io::Result<()> {
        let path = temp_snapshot_path("truncated");
        write_bytes(&path, b"SMS")?;

        let result = load_from_file(path.to_str().unwrap_or_default());
        let _ = std::fs::remove_file(path);

        assert!(matches!(result, Err(FileError::Truncated)));
        Ok(())
    }

    #[test]
    fn load_rejects_bad_magic() -> std::io::Result<()> {
        let path = temp_snapshot_path("bad-magic");
        write_bytes(&path, &[0u8; 72])?;

        let result = load_from_file(path.to_str().unwrap_or_default());
        let _ = std::fs::remove_file(path);

        assert!(matches!(result, Err(FileError::BadMagic)));
        Ok(())
    }

    #[test]
    fn load_rejects_missing_vcpu_state() -> std::io::Result<()> {
        let path = temp_snapshot_path("no-vcpu");
        write_bytes(&path, &minimal_header(0))?;

        let result = load_from_file(path.to_str().unwrap_or_default());
        let _ = std::fs::remove_file(path);

        assert!(matches!(
            result,
            Err(FileError::Malformed("snapshot contains no vCPU state"))
        ));
        Ok(())
    }
}