// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

//! Virtual Machine Monitor that leverages the Linux Kernel-based Virtual Machine (KVM),
//! and other virtualization features to run a single lightweight micro-virtual
//! machine (microVM).
//#![deny(missing_docs)]

#[macro_use]
extern crate log;

/// Handles setup and initialization a `Vmm` object.
pub mod builder;
pub(crate) mod device_manager;
/// Resource store for configured microVM resources.
pub mod resources;
/// Signal handling utilities.
#[cfg(target_os = "linux")]
pub mod signal_handler;
/// Guest-memory checkpoint/restore (the memory portion of a VM snapshot).
pub mod snapshot;
/// Wrappers over structures used to configure the VMM.
pub mod vmm_config;

#[cfg(target_os = "linux")]
mod linux;
#[cfg(target_os = "linux")]
use crate::linux::vstate;
#[cfg(target_os = "macos")]
mod macos;
mod terminal;
#[cfg(target_os = "windows")]
mod windows;
pub mod worker;

#[cfg(target_os = "macos")]
use macos::vstate;
#[cfg(target_os = "windows")]
use windows::vstate;

use std::fmt::{Display, Formatter};
use std::io;
#[cfg(unix)]
use std::os::unix::io::AsRawFd;
use std::sync::atomic::{AtomicI32, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[cfg(target_os = "windows")]
use utils::windows::AsRawFd;

#[cfg(target_arch = "x86_64")]
use crate::device_manager::legacy::PortIODeviceManager;
use crate::device_manager::mmio::MMIODeviceManager;
use crate::vstate::VcpuEvent;
use crate::vstate::{Vcpu, VcpuHandle, VcpuResponse, Vm};

use arch::{ArchMemoryInfo, InitrdConfig};
#[cfg(target_os = "macos")]
use crossbeam_channel::Sender;
#[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
use devices::fdt;
use devices::legacy::IrqChip;
use devices::virtio::VmmExitObserver;
use devices::{BusDevice, DeviceType};
use kernel::cmdline::Cmdline as KernelCmdline;
use polly::event_manager::{self, EventManager, Subscriber};
use utils::epoll::{EpollEvent, EventSet};
use utils::eventfd::EventFd;
pub use vm_memory::GuestMemoryMmap;

/// Success exit code.
pub const FC_EXIT_CODE_OK: u8 = 0;
/// Generic error exit code.
pub const FC_EXIT_CODE_GENERIC_ERROR: u8 = 1;
/// Generic exit code for an error considered not possible to occur if the program logic is sound.
pub const FC_EXIT_CODE_UNEXPECTED_ERROR: u8 = 2;
/// Firecracker was shut down after intercepting a restricted system call.
pub const FC_EXIT_CODE_BAD_SYSCALL: u8 = 148;
/// Firecracker was shut down after intercepting `SIGBUS`.
pub const FC_EXIT_CODE_SIGBUS: u8 = 149;
/// Firecracker was shut down after intercepting `SIGSEGV`.
pub const FC_EXIT_CODE_SIGSEGV: u8 = 150;
/// Bad configuration for microvm's resources, when using a single json.
pub const FC_EXIT_CODE_BAD_CONFIGURATION: u8 = 152;
/// Command line arguments parsing error.
pub const FC_EXIT_CODE_ARG_PARSING: u8 = 153;

/// Errors associated with the VMM internal logic. These errors cannot be generated by direct user
/// input, but can result from bad configuration of the host (for example if Firecracker doesn't
/// have permissions to open the KVM fd).
#[derive(Debug)]
pub enum Error {
    /// This error is thrown by the minimal boot loader implementation.
    ConfigureSystem(arch::Error),
    /// Legacy devices work with Event file descriptors and the creation can fail because
    /// of resource exhaustion.
    #[cfg(target_arch = "x86_64")]
    CreateLegacyDevice(device_manager::legacy::Error),
    /// Cannot read from an Event file descriptor.
    EventFd(io::Error),
    /// Polly error wrapper.
    EventManager(event_manager::Error),
    /// I8042 Error.
    #[cfg(target_arch = "x86_64")]
    I8042Error(devices::legacy::I8042DeviceError),
    /// Cannot access kernel file.
    KernelFile(io::Error),
    /// Cannot open /dev/kvm. Either the host does not have KVM or Firecracker does not have
    /// permission to open the file descriptor.
    KvmContext(vstate::Error),
    #[cfg(target_arch = "x86_64")]
    /// Cannot add devices to the Legacy I/O Bus.
    LegacyIOBus(device_manager::legacy::Error),
    /// Cannot load command line.
    LoadCommandline(kernel::cmdline::Error),
    /// Cannot add a device to the MMIO Bus.
    RegisterMMIODevice(device_manager::mmio::Error),
    /// Write to the serial console failed.
    Serial(io::Error),
    #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
    /// Cannot generate or write FDT
    SetupFDT(devices::fdt::Error),
    /// Cannot create Timer file descriptor.
    TimerFd(io::Error),
    /// Vcpu error.
    Vcpu(vstate::Error),
    /// Cannot send event to vCPU.
    VcpuEvent(vstate::Error),
    /// Cannot create a vCPU handle.
    VcpuHandle(vstate::Error),
    /// vCPU pause failed.
    VcpuPause,
    /// vCPU resume failed.
    VcpuResume,
    /// vCPU state checkpoint/restore failed.
    VcpuSnapshot(String),
    /// VM checkpoint/restore (memory, device, or VM state) failed.
    Snapshot(String),
    /// Fork snapshots require guest RAM backed by shareable file descriptors.
    #[cfg(fork_supported)]
    ForkRequiresMemfd,
    /// Cannot spawn a new Vcpu thread.
    VcpuSpawn(std::io::Error),
    /// Vm error.
    Vm(vstate::Error),
    /// Error thrown by observer object on Vmm initialization.
    VmmObserverInit(utils::errno::Error),
    /// Error thrown by observer object on Vmm teardown.
    VmmObserverTeardown(utils::errno::Error),
}

impl Display for Error {
    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
        use self::Error::*;

        match self {
            ConfigureSystem(e) => write!(f, "System configuration error: {e:?}"),
            #[cfg(target_arch = "x86_64")]
            CreateLegacyDevice(e) => write!(f, "Error creating legacy device: {e:?}"),
            EventFd(e) => write!(f, "Event fd error: {e}"),
            EventManager(e) => write!(f, "Event manager error: {e:?}"),
            #[cfg(target_arch = "x86_64")]
            I8042Error(e) => write!(f, "I8042 error: {e}"),
            KernelFile(e) => write!(f, "Cannot access kernel file: {e}"),
            KvmContext(e) => write!(f, "Failed to validate KVM support: {e:?}"),
            #[cfg(target_arch = "x86_64")]
            LegacyIOBus(e) => write!(f, "Cannot add devices to the legacy I/O Bus. {e}"),
            LoadCommandline(e) => write!(f, "Cannot load command line: {e}"),
            RegisterMMIODevice(e) => write!(f, "Cannot add a device to the MMIO Bus. {e}"),
            Serial(e) => write!(f, "Error writing to the serial console: {e:?}"),
            #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
            SetupFDT(e) => write!(f, "Error generating or writing FDT: {e:?}"),
            TimerFd(e) => write!(f, "Error creating timer fd: {e}"),
            Vcpu(e) => write!(f, "Vcpu error: {e}"),
            VcpuEvent(e) => write!(f, "Cannot send event to vCPU. {e:?}"),
            VcpuHandle(e) => write!(f, "Cannot create a vCPU handle. {e}"),
            VcpuPause => write!(f, "vCPUs pause failed."),
            VcpuResume => write!(f, "vCPUs resume failed."),
            VcpuSnapshot(e) => write!(f, "vCPU snapshot/restore failed: {e}"),
            Snapshot(e) => write!(f, "VM snapshot/restore failed: {e}"),
            #[cfg(fork_supported)]
            ForkRequiresMemfd => write!(f, "fork requires memfd-backed guest RAM"),
            VcpuSpawn(e) => write!(f, "Cannot spawn Vcpu thread: {e}"),
            Vm(e) => write!(f, "Vm error: {e}"),
            VmmObserverInit(e) => write!(
                f,
                "Error thrown by observer object on Vmm initialization: {e}"
            ),
            VmmObserverTeardown(e) => {
                write!(f, "Error thrown by observer object on Vmm teardown: {e}")
            }
        }
    }
}

/// Trait for objects that need custom initialization and teardown during the Vmm lifetime.
pub trait VmmEventsObserver {
    /// This function will be called during microVm boot.
    fn on_vmm_boot(&mut self) -> std::result::Result<(), utils::errno::Error> {
        Ok(())
    }
    /// This function will be called on microVm teardown.
    fn on_vmm_stop(&mut self) -> std::result::Result<(), utils::errno::Error> {
        Ok(())
    }
}

/// Shorthand result type for internal VMM commands.
pub type Result<T> = std::result::Result<T, Error>;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum VmmRunState {
    Paused,
    Running,
    Pausing,
    Resuming,
}

/// A full VM checkpoint, minus the guest-memory image (captured/streamed
/// separately because it is large). Combines the VM-level KVM state, per-vCPU
/// register state, and virtio device state. Produced by [`Vmm::checkpoint`] and
/// consumed by [`Vmm::restore`].
#[cfg(snapshot_supported)]
pub struct VmCheckpoint {
    /// VM-level KVM state (PIT, clock, PIC/IOAPIC).
    pub vm_state: vstate::VmState,
    /// Per-vCPU register state, in vCPU-index order.
    pub vcpu_states: Vec<vstate::VcpuState>,
    /// Virtio device runtime state (queues, negotiated features, etc.).
    pub devices: devices::virtio::persist::VmDevicesState,
    /// Userspace interrupt-controller registers. `None` on hypervisors whose
    /// IRQ chip is in-kernel and already inside `vm_state` (KVM) or captured
    /// separately (HVF GIC distributor). On WHP the IOAPIC is a software device
    /// in this process; without this section a restored clone boots with an
    /// all-masked redirection table and device IRQs (e.g. vsock) never arrive.
    pub ioapic: Option<Vec<u8>>,
}

#[cfg(snapshot_supported)]
impl VmCheckpoint {
    /// Serialize the checkpoint (VM + vCPU + device state, *not* guest RAM —
    /// that is shared/streamed separately) into a single blob for cross-process
    /// fork or on-disk hibernate. Length-prefixed sections; no cross-version
    /// compatibility promised.
    pub fn serialize(&self) -> Vec<u8> {
        fn put(out: &mut Vec<u8>, b: &[u8]) {
            out.extend_from_slice(&(b.len() as u32).to_le_bytes());
            out.extend_from_slice(b);
        }
        let mut out = Vec::new();
        put(&mut out, &self.vm_state.serialize());
        out.extend_from_slice(&(self.vcpu_states.len() as u32).to_le_bytes());
        for v in &self.vcpu_states {
            put(&mut out, &v.serialize());
        }
        put(&mut out, &self.devices.to_bytes().unwrap_or_default());
        // Userspace IRQ-chip section; empty means "none" (in-kernel chip).
        put(&mut out, self.ioapic.as_deref().unwrap_or(&[]));
        out
    }

    /// Reconstruct from a blob produced by [`Self::serialize`].
    pub fn deserialize(bytes: &[u8]) -> std::result::Result<VmCheckpoint, String> {
        fn take<'a>(b: &'a [u8], pos: &mut usize) -> std::result::Result<&'a [u8], String> {
            if *pos + 4 > b.len() {
                return Err("checkpoint truncated (length header)".to_string());
            }
            let len = u32::from_le_bytes(b[*pos..*pos + 4].try_into().unwrap()) as usize;
            *pos += 4;
            if *pos + len > b.len() {
                return Err("checkpoint truncated (section)".to_string());
            }
            let s = &b[*pos..*pos + len];
            *pos += len;
            Ok(s)
        }
        let mut pos = 0usize;
        let vm_state = vstate::VmState::deserialize(take(bytes, &mut pos)?)?;
        if pos + 4 > bytes.len() {
            return Err("checkpoint truncated (vcpu count)".to_string());
        }
        let n = u32::from_le_bytes(bytes[pos..pos + 4].try_into().unwrap()) as usize;
        pos += 4;
        let mut vcpu_states = Vec::with_capacity(n);
        for _ in 0..n {
            vcpu_states.push(vstate::VcpuState::deserialize(take(bytes, &mut pos)?)?);
        }
        let devices = devices::virtio::persist::VmDevicesState::from_bytes(take(bytes, &mut pos)?)?;
        // Userspace IRQ-chip section: empty (or absent, for blobs written before
        // the section existed) means the chip is in-kernel and rides vm_state.
        let ioapic = if pos < bytes.len() {
            let b = take(bytes, &mut pos)?;
            (!b.is_empty()).then(|| b.to_vec())
        } else {
            None
        };
        Ok(VmCheckpoint {
            vm_state,
            vcpu_states,
            devices,
            ioapic,
        })
    }
}

/// Contains the state and associated methods required for the Firecracker VMM.
pub struct Vmm {
    // Guest VM core resources.
    guest_memory: GuestMemoryMmap,
    arch_memory_info: ArchMemoryInfo,

    kernel_cmdline: KernelCmdline,

    vcpus_handles: Vec<VcpuHandle>,
    run_state: VmmRunState,
    paused_at: Option<Instant>,
    devices_quiesced: bool,
    exit_evt: EventFd,
    vm: Vm,
    exit_observers: Vec<Arc<Mutex<dyn VmmExitObserver>>>,
    exit_code: Arc<AtomicI32>,

    // Guest VM devices.
    mmio_device_manager: MMIODeviceManager,
    /// Balloon device handle for runtime target control (idle-squeeze).
    #[cfg(not(feature = "tee"))]
    balloon: Option<Arc<Mutex<devices::virtio::Balloon>>>,
    #[cfg(target_arch = "x86_64")]
    pio_device_manager: PortIODeviceManager,
    /// The interrupt controller, held for checkpoint/restore of userspace IRQ
    /// chips (the WHP software IOAPIC — see [`VmCheckpoint::ioapic`]).
    #[cfg(all(target_arch = "x86_64", target_os = "windows"))]
    intc: devices::legacy::IrqChip,
}

#[cfg(fork_supported)]
fn has_memfd_backed_memory(descs: &[snapshot::MemfdRegionDesc]) -> bool {
    descs.iter().any(|desc| desc.fd >= 0)
}

impl Vmm {
    #[cfg(not(feature = "tee"))]
    pub(crate) fn set_balloon(&mut self, balloon: Arc<Mutex<devices::virtio::Balloon>>) {
        self.balloon = Some(balloon);
    }

    /// Ask the guest to inflate/deflate its balloon to `mib` MiB (0 fully
    /// deflates). Returns `(target_pages, actual_pages)` after applying.
    #[cfg(not(feature = "tee"))]
    pub fn balloon_set_target_mib(
        &mut self,
        mib: u32,
    ) -> std::result::Result<(u32, u32), &'static str> {
        let balloon = self.balloon.as_ref().ok_or("no balloon device")?;
        let mut b = balloon.lock().unwrap();
        b.set_target_pages(mib.saturating_mul(256));
        Ok(b.pages())
    }

    /// Current balloon `(target_pages, actual_pages)`.
    #[cfg(not(feature = "tee"))]
    pub fn balloon_pages(&self) -> std::result::Result<(u32, u32), &'static str> {
        let balloon = self.balloon.as_ref().ok_or("no balloon device")?;
        let b = balloon.lock().unwrap();
        Ok(b.pages())
    }

    /// The guest RAM regions as `(guest_physical_start, host_virtual_addr, len)`
    /// triples.
    ///
    /// Because an embedder-hosted server (e.g. a CUDA forwarding server sharing
    /// this process) can read guest memory directly at `hva + (gpa - gpa_start)`,
    /// this is the primitive that enables zero-copy guest↔host transfers: the
    /// guest ships a guest-physical descriptor instead of the bytes. Guest RAM is
    /// typically split into a low and a high region around the 4 GiB PCI hole, so
    /// all regions are returned.
    pub fn guest_ram_regions(&self) -> Vec<(u64, u64, u64)> {
        use vm_memory::{Address, GuestMemory, GuestMemoryRegion};
        self.guest_memory
            .iter()
            .filter_map(|r| {
                let gpa = r.start_addr().raw_value();
                let hva = self.guest_memory.get_host_address(r.start_addr()).ok()? as u64;
                Some((gpa, hva, r.len()))
            })
            .collect()
    }

    fn wait_for_vcpu_response(
        handle: &VcpuHandle,
        expected: VcpuResponse,
        deadline: Instant,
    ) -> Result<()> {
        loop {
            let remaining = deadline
                .checked_duration_since(Instant::now())
                .ok_or(Error::VcpuPause)?;

            match handle.response_receiver().recv_timeout(remaining) {
                // Compare by variant, not value: some responses carry a payload
                // (e.g. captured vCPU state) that is not comparable.
                Ok(response)
                    if std::mem::discriminant(&response) == std::mem::discriminant(&expected) =>
                {
                    return Ok(());
                }
                Ok(VcpuResponse::Exited(_)) => return Err(Error::VcpuPause),
                Ok(_) => {}
                Err(_) => return Err(Error::VcpuPause),
            }
        }
    }

    fn paused_duration_ns(paused_duration: Duration) -> u64 {
        u64::try_from(paused_duration.as_nanos()).unwrap_or(u64::MAX)
    }

    fn adjust_guest_clock_after_pause(&self, paused_duration: Duration) -> Result<()> {
        self.vm
            .adjust_clock_after_pause(Self::paused_duration_ns(paused_duration))
            .map_err(Error::Vm)
    }

    /// Gets the the specified bus device.
    pub fn get_bus_device(
        &self,
        device_type: DeviceType,
        device_id: &str,
    ) -> Option<&Mutex<dyn BusDevice>> {
        self.mmio_device_manager.get_device(device_type, device_id)
    }

    /// Capture the runtime state of all snapshot-supporting virtio devices.
    ///
    /// This is the device portion of a VM checkpoint/fork. A full checkpoint
    /// composes this with (a) paused-vCPU state via `vstate` save_state and
    /// (b) the guest-memory image; restore reverses it after the VM is rebuilt.
    /// Callers must pause the vCPUs and drain device workers first so the
    /// captured queue indices are at a clean boundary.
    pub fn snapshot_devices(&self) -> devices::virtio::persist::VmDevicesState {
        self.mmio_device_manager.snapshot_devices()
    }

    /// Restore device state (captured by [`Self::snapshot_devices`]) onto the
    /// freshly-rebuilt, not-yet-activated devices during VM restore.
    pub fn restore_devices(
        &self,
        state: &devices::virtio::persist::VmDevicesState,
    ) -> std::result::Result<(), String> {
        self.mmio_device_manager.restore_devices(state)
    }

    /// Re-activate virtio devices on a freshly-built clone from a checkpoint —
    /// reconstruct each device's queues from saved state and activate, bypassing
    /// the guest handshake (a restored guest resumes past boot). Used by
    /// restore-into-a-fresh-VM (cross-process fork); not used by in-place rewind
    /// (there the devices are already activated).
    #[cfg(snapshot_supported)]
    pub fn restore_activate_devices(
        &self,
        state: &devices::virtio::persist::VmDevicesState,
    ) -> std::result::Result<(), String> {
        self.mmio_device_manager.restore_activate_devices(state)
    }

    /// Apply a checkpoint onto a freshly-built clone whose vCPUs were started
    /// **paused** ([`Self::start_vcpus_paused`]) and whose guest RAM is the
    /// restored image: restore VM-level state, re-activate the virtio devices
    /// from saved queue state, and load the vCPU registers. The caller resumes
    /// (e.g. [`Self::resume`]) to start the clone running.
    #[cfg(snapshot_supported)]
    pub fn apply_restore(&mut self, checkpoint: VmCheckpoint) -> Result<()> {
        self.vm
            .restore_state(&checkpoint.vm_state)
            .map_err(Error::Vm)?;
        // Restore the userspace IRQ chip (WHP software IOAPIC) before devices
        // re-activate: a fresh chip has every redirection entry masked, so an
        // IRQ raised earlier than this would be dropped and the clone's vsock
        // never delivers.
        #[cfg(all(target_arch = "x86_64", target_os = "windows"))]
        if let Some(bytes) = &checkpoint.ioapic {
            self.intc.lock().unwrap().restore_state(bytes);
        }
        self.restore_activate_devices(&checkpoint.devices)
            .map_err(Error::Snapshot)?;
        self.restore_vcpu_states(checkpoint.vcpu_states)?;
        Ok(())
    }

    /// Drain + reclaim worker-owned virtqueues (block/net) so their indices can
    /// be captured/restored at a clean boundary. vCPUs must be paused first.
    pub fn quiesce_devices(&mut self) {
        if !self.devices_quiesced {
            self.mmio_device_manager.quiesce_devices();
            self.devices_quiesced = true;
        }
    }

    /// Re-arm device workers quiesced by [`Self::quiesce_devices`].
    pub fn rearm_devices(&mut self) {
        if self.devices_quiesced {
            self.mmio_device_manager.rearm_devices();
            self.devices_quiesced = false;
        }
    }

    /// Starts the microVM vcpus (running).
    pub fn start_vcpus(&mut self, vcpus: Vec<Vcpu>) -> Result<()> {
        self.start_vcpus_inner(vcpus, true)
    }

    /// Spawn the vCPU threads but leave them **paused** (Linux: they sit in the
    /// paused state machine). Used by restore-into-a-fresh-VM so VM/device/vCPU
    /// state can be restored before the guest runs; resume via [`Self::resume`].
    #[cfg(snapshot_supported)]
    pub fn start_vcpus_paused(&mut self, vcpus: Vec<Vcpu>) -> Result<()> {
        self.start_vcpus_inner(vcpus, false)
    }

    fn start_vcpus_inner(&mut self, mut vcpus: Vec<Vcpu>, resume_after: bool) -> Result<()> {
        let t_sv = std::time::Instant::now();
        let sv_timing_on = std::env::var("LIBKRUN_TIMING").is_ok();
        macro_rules! sv_timing {
            ($label:expr) => {
                if sv_timing_on {
                    eprintln!("[vcpu] {:28} {}ms", $label, t_sv.elapsed().as_millis());
                }
            };
        }

        let vcpu_count = vcpus.len();

        Vcpu::register_kick_signal_handler();

        self.vcpus_handles.reserve(vcpu_count);

        for mut vcpu in vcpus.drain(..) {
            vcpu.set_mmio_bus(self.mmio_device_manager.bus.clone());

            // HVF (macOS) and WHP (Windows): vCPU threads run guest code
            // immediately after spawn, so a restore-into-a-clone (paused) start
            // must mark the vCPU to hold in its paused event loop until the
            // orchestrator restores registers and resumes. On KVM the paused
            // state machine handles this already.
            #[cfg(any(
                all(target_os = "macos", target_arch = "aarch64"),
                all(target_os = "windows", target_arch = "x86_64")
            ))]
            vcpu.set_start_paused(!resume_after);

            self.vcpus_handles
                .push(vcpu.start_threaded().map_err(Error::VcpuHandle)?);
        }
        sv_timing!("threads spawned (hv_vcpu_create)");

        self.run_state = if resume_after {
            VmmRunState::Running
        } else {
            VmmRunState::Paused
        };

        // HVF (macOS): vCPU threads are not in a Paused state after spawn —
        // handle_pending_event() uses a non-blocking try_recv on the first
        // iteration, finds nothing, and enters hv_vcpu_run() immediately.
        // Sending VcpuEvent::Resume would queue a VcpuResponse::Resumed that
        // sits in the channel and is silently discarded by the next
        // wait_for_vcpu_response call, so skip it.
        //
        // KVM (Linux): vCPU threads start in the Paused state machine
        // (vstate.rs: StateMachine::run(self, Self::paused)). They block
        // waiting for a Resume event and will never execute guest code until
        // one is sent.
        #[cfg(target_os = "linux")]
        if resume_after {
            self.resume_vcpus(Duration::ZERO)?;
        }

        sv_timing!("vcpus started");

        Ok(())
    }

    /// Sends a pause command to the vcpus.
    pub fn pause_vcpus(&mut self) -> Result<()> {
        for handle in self.vcpus_handles.iter() {
            handle
                .send_event(VcpuEvent::Pause)
                .map_err(Error::VcpuEvent)?;
        }

        let deadline = Instant::now() + Duration::from_millis(1000);
        for handle in self.vcpus_handles.iter() {
            Self::wait_for_vcpu_response(handle, VcpuResponse::Paused, deadline)?;
        }
        Ok(())
    }

    /// Sends a resume command to the vcpus.
    pub fn resume_vcpus(&mut self, paused_duration: Duration) -> Result<()> {
        let paused_ns = Self::paused_duration_ns(paused_duration);
        for handle in self.vcpus_handles.iter() {
            handle
                .send_event(VcpuEvent::Resume { paused_ns })
                .map_err(Error::VcpuEvent)?;
        }

        let deadline = Instant::now() + Duration::from_millis(1000);
        for handle in self.vcpus_handles.iter() {
            Self::wait_for_vcpu_response(handle, VcpuResponse::Resumed, deadline)
                .map_err(|_| Error::VcpuResume)?;
        }
        Ok(())
    }

    /// Capture the full register state of every vCPU for a checkpoint.
    ///
    /// The vCPUs must already be paused (call [`Self::pause`] first) — capture
    /// runs in each vCPU thread while it sits in the paused state loop, which is
    /// the only safe point to issue the KVM GET ioctls (`KVM_GET_VCPU_EVENTS` is
    /// unsafe while other vCPUs run). Returns one [`vstate::VcpuState`] per vCPU,
    /// in vCPU-index order.
    #[cfg(snapshot_supported)]
    pub fn save_vcpu_states(&mut self) -> Result<Vec<vstate::VcpuState>> {
        if self.run_state != VmmRunState::Paused {
            return Err(Error::VcpuSnapshot(
                "vCPUs must be paused before capturing state".to_string(),
            ));
        }
        for handle in self.vcpus_handles.iter() {
            handle
                .send_event(VcpuEvent::SaveState)
                .map_err(Error::VcpuEvent)?;
        }
        let deadline = Instant::now() + Duration::from_secs(5);
        let mut states = Vec::with_capacity(self.vcpus_handles.len());
        for handle in self.vcpus_handles.iter() {
            states.push(Self::wait_for_saved_state(handle, deadline)?);
        }
        Ok(states)
    }

    #[cfg(snapshot_supported)]
    fn wait_for_saved_state(handle: &VcpuHandle, deadline: Instant) -> Result<vstate::VcpuState> {
        loop {
            let remaining = deadline
                .checked_duration_since(Instant::now())
                .ok_or_else(|| Error::VcpuSnapshot("timed out capturing vCPU state".to_string()))?;
            match handle.response_receiver().recv_timeout(remaining) {
                Ok(VcpuResponse::SavedState(state)) => return Ok(*state),
                Ok(VcpuResponse::Exited(_)) => {
                    return Err(Error::VcpuSnapshot(
                        "vCPU exited during capture".to_string(),
                    ));
                }
                Ok(_) => {}
                Err(_) => {
                    return Err(Error::VcpuSnapshot(
                        "channel closed while capturing vCPU state".to_string(),
                    ));
                }
            }
        }
    }

    /// Restore previously-captured vCPU register state onto the paused vCPUs
    /// during VM restore. `states` must hold exactly one entry per vCPU, in
    /// vCPU-index order (as produced by [`Self::save_vcpu_states`]). The vCPUs
    /// must be paused; the caller resumes them afterwards.
    #[cfg(snapshot_supported)]
    pub fn restore_vcpu_states(&mut self, states: Vec<vstate::VcpuState>) -> Result<()> {
        if states.len() != self.vcpus_handles.len() {
            return Err(Error::VcpuSnapshot(format!(
                "vCPU state count {} does not match vCPU count {}",
                states.len(),
                self.vcpus_handles.len()
            )));
        }
        for (handle, state) in self.vcpus_handles.iter().zip(states) {
            handle
                .send_event(VcpuEvent::RestoreState(Box::new(state)))
                .map_err(Error::VcpuEvent)?;
        }
        let deadline = Instant::now() + Duration::from_secs(5);
        for handle in self.vcpus_handles.iter() {
            Self::wait_for_vcpu_response(handle, VcpuResponse::RestoredState, deadline)
                .map_err(|_| Error::VcpuSnapshot("vCPU restore not acknowledged".to_string()))?;
        }
        Ok(())
    }

    /// Capture a full VM checkpoint: pause at a clean boundary, then snapshot
    /// the VM-level KVM state, every vCPU's register state, the virtio device
    /// state, and the guest-memory image (streamed to `mem_out`).
    ///
    /// On success the VM is left **paused** — the caller decides whether to
    /// resume the original (e.g. after a save-to-file), retire it (fork:
    /// the clones continue from this image), or tear it down. The returned
    /// [`VmCheckpoint`] plus the memory image and its [`MemoryRegionDesc`] list
    /// is everything needed to reconstruct the VM (see [`Self::restore`]).
    ///
    /// For the in-process fork fast path, capture once and clone the guest
    /// memory (CoW) per child instead of streaming it; the CPU/device state in
    /// [`VmCheckpoint`] is shared by all children (each then rejuvenated).
    #[cfg(snapshot_supported)]
    pub fn checkpoint<W: std::io::Write>(
        &mut self,
        mem_out: &mut W,
    ) -> Result<(VmCheckpoint, Vec<snapshot::MemoryRegionDesc>)> {
        // Quiesce all vCPUs at an instruction boundary first, then drain the
        // device workers so block/net virtqueue indices are captured at a clean
        // boundary (no in-flight I/O).
        self.pause()?;
        self.quiesce_devices();
        let vcpu_states = self.save_vcpu_states()?;
        let vm_state = self.vm.save_state().map_err(Error::Vm)?;
        let devices = self.snapshot_devices();
        let mem_descs = snapshot::write_guest_memory(&self.guest_memory, mem_out)
            .map_err(|e| Error::Snapshot(format!("guest-memory dump: {e}")))?;
        // Re-arm workers so the original VM can resume from this checkpoint.
        self.rearm_devices();
        #[cfg(all(target_arch = "x86_64", target_os = "windows"))]
        let ioapic = self.intc.lock().unwrap().save_state();
        #[cfg(not(all(target_arch = "x86_64", target_os = "windows")))]
        let ioapic = None;
        Ok((
            VmCheckpoint {
                vm_state,
                vcpu_states,
                devices,
                ioapic,
            },
            mem_descs,
        ))
    }

    /// Restore a VM from a [`VmCheckpoint`] and its guest-memory image onto this
    /// (freshly-built, vCPU-paused) `Vmm`: load guest RAM, restore VM-level and
    /// device state, then load every vCPU's registers. The caller resumes the
    /// vCPUs afterwards (e.g. via [`Self::resume`]).
    ///
    /// `self` must have been built with the same VM config as the checkpoint
    /// (vCPU count, memory layout, device set) so the layout matches.
    #[cfg(snapshot_supported)]
    pub fn restore<R: std::io::Read>(
        &mut self,
        checkpoint: VmCheckpoint,
        mem_descs: &[snapshot::MemoryRegionDesc],
        mem_in: &mut R,
    ) -> Result<()> {
        // Devices must be in their not-yet-running, paused state.
        if self.run_state != VmmRunState::Paused {
            return Err(Error::Snapshot(
                "VM must be paused before restore".to_string(),
            ));
        }
        // Drain + reclaim worker-owned queues so restored indices can be applied
        // to them (and so they match the rewound guest rings).
        self.quiesce_devices();
        snapshot::read_guest_memory_into(&self.guest_memory, mem_descs, mem_in)
            .map_err(|e| Error::Snapshot(format!("guest-memory load: {e}")))?;
        self.vm
            .restore_state(&checkpoint.vm_state)
            .map_err(Error::Vm)?;
        // See apply_restore: the WHP software IOAPIC must be restored before
        // devices signal or their IRQs hit an all-masked redirection table.
        #[cfg(all(target_arch = "x86_64", target_os = "windows"))]
        if let Some(bytes) = &checkpoint.ioapic {
            self.intc.lock().unwrap().restore_state(bytes);
        }
        self.restore_devices(&checkpoint.devices)
            .map_err(Error::Snapshot)?;
        self.restore_vcpu_states(checkpoint.vcpu_states)?;
        // Re-arm workers from the restored queue indices.
        self.rearm_devices();
        Ok(())
    }

    /// Capture VM/vCPU/device state plus an in-memory **CoW clone** of guest RAM,
    /// for **fork** (1→N). For memfd-backed RAM the clone shares the parent's
    /// clean pages (cheap, bounded, size-independent) and copies on write.
    ///
    /// **Fork-only — the parent must be frozen/retired after this call.** This is
    /// NOT valid for an in-process *rewind* where the parent resumes: a resuming
    /// parent keeps writing the shared memfd, and a `MAP_PRIVATE` clone observes
    /// those later writes on pages it hasn't copied yet (plan §4 / FINDINGS Exp 1),
    /// so the clone would not preserve the checkpoint snapshot. Rewind/hibernate
    /// must use the eager [`Self::checkpoint`] (which copies the bytes at
    /// checkpoint time and is immune to later parent writes).
    #[cfg(all(target_os = "linux", target_arch = "x86_64"))]
    pub fn checkpoint_cow(&mut self) -> Result<(VmCheckpoint, GuestMemoryMmap)> {
        self.pause()?;
        self.quiesce_devices();
        let vcpu_states = self.save_vcpu_states()?;
        let vm_state = self.vm.save_state().map_err(Error::Vm)?;
        let devices = self.snapshot_devices();
        let mem_clone = snapshot::cow_clone_guest_memory(&self.guest_memory)
            .map_err(|e| Error::Snapshot(format!("cow-clone guest memory: {e}")))?;
        self.rearm_devices();
        Ok((
            VmCheckpoint {
                vm_state,
                vcpu_states,
                devices,
                // Linux-only path; KVM's IOAPIC is in-kernel and rides vm_state.
                ioapic: None,
            },
            mem_clone,
        ))
    }

    /// Restore this VM from a CoW snapshot captured by [`Self::checkpoint_cow`]
    /// **whose source was frozen after capture** (the clone's pages must still
    /// reflect the checkpoint — see the caveat on `checkpoint_cow`): copy the
    /// snapshot's guest RAM back into the live memory, then restore
    /// VM/device/vCPU state. The VM must be paused; the caller resumes. (The
    /// general rewind path is the eager [`Self::restore`].)
    #[cfg(all(target_os = "linux", target_arch = "x86_64"))]
    pub fn restore_cow(
        &mut self,
        checkpoint: VmCheckpoint,
        mem_clone: &GuestMemoryMmap,
    ) -> Result<()> {
        if self.run_state != VmmRunState::Paused {
            return Err(Error::Snapshot(
                "VM must be paused before restore".to_string(),
            ));
        }
        self.quiesce_devices();
        snapshot::copy_guest_memory(mem_clone, &self.guest_memory)
            .map_err(|e| Error::Snapshot(format!("guest-memory restore copy: {e}")))?;
        self.vm
            .restore_state(&checkpoint.vm_state)
            .map_err(Error::Vm)?;
        self.restore_devices(&checkpoint.devices)
            .map_err(Error::Snapshot)?;
        self.restore_vcpu_states(checkpoint.vcpu_states)?;
        self.rearm_devices();
        Ok(())
    }

    /// Capture a checkpoint for **fork**, leaving the VM **frozen** (paused, not
    /// re-armed) so its memfd-backed guest RAM is a stable CoW base that clone
    /// processes map `MAP_PRIVATE` off. Returns the `VmCheckpoint` plus the
    /// guest-RAM region descriptors (gpa/len/memfd-fd/offset) a clone needs to
    /// reach the fds via `/proc/<this_pid>/fd`. The caller keeps this process
    /// alive as the golden base.
    #[cfg(fork_supported)]
    pub fn checkpoint_for_fork(
        &mut self,
    ) -> Result<(VmCheckpoint, Vec<snapshot::MemfdRegionDesc>)> {
        // Validate the immutable prerequisite before pausing vCPUs or draining
        // device workers. A failed FORK must not alter the running VM.
        let descs = snapshot::memfd_region_descs(&self.guest_memory);
        if !has_memfd_backed_memory(&descs) {
            return Err(Error::ForkRequiresMemfd);
        }

        self.pause()?;
        self.quiesce_devices();
        let checkpoint = (|| {
            let vcpu_states = self.save_vcpu_states()?;
            let vm_state = self.vm.save_state().map_err(Error::Vm)?;
            let devices = self.snapshot_devices();
            #[cfg(all(target_arch = "x86_64", target_os = "windows"))]
            let ioapic = self.intc.lock().unwrap().save_state();
            #[cfg(not(all(target_arch = "x86_64", target_os = "windows")))]
            let ioapic = None;
            Ok(VmCheckpoint {
                vm_state,
                vcpu_states,
                devices,
                ioapic,
            })
        })();

        match checkpoint {
            // Intentionally leave a successful checkpoint frozen as the stable
            // shared CoW base for its clones.
            Ok(checkpoint) => Ok((checkpoint, descs)),
            Err(checkpoint_error) => match self.resume() {
                Ok(()) => Err(checkpoint_error),
                Err(resume_error) => Err(Error::Snapshot(format!(
                    "fork checkpoint failed ({checkpoint_error}); rollback failed ({resume_error})"
                ))),
            },
        }
    }

    /// Reapply a successful fork checkpoint and resume the original VM.
    ///
    /// Fork capture intentionally leaves device workers stopped and vCPUs
    /// paused. A failure after capture must restore the captured VM, vCPU, and
    /// device state before resuming; a plain resume can leave guest I/O wedged.
    #[cfg(fork_supported)]
    pub fn rollback_fork_checkpoint(&mut self, checkpoint: VmCheckpoint) -> Result<()> {
        if self.run_state != VmmRunState::Paused {
            return Err(Error::Snapshot(
                "fork rollback requires a paused VM".to_string(),
            ));
        }

        let VmCheckpoint {
            vm_state,
            vcpu_states,
            devices,
            ioapic,
        } = checkpoint;
        let restore = (|| {
            self.vm.restore_state(&vm_state).map_err(Error::Vm)?;
            #[cfg(all(target_arch = "x86_64", target_os = "windows"))]
            if let Some(bytes) = &ioapic {
                self.intc.lock().unwrap().restore_state(bytes);
            }
            #[cfg(not(all(target_arch = "x86_64", target_os = "windows")))]
            let _ = ioapic;
            self.restore_devices(&devices).map_err(Error::Snapshot)?;
            self.restore_vcpu_states(vcpu_states)
        })();

        // Attempt both halves even if state restoration fails, so recoverable
        // errors do not silently strand stopped workers and paused vCPUs.
        self.rearm_devices();
        let resume = self.resume();
        match (restore, resume) {
            (Ok(()), Ok(())) => Ok(()),
            (Err(restore_error), Ok(())) => Err(restore_error),
            (Ok(()), Err(resume_error)) => Err(resume_error),
            (Err(restore_error), Err(resume_error)) => Err(Error::Snapshot(format!(
                "fork rollback restore failed ({restore_error}); resume failed ({resume_error})"
            ))),
        }
    }

    /// Pause the microVM.
    pub fn pause(&mut self) -> Result<()> {
        match self.run_state {
            VmmRunState::Paused => Ok(()),
            VmmRunState::Running | VmmRunState::Resuming => {
                self.run_state = VmmRunState::Pausing;
                match self.pause_vcpus() {
                    Ok(()) => {
                        self.run_state = VmmRunState::Paused;
                        self.paused_at = Some(Instant::now());
                        Ok(())
                    }
                    Err(e) => {
                        if self.resume_vcpus(Duration::ZERO).is_ok() {
                            self.run_state = VmmRunState::Running;
                            self.paused_at = None;
                        }
                        Err(e)
                    }
                }
            }
            VmmRunState::Pausing => Err(Error::VcpuPause),
        }
    }

    /// Resume the microVM.
    pub fn resume(&mut self) -> Result<()> {
        match self.run_state {
            VmmRunState::Running => Ok(()),
            VmmRunState::Paused | VmmRunState::Pausing => {
                // A fork checkpoint drains and stops virtio workers while the
                // golden is frozen. A failed fork may resume that golden, so
                // restore device service before its vCPUs can issue more I/O.
                // Ordinary PAUSE/RESUME leaves this as a no-op.
                self.rearm_devices();
                self.run_state = VmmRunState::Resuming;
                let paused_duration = self
                    .paused_at
                    .map(|paused_at| paused_at.elapsed())
                    .unwrap_or_default();
                match self
                    .adjust_guest_clock_after_pause(paused_duration)
                    .and_then(|_| self.resume_vcpus(paused_duration))
                {
                    Ok(()) => {
                        self.run_state = VmmRunState::Running;
                        self.paused_at = None;
                        Ok(())
                    }
                    Err(e) => {
                        if self.pause_vcpus().is_ok() {
                            self.run_state = VmmRunState::Paused;
                            self.paused_at = Some(Instant::now());
                        }
                        Err(e)
                    }
                }
            }
            VmmRunState::Resuming => Err(Error::VcpuResume),
        }
    }

    /// Return the current runtime state.
    pub fn run_state(&self) -> &'static str {
        match self.run_state {
            VmmRunState::Paused => "paused",
            VmmRunState::Running => "running",
            VmmRunState::Pausing => "pausing",
            VmmRunState::Resuming => "resuming",
        }
    }

    /// Configures the system for boot.
    pub fn configure_system(
        &self,
        vcpus: &[Vcpu],
        _intc: &IrqChip,
        initrd: &Option<InitrdConfig>,
        _smbios_oem_strings: &Option<Vec<String>>,
        _pvh: bool,
    ) -> Result<()> {
        #[cfg(target_arch = "x86_64")]
        {
            let cmdline_len = if cfg!(feature = "tee") {
                arch::x86_64::layout::CMDLINE_SEV_SIZE
            } else {
                self.kernel_cmdline.len() + 1
            };

            arch::x86_64::configure_system(
                &self.guest_memory,
                &self.arch_memory_info,
                vm_memory::GuestAddress(arch::x86_64::layout::CMDLINE_START),
                cmdline_len,
                initrd,
                vcpus.len() as u8,
                _pvh,
            )
            .map_err(Error::ConfigureSystem)?;
        }

        #[cfg(target_arch = "aarch64")]
        {
            let vcpu_mpidr = vcpus.iter().map(|cpu| cpu.get_mpidr()).collect();
            fdt::create_fdt(
                &self.guest_memory,
                &self.arch_memory_info,
                vcpu_mpidr,
                self.kernel_cmdline.as_str(),
                self.mmio_device_manager.get_device_info(),
                _intc,
                initrd,
            )
            .map_err(Error::SetupFDT)?;
        }

        #[cfg(target_arch = "aarch64")]
        {
            arch::aarch64::configure_system(
                &self.guest_memory,
                &self.arch_memory_info,
                _smbios_oem_strings,
            )
            .map_err(Error::ConfigureSystem)?;
        }

        #[cfg(target_arch = "riscv64")]
        {
            fdt::create_fdt(
                &self.guest_memory,
                &self.arch_memory_info,
                vcpus.len() as u32,
                self.kernel_cmdline.as_str(),
                self.mmio_device_manager.get_device_info(),
                _intc,
                initrd,
            )
            .map_err(Error::SetupFDT)?;

            arch::riscv64::configure_system(&self.guest_memory, _smbios_oem_strings)
                .map_err(Error::ConfigureSystem)?;
        }

        Ok(())
    }

    /// Returns a reference to the inner `GuestMemoryMmap` object if present, or `None` otherwise.
    pub fn guest_memory(&self) -> &GuestMemoryMmap {
        &self.guest_memory
    }

    /// Injects CTRL+ALT+DEL keystroke combo in the i8042 device.
    #[cfg(target_arch = "x86_64")]
    pub fn send_ctrl_alt_del(&mut self) -> Result<()> {
        self.pio_device_manager
            .i8042
            .lock()
            .expect("i8042 lock was poisoned")
            .trigger_ctrl_alt_del()
            .map_err(Error::I8042Error)
    }

    /// Waits for all vCPUs to exit and terminates the Firecracker process.
    pub fn stop(&mut self, exit_code: i32) {
        info!("Vmm is stopping.");

        for observer in &self.exit_observers {
            observer
                .lock()
                .expect("Poisoned mutex for exit observer")
                .on_vmm_exit();
        }

        // Exit from Firecracker using the provided exit code. Safe because we're terminating
        // the process anyway.
        unsafe {
            libc::_exit(exit_code);
        }
    }

    /// Returns a reference to the inner KVM Vm object.
    pub fn kvm_vm(&self) -> &Vm {
        &self.vm
    }

    #[cfg(target_os = "macos")]
    pub fn add_mapping(
        &self,
        reply_sender: Sender<bool>,
        host_addr: u64,
        guest_addr: u64,
        len: u64,
    ) {
        self.vm
            .add_mapping(reply_sender, host_addr, guest_addr, len);
    }

    #[cfg(target_os = "macos")]
    pub fn remove_mapping(&self, reply_sender: Sender<bool>, guest_addr: u64, len: u64) {
        self.vm.remove_mapping(reply_sender, guest_addr, len);
    }
}

#[cfg(all(test, fork_supported))]
mod fork_tests {
    use super::*;

    fn region(fd: i32) -> snapshot::MemfdRegionDesc {
        snapshot::MemfdRegionDesc {
            gpa: 0,
            len: 4096,
            fd,
            offset: 0,
            path: String::new(),
        }
    }

    #[test]
    fn fork_requires_at_least_one_shareable_ram_region() {
        assert!(!has_memfd_backed_memory(&[]));
        assert!(!has_memfd_backed_memory(&[region(-1), region(-1)]));
        assert!(has_memfd_backed_memory(&[region(-1), region(7)]));
    }
}

impl Subscriber for Vmm {
    /// Handle a read event (EPOLLIN).
    fn process(&mut self, event: &EpollEvent, _: &mut EventManager) {
        let source = event.fd();
        let event_set = event.event_set();

        if source == self.exit_evt.as_raw_fd() && event_set == EventSet::IN {
            let _ = self.exit_evt.read();
            // Query each vcpu for the exit_code.
            // If the exit_code can't be found on any vcpu, it means that the exit signal
            // has been issued by the i8042 controller in which case we exit with
            // FC_EXIT_CODE_OK.
            //
            // The exit code set up by the guest takes preference over the one reported
            // by either a vcpu or the i8042 controller.
            let vcpu_exit_code = self
                .vcpus_handles
                .iter()
                .find_map(|handle| match handle.response_receiver().try_recv() {
                    Ok(VcpuResponse::Exited(exit_code)) => Some(exit_code),
                    _ => None,
                })
                .unwrap_or(FC_EXIT_CODE_OK);
            let vmm_exit_code = self.exit_code.load(Ordering::SeqCst);
            let exit_code = if vmm_exit_code != i32::MAX {
                debug!("using vmm exit code: {vmm_exit_code}");
                vmm_exit_code
            } else {
                debug!("using vcpu exit code: {vcpu_exit_code}");
                vcpu_exit_code as i32
            };
            self.stop(exit_code);
        } else {
            error!("Spurious EventManager event for handler: Vmm");
        }
    }

    fn interest_list(&self) -> Vec<EpollEvent> {
        vec![EpollEvent::new(
            EventSet::IN,
            self.exit_evt.as_raw_fd() as u64,
        )]
    }
}
