arcbox-vm 0.4.9

//! `SandboxManager` — orchestrates sandbox microVM lifecycle.
//!
//! A sandbox is a short-lived, strongly-isolated microVM decoupled from its
//! workload: when the initial `cmd` process exits the sandbox transitions back
//! to `Ready` rather than stopping, and continues accepting `Run` calls until
//! an explicit `Stop`/`Remove` or TTL expiry.
//!
//! `create_sandbox` returns immediately with state `"starting"`.  The VM boots
//! in a background task which broadcasts a `"ready"` event on success.

use std::collections::HashMap;
use std::num::NonZeroU64;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, RwLock};
use std::time::Duration;

use chrono::{DateTime, Utc};
use fc_sdk::VmBuilder;
use fc_sdk::types::{BootSource, Drive, NetworkInterface, Vsock};
use nix::unistd::{Gid, Uid, chown};
use tokio::sync::broadcast;
use tracing::{debug, error, info, warn};
use uuid::Uuid;

use crate::boot_proto::KernelIpParam;
use crate::config::VmmConfig;
use crate::error::{Result, VmmError};
use crate::network::{NetworkAllocation, NetworkManager};
use crate::snapshot::SnapshotCatalog;
use crate::snapshot_cow::{CowHandle, CowManager};
use crate::spawn::{spawn_direct, spawn_jailer};
use crate::vsock::{self, ExecInputMsg, OutputChunk, StartCommand};

/// Unique sandbox identifier (UUID string).
pub type SandboxId = String;

// =============================================================================
// State
// =============================================================================

/// Lifecycle state of a sandbox.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SandboxState {
    /// Firecracker process spawned; VM still booting.
    Starting,
    /// VM booted and ready to accept workloads (or last workload exited).
    Ready,
    /// A workload (cmd / Run) is currently executing inside the VM.
    Running,
    /// `Stop` called; draining workload and shutting down VM.
    Stopping,
    /// VM has shut down cleanly.
    Stopped,
    /// Unrecoverable error occurred.
    Failed,
}

impl std::fmt::Display for SandboxState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Starting => write!(f, "starting"),
            Self::Ready => write!(f, "ready"),
            Self::Running => write!(f, "running"),
            Self::Stopping => write!(f, "stopping"),
            Self::Stopped => write!(f, "stopped"),
            Self::Failed => write!(f, "failed"),
        }
    }
}

// =============================================================================
// Spec types (input to SandboxManager methods)
// =============================================================================

/// Network configuration supplied at sandbox creation time.
#[derive(Debug, Clone, Default)]
pub struct SandboxNetworkSpec {
    /// `"tap"` (default) or `"none"`.
    pub mode: String,
}

/// A single bind-mount into the sandbox.
#[derive(Debug, Clone)]
pub struct SandboxMountSpec {
    pub source: String,
    pub target: String,
    pub readonly: bool,
}

/// Full sandbox creation parameters.
///
/// Fields related to workload execution (`cmd`, `env`, `working_dir`, `user`,
/// `mounts`) are **not consumed by `create_sandbox`**.  They are stored in the
/// spec for later use by `run_in_sandbox` / `exec_in_sandbox` or passed through
/// the gRPC layer.  `image` and `ssh_public_key` are reserved for future use.
#[derive(Debug, Clone, Default)]
pub struct SandboxSpec {
    /// Caller-supplied ID; auto-generated (UUID) when `None` or empty.
    pub id: Option<String>,
    /// Arbitrary key-value metadata (filtering, listing).
    pub labels: HashMap<String, String>,
    /// Kernel image path (empty = daemon default).
    pub kernel: String,
    /// Root filesystem image path (empty = daemon default).
    pub rootfs: String,
    /// Kernel command-line arguments (empty = daemon default).
    pub boot_args: String,
    /// Number of vCPUs (0 = daemon default).
    pub vcpus: u32,
    /// Memory in MiB (0 = daemon default).
    pub memory_mib: u64,
    /// OCI image reference (empty = use rootfs directly; reserved for future use).
    pub image: String,
    /// Initial command launched automatically after boot (empty = none).
    pub cmd: Vec<String>,
    /// Environment variables for the initial command.
    pub env: HashMap<String, String>,
    /// Working directory for the initial command.
    pub working_dir: String,
    /// User to run the initial command as.
    pub user: String,
    /// Bind mounts into the sandbox.
    pub mounts: Vec<SandboxMountSpec>,
    /// Network configuration.
    pub network: SandboxNetworkSpec,
    /// Auto-destroy TTL in seconds (0 = no limit).
    pub ttl_seconds: u32,
    /// SSH public key injected via MMDS (None = no SSH setup).
    pub ssh_public_key: Option<String>,
}

/// Parameters to restore a sandbox from a checkpoint.
#[derive(Debug, Clone, Default)]
pub struct RestoreSandboxSpec {
    /// Caller-supplied ID (None = auto-generate).
    pub id: Option<String>,
    /// Source checkpoint/snapshot ID.
    pub snapshot_id: String,
    /// Labels to assign to the restored sandbox.
    pub labels: HashMap<String, String>,
    /// Assign a fresh TAP + IP to the restored sandbox.
    pub network_override: bool,
    /// Auto-destroy TTL in seconds (0 = no limit).
    pub ttl_seconds: u32,
}

// =============================================================================
// Runtime instance
// =============================================================================

/// Per-sandbox runtime state.
pub struct SandboxInstance {
    /// Unique identifier.
    pub id: SandboxId,
    /// User-supplied labels.
    pub labels: HashMap<String, String>,
    /// Original creation spec.
    pub spec: SandboxSpec,
    /// Current lifecycle state.
    pub state: SandboxState,
    /// Handle to the Firecracker process.
    pub process: Option<fc_sdk::FirecrackerProcess>,
    /// Post-boot API handle (present once the VM has booted).
    pub vm: Option<Arc<fc_sdk::Vm>>,
    /// Allocated network resources.
    pub network: Option<NetworkAllocation>,
    /// Directory holding the VM's runtime files (socket, logs, metrics).
    pub vm_dir: PathBuf,
    /// Path to the Firecracker vsock Unix domain socket (host side).
    /// `None` until the VM is booted.
    pub vsock_uds_path: Option<PathBuf>,
    /// When the sandbox record was created.
    pub created_at: DateTime<Utc>,
    /// When the sandbox first became ready.
    pub ready_at: Option<DateTime<Utc>>,
    /// When the last workload exited.
    pub last_exited_at: Option<DateTime<Utc>>,
    /// Exit code of the last workload.
    pub last_exit_code: Option<i32>,
    /// Human-readable error (only set when state == `Failed`).
    pub error: Option<String>,
    /// dm-snapshot CoW handle (present when snapshot-based rootfs is active).
    pub cow_handle: Option<CowHandle>,
    /// For restored sandboxes only: the original sandbox's vm_dir, recreated
    /// so the vmstate-recorded `rootfs.link` symlink (and FC vsock socket)
    /// resolve correctly.  Removed alongside the sandbox.
    pub restore_origin_dir: Option<PathBuf>,
}

impl SandboxInstance {
    fn new(
        id: SandboxId,
        spec: SandboxSpec,
        network: Option<NetworkAllocation>,
        vm_dir: PathBuf,
    ) -> Self {
        Self {
            id,
            labels: spec.labels.clone(),
            spec,
            state: SandboxState::Starting,
            process: None,
            vm: None,
            network,
            vm_dir,
            vsock_uds_path: None,
            created_at: Utc::now(),
            ready_at: None,
            last_exited_at: None,
            last_exit_code: None,
            error: None,
            cow_handle: None,
            restore_origin_dir: None,
        }
    }

    /// Path to the Firecracker API socket for this sandbox.
    pub fn socket_path(&self) -> PathBuf {
        self.vm_dir.join("firecracker.sock")
    }
}

// =============================================================================
// Public output types (returned to callers / gRPC layer)
// =============================================================================

/// Lightweight summary for `List` operations.
pub struct SandboxSummary {
    pub id: SandboxId,
    pub state: SandboxState,
    pub labels: HashMap<String, String>,
    /// Allocated IP address (empty when network mode is `"none"`).
    pub ip_address: String,
    pub created_at: DateTime<Utc>,
}

/// Detailed sandbox state for `Inspect`.
pub struct SandboxInfo {
    pub id: SandboxId,
    pub state: SandboxState,
    pub labels: HashMap<String, String>,
    pub vcpus: u32,
    pub memory_mib: u64,
    pub network: Option<SandboxNetworkInfo>,
    pub created_at: DateTime<Utc>,
    pub ready_at: Option<DateTime<Utc>>,
    pub last_exited_at: Option<DateTime<Utc>>,
    pub last_exit_code: Option<i32>,
    pub error: Option<String>,
}

/// Network details within `SandboxInfo`.
pub struct SandboxNetworkInfo {
    pub ip_address: String,
    pub gateway: String,
    pub tap_name: String,
}

// =============================================================================
// Events
// =============================================================================

/// A sandbox lifecycle event broadcast to subscribers.
#[derive(Debug, Clone)]
pub struct SandboxEvent {
    pub sandbox_id: SandboxId,
    /// Action: `"created"` | `"ready"` | `"running"` | `"idle"` |
    ///         `"stopping"` | `"stopped"` | `"failed"` | `"removed"`
    pub action: String,
    /// Unix nanoseconds.
    pub timestamp_ns: i64,
    /// Extra context (e.g. `"exit_code"` on `"idle"`, `"error"` on `"failed"`).
    pub attributes: HashMap<String, String>,
}

impl SandboxEvent {
    fn new(sandbox_id: &str, action: &str) -> Self {
        Self {
            sandbox_id: sandbox_id.to_owned(),
            action: action.to_owned(),
            timestamp_ns: Utc::now().timestamp_nanos_opt().unwrap_or(0),
            attributes: HashMap::new(),
        }
    }

    fn with_attr(mut self, key: &str, value: &str) -> Self {
        self.attributes.insert(key.to_owned(), value.to_owned());
        self
    }
}

// =============================================================================
// Checkpoint / Restore output types
// =============================================================================

/// Info returned after a successful checkpoint.
pub struct CheckpointInfo {
    pub snapshot_id: String,
    pub snapshot_dir: String,
    pub created_at: String,
}

/// Lightweight checkpoint summary for `ListSnapshots`.
pub struct CheckpointSummary {
    pub id: String,
    /// ID of the sandbox that was checkpointed.
    pub sandbox_id: String,
    pub name: String,
    pub labels: HashMap<String, String>,
    pub snapshot_dir: String,
    pub created_at: String,
}

// =============================================================================
// SandboxManager
// =============================================================================

const EVENT_CHANNEL_CAPACITY: usize = 256;

/// Manages the full lifecycle of multiple sandbox microVMs.
pub struct SandboxManager {
    instances: Arc<RwLock<HashMap<SandboxId, Arc<Mutex<SandboxInstance>>>>>,
    network: Arc<NetworkManager>,
    snapshots: Arc<SnapshotCatalog>,
    config: Arc<VmmConfig>,
    events_tx: broadcast::Sender<SandboxEvent>,
    cow_manager: Arc<CowManager>,
}

impl SandboxManager {
    /// Create a new manager from the given configuration.
    pub fn new(config: VmmConfig) -> Result<Self> {
        let network = Arc::new(NetworkManager::new(
            &config.network.cidr,
            &config.network.gateway,
            config.network.dns.clone(),
        )?);
        let snapshots = Arc::new(SnapshotCatalog::new(&config.firecracker.data_dir));
        let (events_tx, _) = broadcast::channel(EVENT_CHANNEL_CAPACITY);
        let cow_manager = Arc::new(
            CowManager::new(&config.firecracker.data_dir)
                .map_err(|e| VmmError::Config(format!("CowManager init: {e}")))?,
        );

        // Ensure the jailer chroot base directory exists.
        if let Some(ref jc) = config.firecracker.jailer {
            let base = jc.chroot_base_dir.as_deref().unwrap_or("/srv/jailer");
            std::fs::create_dir_all(base).map_err(VmmError::Io)?;
        }

        Ok(Self {
            instances: Arc::new(RwLock::new(HashMap::new())),
            network,
            snapshots,
            config: Arc::new(config),
            events_tx,
            cow_manager,
        })
    }

    // =========================================================================
    // Core lifecycle
    // =========================================================================

    /// Create a sandbox and return immediately (state = `"starting"`).
    ///
    /// Boots the VM in a background task.  Subscribe to [`subscribe_events`]
    /// or poll [`inspect_sandbox`] to wait for state `"ready"`.
    ///
    /// Returns `(sandbox_id, ip_address)`.  The IP is pre-allocated even
    /// before the VM finishes booting.
    pub async fn create_sandbox(&self, mut spec: SandboxSpec) -> Result<(SandboxId, String)> {
        // Apply daemon defaults for fields not supplied by the caller.
        let defaults = &self.config.defaults;
        if spec.kernel.is_empty() {
            spec.kernel.clone_from(&defaults.kernel);
        }
        if spec.rootfs.is_empty() {
            spec.rootfs.clone_from(&defaults.rootfs);
        }
        if spec.boot_args.is_empty() {
            spec.boot_args.clone_from(&defaults.boot_args);
        }
        if spec.vcpus == 0 {
            spec.vcpus = defaults.vcpus as u32;
        }
        if spec.memory_mib == 0 {
            spec.memory_mib = defaults.memory_mib;
        }
        if spec.network.mode.is_empty() {
            spec.network.mode = "tap".into();
        }

        let id = spec
            .id
            .clone()
            .filter(|s| !s.is_empty())
            .unwrap_or_else(|| Uuid::new_v4().to_string());

        // Sanitize caller-supplied IDs: reject path separators and other
        // dangerous characters to prevent directory traversal.
        if id.contains('/') || id.contains('\\') || id.contains('\0') || id == "." || id == ".." {
            return Err(VmmError::Config(format!(
                "invalid sandbox ID: {id:?} (must not contain path separators)"
            )));
        }

        // Uniqueness check.
        {
            let instances = self.instances.read().unwrap();
            if instances.contains_key(&id) {
                return Err(VmmError::AlreadyExists(id));
            }
        }

        // Allocate network resources (point-to-point TAP).
        let net_alloc = if spec.network.mode == "none" {
            None
        } else {
            Some(self.network.allocate(&id)?)
        };

        let ip_address = net_alloc
            .as_ref()
            .map(|n| n.ip_address.to_string())
            .unwrap_or_default();

        // Create the VM working directory.
        let vm_dir = PathBuf::from(&self.config.firecracker.data_dir)
            .join("sandboxes")
            .join(&id);
        std::fs::create_dir_all(&vm_dir).map_err(VmmError::Io)?;

        // Insert instance in Starting state.
        let instance =
            SandboxInstance::new(id.clone(), spec.clone(), net_alloc.clone(), vm_dir.clone());
        {
            let mut instances = self.instances.write().unwrap();
            instances.insert(id.clone(), Arc::new(Mutex::new(instance)));
        }

        // Broadcast "created" event.
        let _ = self.events_tx.send(SandboxEvent::new(&id, "created"));

        // Spawn background boot task.
        {
            let instances = Arc::clone(&self.instances);
            let network = Arc::clone(&self.network);
            let config = Arc::clone(&self.config);
            let events_tx = self.events_tx.clone();
            let cow_manager = Arc::clone(&self.cow_manager);
            let id_clone = id.clone();
            let spec_clone = spec.clone();
            let net_alloc_clone = net_alloc;
            tokio::spawn(async move {
                boot_sandbox(
                    id_clone,
                    spec_clone,
                    net_alloc_clone,
                    vm_dir,
                    instances,
                    network,
                    config,
                    events_tx,
                    cow_manager,
                )
                .await;
            });
        }

        // Spawn TTL expiry task if requested.
        if spec.ttl_seconds > 0 {
            let instances = Arc::clone(&self.instances);
            let network = Arc::clone(&self.network);
            let events_tx = self.events_tx.clone();
            let config2 = Arc::clone(&self.config);
            let cow2 = Arc::clone(&self.cow_manager);
            let id2 = id.clone();
            let ttl = spec.ttl_seconds;
            tokio::spawn(async move {
                tokio::time::sleep(Duration::from_secs(ttl as u64)).await;
                remove_sandbox_impl(
                    &id2, true, &instances, &network, &events_tx, &config2, &cow2,
                )
                .await;
            });
        }

        info!(sandbox_id = %id, "sandbox create requested (async boot started)");
        Ok((id, ip_address))
    }

    /// Stop a sandbox gracefully.
    ///
    /// Sends Ctrl+Alt+Del to the guest and waits up to `timeout_seconds`
    /// (default 30 s) for the VM to shut down.
    pub async fn stop_sandbox(&self, id: &SandboxId, timeout_seconds: u32) -> Result<()> {
        let vm_handle = {
            let instance = self.get_instance(id)?;
            let mut inst = instance.lock().unwrap();
            match inst.state {
                SandboxState::Ready | SandboxState::Running => {}
                s => {
                    return Err(VmmError::WrongState {
                        id: id.clone(),
                        expected: "Ready or Running".into(),
                        actual: s.to_string(),
                    });
                }
            }
            inst.state = SandboxState::Stopping;
            inst.vm.as_ref().map(Arc::clone)
        };

        let _ = self.events_tx.send(SandboxEvent::new(id, "stopping"));

        if let Some(vm) = vm_handle {
            let timeout = if timeout_seconds > 0 {
                timeout_seconds
            } else {
                30
            };
            // Ignore errors — VM may have already exited.
            let _ =
                tokio::time::timeout(Duration::from_secs(timeout as u64), vm.send_ctrl_alt_del())
                    .await;
        }

        // Force-kill the Firecracker process if it is still alive.
        {
            let instance = self.get_instance(id)?;
            let mut inst = instance.lock().unwrap();
            if let Some(ref mut proc) = inst.process
                && let Some(pid) = proc.pid()
                && pid > 0
            {
                let _ = nix::sys::signal::kill(
                    #[allow(clippy::cast_possible_wrap)]
                    nix::unistd::Pid::from_raw(pid as i32),
                    nix::sys::signal::Signal::SIGKILL,
                );
            }
            inst.state = SandboxState::Stopped;
        }

        let _ = self.events_tx.send(SandboxEvent::new(id, "stopped"));
        info!(sandbox_id = %id, "sandbox stopped");
        Ok(())
    }

    /// Forcibly destroy a sandbox and release all resources immediately.
    pub async fn remove_sandbox(&self, id: &SandboxId, force: bool) -> Result<()> {
        // Verify the sandbox exists.
        let state = {
            let instance = self.get_instance(id)?;
            instance.lock().unwrap().state
        };

        if !force && state == SandboxState::Running {
            return Err(VmmError::WrongState {
                id: id.clone(),
                expected: "non-running (pass force=true to override)".into(),
                actual: state.to_string(),
            });
        }

        remove_sandbox_impl(
            id,
            force,
            &self.instances,
            &self.network,
            &self.events_tx,
            &self.config,
            &self.cow_manager,
        )
        .await;
        info!(sandbox_id = %id, "sandbox removed");
        Ok(())
    }

    /// Return the current state and metadata of a sandbox.
    pub fn inspect_sandbox(&self, id: &SandboxId) -> Result<SandboxInfo> {
        let instance = self.get_instance(id)?;
        let inst = instance.lock().unwrap();
        Ok(inst_to_info(&inst))
    }

    /// List sandboxes, optionally filtered by state string and/or labels.
    pub fn list_sandboxes(
        &self,
        state_filter: Option<&str>,
        label_filter: &HashMap<String, String>,
    ) -> Vec<SandboxSummary> {
        self.instances
            .read()
            .unwrap()
            .values()
            .filter_map(|arc| {
                let inst = arc.lock().unwrap();
                // State filter.
                if let Some(sf) = state_filter
                    && !sf.is_empty()
                    && inst.state.to_string() != sf
                {
                    return None;
                }
                // Label filter: all supplied key-value pairs must match.
                for (k, v) in label_filter {
                    if inst.labels.get(k).map(String::as_str) != Some(v.as_str()) {
                        return None;
                    }
                }
                Some(SandboxSummary {
                    id: inst.id.clone(),
                    state: inst.state,
                    labels: inst.labels.clone(),
                    ip_address: inst
                        .network
                        .as_ref()
                        .map(|n| n.ip_address.to_string())
                        .unwrap_or_default(),
                    created_at: inst.created_at,
                })
            })
            .collect()
    }

    /// Subscribe to sandbox lifecycle events.
    pub fn subscribe_events(&self) -> broadcast::Receiver<SandboxEvent> {
        self.events_tx.subscribe()
    }

    // =========================================================================
    // Workload execution (requires guest agent via vsock)
    // =========================================================================

    /// Run a command inside a ready sandbox and stream its output.
    ///
    /// The sandbox must be in `Ready` state.  It transitions to `Running`
    /// immediately and back to `Ready` (emitting an `"idle"` event) when the
    /// command exits.
    ///
    /// Returns a channel receiver yielding [`OutputChunk`]s.  The final chunk
    /// has `stream == "exit"` and carries the process exit code.
    #[allow(clippy::too_many_arguments)]
    pub async fn run_in_sandbox(
        &self,
        id: &SandboxId,
        cmd: Vec<String>,
        env: HashMap<String, String>,
        working_dir: String,
        user: String,
        tty: bool,
        tty_size: Option<(u16, u16)>,
        timeout_seconds: u32,
    ) -> Result<tokio::sync::mpsc::Receiver<Result<OutputChunk>>> {
        let uds_path = self.require_ready_vsock(id)?;

        let start = StartCommand {
            cmd,
            env,
            working_dir,
            user,
            tty,
            tty_width: tty_size.map_or(80, |(w, _)| w),
            tty_height: tty_size.map_or(24, |(_, h)| h),
            timeout_seconds,
        };

        let inner_rx = vsock::run(&uds_path, start).await?;

        // Transition to Running only after vsock session is established.
        {
            let inst = self.get_instance(id)?;
            inst.lock().unwrap().state = SandboxState::Running;
        }
        let _ = self.events_tx.send(SandboxEvent::new(id, "running"));

        // Wrap the receiver to intercept MSG_EXIT and update state.
        let (wrapped_tx, wrapped_rx) = tokio::sync::mpsc::channel(64);
        let instances = Arc::clone(&self.instances);
        let events_tx = self.events_tx.clone();
        let sandbox_id = id.clone();
        tokio::spawn(async move {
            let mut inner_rx = inner_rx;
            while let Some(result) = inner_rx.recv().await {
                let send_result = match &result {
                    Ok(chunk) if chunk.stream == "exit" => {
                        let exit_code = chunk.exit_code;
                        let value = instances.read().unwrap().get(&sandbox_id).cloned();
                        if let Some(arc) = value {
                            let mut inst = arc.lock().unwrap();
                            inst.state = SandboxState::Ready;
                            inst.last_exit_code = Some(exit_code);
                            inst.last_exited_at = Some(Utc::now());
                        }
                        let _ = events_tx.send(
                            SandboxEvent::new(&sandbox_id, "idle")
                                .with_attr("exit_code", &exit_code.to_string()),
                        );
                        wrapped_tx.send(result).await
                    }
                    _ => wrapped_tx.send(result).await,
                };
                if send_result.is_err() {
                    break;
                }
            }
        });

        Ok(wrapped_rx)
    }

    /// Start an interactive exec session inside a ready sandbox.
    ///
    /// The sandbox must be in `Ready` state.  It transitions to `Running`
    /// immediately and back to `Ready` when the session ends.
    ///
    /// Returns `(input_sender, output_receiver)`:
    /// - Push [`ExecInputMsg`]s (stdin bytes, TTY resize, EOF) into `input_sender`.
    /// - Read [`OutputChunk`]s from `output_receiver` for stdout, stderr, and exit.
    #[allow(clippy::too_many_arguments)]
    pub async fn exec_in_sandbox(
        &self,
        id: &SandboxId,
        cmd: Vec<String>,
        env: HashMap<String, String>,
        working_dir: String,
        user: String,
        tty: bool,
        tty_size: Option<(u16, u16)>,
        timeout_seconds: u32,
    ) -> Result<(
        tokio::sync::mpsc::Sender<ExecInputMsg>,
        tokio::sync::mpsc::Receiver<Result<OutputChunk>>,
    )> {
        let uds_path = self.require_ready_vsock(id)?;

        let start = StartCommand {
            cmd,
            env,
            working_dir,
            user,
            tty,
            tty_width: tty_size.map_or(80, |(w, _)| w),
            tty_height: tty_size.map_or(24, |(_, h)| h),
            timeout_seconds,
        };

        let (in_tx, inner_rx) = vsock::exec(&uds_path, start).await?;

        // Transition to Running only after vsock session is established.
        {
            let inst = self.get_instance(id)?;
            inst.lock().unwrap().state = SandboxState::Running;
        }
        let _ = self.events_tx.send(SandboxEvent::new(id, "running"));

        // Wrap the output receiver to intercept MSG_EXIT and update state.
        let (wrapped_tx, wrapped_rx) = tokio::sync::mpsc::channel(64);
        let instances = Arc::clone(&self.instances);
        let events_tx = self.events_tx.clone();
        let sandbox_id = id.clone();
        tokio::spawn(async move {
            let mut inner_rx = inner_rx;
            while let Some(result) = inner_rx.recv().await {
                let send_result = match &result {
                    Ok(chunk) if chunk.stream == "exit" => {
                        let exit_code = chunk.exit_code;
                        let value = instances.read().unwrap().get(&sandbox_id).cloned();
                        if let Some(arc) = value {
                            let mut inst = arc.lock().unwrap();
                            inst.state = SandboxState::Ready;
                            inst.last_exit_code = Some(exit_code);
                            inst.last_exited_at = Some(Utc::now());
                        }
                        let _ = events_tx.send(
                            SandboxEvent::new(&sandbox_id, "idle")
                                .with_attr("exit_code", &exit_code.to_string()),
                        );
                        wrapped_tx.send(result).await
                    }
                    _ => wrapped_tx.send(result).await,
                };
                if send_result.is_err() {
                    break;
                }
            }
        });

        Ok((in_tx, wrapped_rx))
    }

    // =========================================================================
    // Checkpoint / Restore
    // =========================================================================

    /// Pause, checkpoint, and resume a sandbox.
    ///
    /// The sandbox must be in `Ready` state (no active workload).
    pub async fn checkpoint_sandbox(
        &self,
        sandbox_id: &SandboxId,
        name: String,
    ) -> Result<CheckpointInfo> {
        // Verify state and capture the kernel/rootfs paths for jailer re-staging.
        let (kernel_path, rootfs_path) = {
            let instance = self.get_instance(sandbox_id)?;
            let inst = instance.lock().unwrap();
            if inst.state != SandboxState::Ready {
                return Err(VmmError::WrongState {
                    id: sandbox_id.clone(),
                    expected: "Ready".into(),
                    actual: inst.state.to_string(),
                });
            }
            // Only needed for jailer mode; safe to capture regardless.
            (inst.spec.kernel.clone(), inst.spec.rootfs.clone())
        };

        let vm = self.get_vm_handle(sandbox_id)?;

        // Pause before snapshotting.
        vm.pause().await.map_err(VmmError::from)?;

        let snapshot_id = Uuid::new_v4().to_string();

        // In jailer mode FC runs inside a chroot and can only write to paths
        // within that chroot.  We create a temporary snapshot directory inside
        // the chroot, pass the chroot-relative paths to FC, then move the
        // resulting files to the standard catalog location on the host.
        let (fc_vmstate_path, fc_mem_path, chroot_snap_dir_opt) =
            if let Some(ref jc) = self.config.firecracker.jailer {
                let base = jc.chroot_base_dir.as_deref().unwrap_or("/srv/jailer");
                let cr = chroot_root(&self.config.firecracker.binary, base, sandbox_id);
                let chroot_snap = cr.join("snapshots").join(&snapshot_id);
                std::fs::create_dir_all(&chroot_snap).map_err(VmmError::Io)?;
                // Firecracker runs as jc.uid/jc.gid; chown the directory so it
                // can create the snapshot files.
                let uid = nix::unistd::Uid::from_raw(jc.uid);
                let gid = nix::unistd::Gid::from_raw(jc.gid);
                nix::unistd::chown(&chroot_snap, Some(uid), Some(gid))
                    .map_err(|e| VmmError::Process(format!("chown snapshot dir: {e}")))?;
                // Paths as seen by Firecracker inside the chroot.
                let fc_vmstate = format!("/snapshots/{snapshot_id}/vmstate");
                let fc_mem = format!("/snapshots/{snapshot_id}/mem");
                (fc_vmstate, fc_mem, Some(chroot_snap))
            } else {
                let snap_dir = self.snapshots.prepare_dir(sandbox_id, &snapshot_id)?;
                (
                    snap_dir.join("vmstate").to_str().unwrap().to_owned(),
                    snap_dir.join("mem").to_str().unwrap().to_owned(),
                    None,
                )
            };

        let snap_result = vm.create_snapshot(&fc_vmstate_path, &fc_mem_path).await;

        // Always resume regardless of snapshot success.
        let _ = vm.resume().await;

        snap_result.map_err(VmmError::from)?;

        // If jailer mode, move snapshot files from chroot to the catalog dir.
        let (vmstate_path, mem_path) = if let Some(chroot_snap) = chroot_snap_dir_opt {
            let catalog_dir = self.snapshots.prepare_dir(sandbox_id, &snapshot_id)?;
            let dst_vmstate = catalog_dir.join("vmstate");
            let dst_mem = catalog_dir.join("mem");
            tokio::fs::rename(chroot_snap.join("vmstate"), &dst_vmstate)
                .await
                .map_err(VmmError::Io)?;
            if chroot_snap.join("mem").exists() {
                tokio::fs::rename(chroot_snap.join("mem"), &dst_mem)
                    .await
                    .map_err(VmmError::Io)?;
            }
            let _ = tokio::fs::remove_dir_all(&chroot_snap).await;
            (dst_vmstate, dst_mem)
        } else {
            let snap_dir = self.snapshots.prepare_dir(sandbox_id, &snapshot_id)?;
            (snap_dir.join("vmstate"), snap_dir.join("mem"))
        };

        // Store kernel/rootfs template paths so restore can re-derive them.
        // Jailer mode needs them for chroot staging; direct mode needs the
        // rootfs path to set up a fresh dm-snapshot and retarget the
        // vmstate-recorded symlink.
        let meta = self.snapshots.register(
            sandbox_id,
            Some(name),
            crate::config::SnapshotType::Full,
            vmstate_path,
            Some(mem_path),
            None,
            Some(kernel_path),
            Some(rootfs_path),
        )?;

        let snap_dir_path = meta
            .vmstate_path
            .parent()
            .map(|p| p.to_string_lossy().into_owned())
            .unwrap_or_default();

        info!(sandbox_id, snapshot_id = %meta.id, "sandbox checkpointed");
        Ok(CheckpointInfo {
            snapshot_id: meta.id,
            snapshot_dir: snap_dir_path,
            created_at: meta.created_at.to_rfc3339(),
        })
    }

    /// Restore a new sandbox from a previously created checkpoint.
    ///
    /// The restored sandbox starts in `Ready` state immediately.
    ///
    /// Returns `(sandbox_id, ip_address)`.
    pub async fn restore_sandbox(&self, spec: RestoreSandboxSpec) -> Result<(SandboxId, String)> {
        let new_id = spec
            .id
            .clone()
            .filter(|s| !s.is_empty())
            .unwrap_or_else(|| Uuid::new_v4().to_string());

        if new_id.contains('/')
            || new_id.contains('\\')
            || new_id.contains('\0')
            || new_id == "."
            || new_id == ".."
        {
            return Err(VmmError::Config(format!(
                "invalid sandbox ID: {new_id:?} (must not contain path separators)"
            )));
        }

        // Uniqueness check.
        {
            let instances = self.instances.read().unwrap();
            if instances.contains_key(&new_id) {
                return Err(VmmError::AlreadyExists(new_id.clone()));
            }
        }

        // Allocate network if requested.
        let net_alloc = if spec.network_override {
            Some(self.network.allocate(&new_id)?)
        } else {
            None
        };

        let ip_address = net_alloc
            .as_ref()
            .map(|n| n.ip_address.to_string())
            .unwrap_or_default();

        // Create working directory.
        let vm_dir = PathBuf::from(&self.config.firecracker.data_dir)
            .join("sandboxes")
            .join(&new_id);
        std::fs::create_dir_all(&vm_dir).map_err(VmmError::Io)?;
        let socket_path = vm_dir.join("firecracker.sock");

        // Locate checkpoint on disk.
        let snap_meta = self.snapshots.find_by_id(&spec.snapshot_id)?;
        let vmstate_str = snap_meta.vmstate_path.to_str().unwrap().to_owned();
        let mem_file = snap_meta.mem_path.as_ref().and_then(|p| {
            if p.exists() {
                Some(p.to_str().unwrap().to_owned())
            } else {
                None
            }
        });

        let fc_cfg = &self.config.firecracker;

        // Track resources that need cleanup if anything between this point
        // and the final instance registration fails:
        //
        // - `pending_cow`: a CowHandle has no Drop impl, so a `?` propagating
        //   the error would silently leak the dm device + loop + COW file.
        // - `pending_origin_dir` (direct mode only): we recreate the original
        //   sandbox's vm_dir so the vmstate-recorded symlink + vsock paths
        //   resolve.  Set as soon as the dir is created so any later failure
        //   (CoW setup, fc_sdk::restore) cleans it up, not only the CoW Ok
        //   branch.
        //
        // On success, both are moved onto the SandboxInstance.
        let mut pending_cow: Option<CowHandle> = None;
        let mut pending_origin_dir: Option<PathBuf> = None;

        // Determine the actual host-side vsock UDS path FC will bind to on restore
        // and ensure the socket path is clear before spawning.
        //
        // - Jailer mode: each sandbox has its own chroot; FC sees `/run/firecracker.vsock`
        //   which maps to `{chroot_root}/run/firecracker.vsock` on the host. No conflict
        //   between sandboxes — we just ensure the `run/` directory exists.
        // - Direct mode: the vmstate stores the ABSOLUTE host path from the original
        //   sandbox. We must recreate that directory and delete any stale socket so FC
        //   can bind successfully.
        let (mut process, actual_vsock_path) = if let Some(ref jc) = fc_cfg.jailer {
            let base = jc.chroot_base_dir.as_deref().unwrap_or("/srv/jailer");
            let cr = chroot_root(&fc_cfg.binary, base, &new_id);
            // Ensure the `run/` directory exists inside the new chroot so FC can
            // create the vsock socket there on restore.
            let run_dir = cr.join("run");
            std::fs::create_dir_all(&run_dir).map_err(VmmError::Io)?;
            let vsock_path = cr.join("run/firecracker.vsock");
            let _ = std::fs::remove_file(&vsock_path);

            let proc = spawn_jailer(jc, fc_cfg, &new_id).await?;
            (proc, vsock_path)
        } else {
            // Direct mode: the vmstate embeds the original sandbox's full vsock
            // path.  Firecracker cannot override this on restore, so we must
            // reuse the original directory.  This means concurrent restores from
            // the same checkpoint (or restoring while the source sandbox is
            // still running) will conflict on the vsock socket.
            let original_vm_dir = PathBuf::from(&fc_cfg.data_dir)
                .join("sandboxes")
                .join(&snap_meta.vm_id);
            let original_vsock_path = original_vm_dir.join("firecracker.vsock");

            // Guard: if the vsock socket is already in use (source sandbox or
            // another restore is still running), reject early.
            if original_vsock_path.exists() {
                // Check if the socket is live by attempting a non-blocking connect.
                if std::os::unix::net::UnixStream::connect(&original_vsock_path).is_ok() {
                    return Err(VmmError::Vsock(format!(
                        "vsock path {} is already in use by another sandbox; \
                         direct-mode restore does not support concurrent restores \
                         from the same checkpoint",
                        original_vsock_path.display(),
                    )));
                }
                // Stale socket file — safe to remove.
                let _ = std::fs::remove_file(&original_vsock_path);
            }

            if let Err(e) = std::fs::create_dir_all(&original_vm_dir)
                && e.kind() != std::io::ErrorKind::AlreadyExists
            {
                return Err(VmmError::Io(e));
            }
            // Track the dir for cleanup — any failure past this point (FC spawn,
            // CoW setup, fc_sdk::restore) must remove it, not only the CoW Ok
            // branch.
            pending_origin_dir = Some(original_vm_dir.clone());

            // Pre-create log/metrics files — Firecracker requires them to
            // exist at startup (same as the boot path in do_boot).
            let log_path = vm_dir.join("firecracker.log");
            let metrics_path = vm_dir.join("firecracker.metrics");
            std::fs::File::create(&log_path).map_err(VmmError::Io)?;
            std::fs::File::create(&metrics_path).map_err(VmmError::Io)?;
            let proc =
                spawn_direct(fc_cfg, &new_id, &socket_path, &log_path, &metrics_path).await?;
            (proc, original_vsock_path)
        };

        // In jailer mode the restored FC process also runs inside a chroot and
        // cannot access the catalog's host-absolute paths.  Copy the snapshot
        // files into the new sandbox's chroot and use chroot-relative paths.
        let setup_result: Result<(String, Option<String>)> = async {
            if let Some(ref jc) = fc_cfg.jailer {
                let base = jc.chroot_base_dir.as_deref().unwrap_or("/srv/jailer");
                let cr = chroot_root(&fc_cfg.binary, base, &new_id);
                let snap_in_chroot = cr.join("snapshots").join(&spec.snapshot_id);
                std::fs::create_dir_all(&snap_in_chroot).map_err(VmmError::Io)?;
                let uid = nix::unistd::Uid::from_raw(jc.uid);
                let gid = nix::unistd::Gid::from_raw(jc.gid);
                nix::unistd::chown(&snap_in_chroot, Some(uid), Some(gid))
                    .map_err(|e| VmmError::Process(format!("chown snap dir: {e}")))?;

                // Stage kernel (always copied, ~16MB).
                if let Some(k) = snap_meta.kernel_path.as_deref() {
                    stage_kernel_for_jailer(&cr, k, jc.uid, jc.gid).await?;
                }

                // Stage rootfs: try dm-snapshot + mknod, fall back to full copy.
                // Mirrors the boot path so restored sandboxes get the same CoW
                // semantics (block-level sharing of the template, sparse COW).
                if let Some(r) = snap_meta.rootfs_path.as_deref() {
                    match self.cow_manager.setup(&new_id, r).await {
                        Ok(handle) => {
                            match stage_rootfs_device_for_jailer(
                                &cr,
                                &handle.dm_device,
                                jc.uid,
                                jc.gid,
                            )
                            .await
                            {
                                Ok(_) => pending_cow = Some(handle),
                                Err(e) => {
                                    debug!(
                                        sandbox_id = %new_id,
                                        error = %e,
                                        "mknod failed on restore, falling back to rootfs copy"
                                    );
                                    self.cow_manager.teardown(&handle).await;
                                    stage_rootfs_copy_for_jailer(&cr, r, jc.uid, jc.gid).await?;
                                }
                            }
                        }
                        Err(e) => {
                            debug!(
                                sandbox_id = %new_id,
                                error = %e,
                                "dm-snapshot unavailable on restore, copying rootfs"
                            );
                            stage_rootfs_copy_for_jailer(&cr, r, jc.uid, jc.gid).await?;
                        }
                    }
                }

                // Copy vmstate into chroot.
                let dst_vmstate = snap_in_chroot.join("vmstate");
                tokio::fs::copy(&snap_meta.vmstate_path, &dst_vmstate)
                    .await
                    .map_err(VmmError::Io)?;
                nix::unistd::chown(&dst_vmstate, Some(uid), Some(gid))
                    .map_err(|e| VmmError::Process(format!("chown vmstate: {e}")))?;

                let effective_mem = if let Some(ref mf) = snap_meta.mem_path
                    && mf.exists()
                {
                    let dst_mem = snap_in_chroot.join("mem");
                    tokio::fs::copy(mf, &dst_mem).await.map_err(VmmError::Io)?;
                    nix::unistd::chown(&dst_mem, Some(uid), Some(gid))
                        .map_err(|e| VmmError::Process(format!("chown mem: {e}")))?;
                    Some(format!("/snapshots/{}/mem", spec.snapshot_id))
                } else {
                    None
                };

                Ok((
                    format!("/snapshots/{}/vmstate", spec.snapshot_id),
                    effective_mem,
                ))
            } else {
                // Direct mode: set up a fresh dm-snapshot for the restored sandbox
                // and retarget the vmstate-recorded `{original_vm_dir}/rootfs.link`
                // at the new device.  FC reopens the symlink path on restore.
                //
                // Unlike boot, direct-mode restore has no usable fallback:  the
                // vmstate-recorded path is the symlink, so the only way for FC to
                // open the rootfs is for that symlink to exist.  If dm-snapshot
                // isn't available we fail explicitly rather than silently letting
                // `fc_sdk::restore` fault on a missing file.
                let rootfs = snap_meta.rootfs_path.as_deref().ok_or_else(|| {
                    VmmError::Config(
                        "snapshot has no rootfs_path; cannot restore in direct mode".into(),
                    )
                })?;
                let handle = self.cow_manager.setup(&new_id, rootfs).await.map_err(|e| {
                    VmmError::DeviceMapper(format!(
                        "dm-snapshot setup failed during direct-mode restore: {e}"
                    ))
                })?;
                let original_vm_dir = PathBuf::from(&fc_cfg.data_dir)
                    .join("sandboxes")
                    .join(&snap_meta.vm_id);
                if let Err(e) = create_rootfs_symlink(&original_vm_dir, &handle.dm_device) {
                    self.cow_manager.teardown(&handle).await;
                    return Err(e);
                }
                pending_cow = Some(handle);

                Ok((vmstate_str, mem_file))
            }
        }
        .await;

        let (effective_vmstate, effective_mem) = match setup_result {
            Ok(x) => x,
            Err(e) => {
                // FC was spawned but hasn't yet been told to load the vmstate,
                // so it shouldn't have the dm device open.  Kill it anyway
                // before teardown so the cleanup is unconditionally safe.
                kill_and_reap_fc(&mut process).await;
                cleanup_pending_restore(
                    &self.cow_manager,
                    pending_cow,
                    pending_origin_dir.as_deref(),
                )
                .await;
                return Err(e);
            }
        };

        // Build the restore parameters.
        let mut load_params = fc_sdk::types::SnapshotLoadParams {
            snapshot_path: effective_vmstate,
            mem_file_path: effective_mem,
            mem_backend: None,
            enable_diff_snapshots: None,
            track_dirty_pages: None,
            resume_vm: Some(true),
            network_overrides: vec![],
        };

        if let Some(ref net) = net_alloc {
            load_params.network_overrides = vec![fc_sdk::types::NetworkOverride {
                iface_id: "eth0".into(),
                host_dev_name: net.tap_name.clone(),
            }];
        }

        // In jailer mode, the actual socket path is inside the chroot; use the
        // path reported by the process handle instead of vm_dir's socket_path.
        let effective_socket = process.socket_path().to_owned();
        let vm = match fc_sdk::restore(effective_socket.to_str().unwrap(), load_params).await {
            Ok(v) => Arc::new(v),
            Err(e) => {
                // FC has likely opened the dm-snapshot block device by now
                // (vmstate load reopens all recorded drives).  Kill and wait
                // before teardown so `dmsetup remove` doesn't hit EBUSY.
                kill_and_reap_fc(&mut process).await;
                cleanup_pending_restore(
                    &self.cow_manager,
                    pending_cow.take(),
                    pending_origin_dir.as_deref(),
                )
                .await;
                return Err(VmmError::from(e));
            }
        };

        // Synchronise the guest clock to the host after restore.  The sandbox
        // clock is frozen at snapshot creation time; correct it before any
        // workload runs.  A failure here is non-fatal — the sandbox is still
        // usable, just with a potentially stale clock.
        //
        // Use a short timeout so clock sync never dominates restore latency.
        // sync_clock itself has a 5s read timeout, but connect_to_agent can
        // retry for up to AGENT_READY_TIMEOUT (30s).  Cap the whole operation.
        match tokio::time::timeout(
            std::time::Duration::from_secs(10),
            vsock::sync_clock(&actual_vsock_path),
        )
        .await
        {
            Ok(Err(e)) => warn!(sandbox_id = %new_id, "clock sync after restore failed: {e}"),
            Err(_) => warn!(sandbox_id = %new_id, "clock sync after restore timed out"),
            Ok(Ok(())) => {}
        }

        // Build and register the new sandbox instance.
        let restore_spec = SandboxSpec {
            id: Some(new_id.clone()),
            labels: spec.labels,
            ttl_seconds: spec.ttl_seconds,
            ..Default::default()
        };
        let mut instance =
            SandboxInstance::new(new_id.clone(), restore_spec, net_alloc.clone(), vm_dir);
        instance.process = Some(process);
        instance.vm = Some(vm);
        instance.vsock_uds_path = Some(actual_vsock_path);
        // Hand off pending resources to the instance — they're now tracked
        // for teardown via `remove_sandbox_impl` and won't leak.
        instance.cow_handle = pending_cow.take();
        instance.restore_origin_dir = pending_origin_dir.take();
        instance.state = SandboxState::Ready;
        instance.ready_at = Some(Utc::now());

        {
            let mut instances = self.instances.write().unwrap();
            instances.insert(new_id.clone(), Arc::new(Mutex::new(instance)));
        }

        let _ = self.events_tx.send(SandboxEvent::new(&new_id, "ready"));

        // TTL expiry task.
        if spec.ttl_seconds > 0 {
            let instances = Arc::clone(&self.instances);
            let network = Arc::clone(&self.network);
            let events_tx = self.events_tx.clone();
            let config2 = Arc::clone(&self.config);
            let cow2 = Arc::clone(&self.cow_manager);
            let id2 = new_id.clone();
            let ttl = spec.ttl_seconds;
            tokio::spawn(async move {
                tokio::time::sleep(Duration::from_secs(ttl as u64)).await;
                remove_sandbox_impl(
                    &id2, true, &instances, &network, &events_tx, &config2, &cow2,
                )
                .await;
            });
        }

        info!(
            sandbox_id = %new_id,
            snapshot_id = %spec.snapshot_id,
            "sandbox restored from checkpoint"
        );
        Ok((new_id, ip_address))
    }

    /// List checkpoints, optionally filtered by origin sandbox ID.
    pub fn list_checkpoints(&self, sandbox_id: Option<&str>) -> Result<Vec<CheckpointSummary>> {
        let infos = match sandbox_id {
            Some(sid) => self.snapshots.list(sid)?,
            None => self.snapshots.list_all()?,
        };
        Ok(infos
            .into_iter()
            .map(|s| CheckpointSummary {
                id: s.id,
                sandbox_id: s.vm_id,
                name: s.name.unwrap_or_default(),
                labels: HashMap::new(),
                snapshot_dir: s
                    .vmstate_path
                    .parent()
                    .map(|p| p.to_string_lossy().into_owned())
                    .unwrap_or_default(),
                created_at: s.created_at.to_rfc3339(),
            })
            .collect())
    }

    /// Delete a checkpoint by its ID.
    pub fn delete_checkpoint(&self, snapshot_id: &str) -> Result<()> {
        self.snapshots.delete_by_id(snapshot_id)
    }

    // =========================================================================
    // Private helpers
    // =========================================================================

    fn get_instance(&self, id: &SandboxId) -> Result<Arc<Mutex<SandboxInstance>>> {
        self.instances
            .read()
            .unwrap()
            .get(id)
            .cloned()
            .ok_or_else(|| VmmError::NotFound(id.clone()))
    }

    /// Verify the sandbox is `Ready` and return its vsock UDS path.
    fn require_ready_vsock(&self, id: &SandboxId) -> Result<PathBuf> {
        let instance = self.get_instance(id)?;
        let inst = instance.lock().unwrap();
        match inst.state {
            SandboxState::Ready => {}
            s => {
                return Err(VmmError::WrongState {
                    id: id.clone(),
                    expected: "Ready".into(),
                    actual: s.to_string(),
                });
            }
        }
        inst.vsock_uds_path
            .clone()
            .ok_or_else(|| VmmError::Vsock(format!("sandbox {id} has no vsock configured")))
    }

    fn get_vm_handle(&self, id: &SandboxId) -> Result<Arc<fc_sdk::Vm>> {
        let instance = self.get_instance(id)?;
        let inst = instance.lock().unwrap();
        inst.vm
            .as_ref()
            .map(Arc::clone)
            .ok_or_else(|| VmmError::WrongState {
                id: id.clone(),
                expected: "Ready or Running (VM handle not yet available)".into(),
                actual: inst.state.to_string(),
            })
    }
}

// =============================================================================
// Background task: boot a sandbox VM
// =============================================================================

/// Spawned by `create_sandbox`; boots the Firecracker VM and updates state.
#[allow(clippy::too_many_arguments)]
async fn boot_sandbox(
    id: SandboxId,
    spec: SandboxSpec,
    net_alloc: Option<NetworkAllocation>,
    vm_dir: PathBuf,
    instances: Arc<RwLock<HashMap<SandboxId, Arc<Mutex<SandboxInstance>>>>>,
    network: Arc<NetworkManager>,
    config: Arc<VmmConfig>,
    events_tx: broadcast::Sender<SandboxEvent>,
    cow_manager: Arc<CowManager>,
) {
    match do_boot(
        &id,
        &spec,
        net_alloc.as_ref(),
        &vm_dir,
        &config,
        &cow_manager,
    )
    .await
    {
        Ok((process, vm, vsock_uds_path, cow_handle)) => {
            let ready_at = Utc::now();

            let value = instances.read().unwrap().get(&id).cloned();
            if let Some(arc) = value {
                let mut inst = arc.lock().unwrap();
                // If stop was requested while booting, do not transition to Ready.
                if inst.state == SandboxState::Stopping || inst.state == SandboxState::Stopped {
                    info!(sandbox_id = %id, "sandbox boot completed but stop was requested; staying stopped");
                    return;
                }
                inst.process = Some(process);
                inst.vm = Some(vm);
                inst.vsock_uds_path = Some(vsock_uds_path);
                inst.cow_handle = cow_handle;
                inst.state = SandboxState::Ready;
                inst.ready_at = Some(ready_at);
            }
            let _ = events_tx.send(SandboxEvent::new(&id, "ready"));
            info!(sandbox_id = %id, "sandbox booted and ready");
        }
        Err(e) => {
            let value = instances.read().unwrap().get(&id).cloned();
            if let Some(arc) = value {
                let mut inst = arc.lock().unwrap();
                inst.state = SandboxState::Failed;
                inst.error = Some(e.to_string());
            }
            // Release network on boot failure.
            if let Some(ref net) = net_alloc {
                network.release(net);
            }
            let _ =
                events_tx.send(SandboxEvent::new(&id, "failed").with_attr("error", &e.to_string()));
            error!(sandbox_id = %id, error = %e, "sandbox boot failed");
        }
    }
}

/// Compute the host-side absolute path to the jailer chroot root directory.
///
/// Returns `{chroot_base_dir}/{fc_binary_filename}/{id}/root`.
fn chroot_root(fc_binary: &str, chroot_base_dir: &str, id: &str) -> PathBuf {
    let exec_name = Path::new(fc_binary)
        .file_name()
        .expect("fc_binary must have a filename")
        .to_string_lossy();
    PathBuf::from(chroot_base_dir)
        .join(exec_name.as_ref())
        .join(id)
        .join("root")
}

/// Copy kernel into the jailer chroot and set ownership.
///
/// Returns the chroot-relative kernel path (e.g. `"/vmlinux"`).
async fn stage_kernel_for_jailer(
    chroot_root: &Path,
    kernel_src: &str,
    uid: u32,
    gid: u32,
) -> Result<String> {
    tokio::fs::create_dir_all(chroot_root)
        .await
        .map_err(VmmError::Io)?;
    let kernel_dst = chroot_root.join("vmlinux");
    tokio::fs::copy(kernel_src, &kernel_dst)
        .await
        .map_err(VmmError::Io)?;
    chown(
        &kernel_dst,
        Some(Uid::from_raw(uid)),
        Some(Gid::from_raw(gid)),
    )
    .map_err(|e| VmmError::Process(format!("chown kernel: {e}")))?;
    Ok("/vmlinux".to_string())
}

/// Copy rootfs into the jailer chroot and set ownership.
///
/// Returns the chroot-relative rootfs path (e.g. `"/rootfs.ext4"`).
async fn stage_rootfs_copy_for_jailer(
    chroot_root: &Path,
    rootfs_src: &str,
    uid: u32,
    gid: u32,
) -> Result<String> {
    tokio::fs::create_dir_all(chroot_root)
        .await
        .map_err(VmmError::Io)?;
    let rootfs_dst = chroot_root.join("rootfs.ext4");
    // Remove any stale entry — a previous crash or a failed mknod-then-chown
    // fallback may have left a block device node here, in which case
    // `tokio::fs::copy` would write into the device instead of replacing it.
    if let Err(e) = tokio::fs::remove_file(&rootfs_dst).await
        && e.kind() != std::io::ErrorKind::NotFound
    {
        return Err(VmmError::Io(e));
    }
    tokio::fs::copy(rootfs_src, &rootfs_dst)
        .await
        .map_err(VmmError::Io)?;
    chown(
        &rootfs_dst,
        Some(Uid::from_raw(uid)),
        Some(Gid::from_raw(gid)),
    )
    .map_err(|e| VmmError::Process(format!("chown rootfs: {e}")))?;
    Ok("/rootfs.ext4".to_string())
}

/// Create a block device node in the jailer chroot pointing to a dm device.
///
/// Returns the chroot-relative rootfs path (`"/rootfs.ext4"`).
async fn stage_rootfs_device_for_jailer(
    chroot_root: &Path,
    dm_device: &str,
    uid: u32,
    gid: u32,
) -> Result<String> {
    tokio::fs::create_dir_all(chroot_root)
        .await
        .map_err(VmmError::Io)?;
    let (major, minor) = crate::snapshot_cow::device_major_minor(dm_device).await?;
    let node_path = chroot_root.join("rootfs.ext4");
    // Remove any leftover entry from a previous crash so mknod can succeed
    // (and so we never end up writing into a stale device node).
    if let Err(e) = tokio::fs::remove_file(&node_path).await
        && e.kind() != std::io::ErrorKind::NotFound
    {
        return Err(VmmError::Io(e));
    }
    crate::snapshot_cow::mknod_blkdev(&node_path, major, minor).await?;
    chown(
        &node_path,
        Some(Uid::from_raw(uid)),
        Some(Gid::from_raw(gid)),
    )
    .map_err(|e| VmmError::Process(format!("chown rootfs device: {e}")))?;
    Ok("/rootfs.ext4".to_string())
}

/// Create a stable `{vm_dir}/rootfs.link` symlink pointing at the dm-snapshot
/// device.  Returns the symlink path as a string for Firecracker to use as the
/// rootfs.  The vmstate records this path verbatim, so on restore we can
/// retarget the symlink at a freshly-created dm-snapshot without FC noticing.
///
/// Removes any stale symlink first so a previous crash doesn't cause EEXIST.
fn create_rootfs_symlink(vm_dir: &Path, dm_device: &str) -> Result<String> {
    let link_path = vm_dir.join("rootfs.link");
    let _ = std::fs::remove_file(&link_path);
    std::os::unix::fs::symlink(dm_device, &link_path).map_err(VmmError::Io)?;
    link_path
        .to_str()
        .map(str::to_owned)
        .ok_or_else(|| VmmError::Config(format!("non-UTF-8 path: {}", link_path.display())))
}

/// Release dm-snapshot + recreated origin directory after a failed restore.
///
/// `CowHandle` has no Drop impl, so dropping it would leak the dm device,
/// loop device, and sparse COW file.  This must be called on every error
/// path between `cow_manager.setup` and the point where the handle is
/// handed off to the SandboxInstance.
async fn cleanup_pending_restore(
    cow_manager: &CowManager,
    cow: Option<CowHandle>,
    origin_dir: Option<&Path>,
) {
    if let Some(handle) = cow {
        cow_manager.teardown(&handle).await;
    }
    if let Some(dir) = origin_dir
        && let Err(e) = tokio::fs::remove_dir_all(dir).await
        && e.kind() != std::io::ErrorKind::NotFound
    {
        warn!(dir = %dir.display(), err = %e, "failed to clean up restore origin dir");
    }
}

/// SIGKILL Firecracker and wait for it to exit (bounded timeout).
///
/// Required before `cow_manager.teardown` on any failure path where FC may
/// have opened the dm-snapshot block device: `dmsetup remove` returns EBUSY
/// while a process still holds the device, leaking the dm device + loop +
/// sparse COW file.  `FirecrackerProcess::drop` sends SIGKILL but never
/// reaps, so by the time teardown runs FC may still be alive.
async fn kill_and_reap_fc(process: &mut fc_sdk::FirecrackerProcess) {
    if let Some(pid) = process.pid()
        && pid > 0
    {
        let _ = nix::sys::signal::kill(
            #[allow(clippy::cast_possible_wrap)]
            nix::unistd::Pid::from_raw(pid as i32),
            nix::sys::signal::Signal::SIGKILL,
        );
    }
    let _ = tokio::time::timeout(std::time::Duration::from_secs(5), process.wait()).await;
}

/// Perform the actual Firecracker boot: spawn process, configure, start VM.
///
/// Returns `(FirecrackerProcess, Arc<Vm>, vsock_uds_path, Option<CowHandle>)`
/// on success.  The `CowHandle` is `Some` whenever dm-snapshot CoW is
/// active — both direct mode and jailer mode (when the snapshot device
/// node is successfully created inside the chroot).
#[allow(clippy::type_complexity)]
async fn do_boot(
    id: &str,
    spec: &SandboxSpec,
    net_alloc: Option<&NetworkAllocation>,
    vm_dir: &Path,
    config: &VmmConfig,
    cow_manager: &CowManager,
) -> Result<(
    fc_sdk::FirecrackerProcess,
    Arc<fc_sdk::Vm>,
    PathBuf,
    Option<CowHandle>,
)> {
    let log_path = vm_dir.join("firecracker.log");
    let metrics_path = vm_dir.join("firecracker.metrics");
    // socket_path is used only for the direct (non-jailer) mode spawn.
    let socket_path = vm_dir.join("firecracker.sock");

    let fc_cfg = &config.firecracker;

    // Some Firecracker builds expect log/metrics targets to pre-exist when
    // --log-path/--metrics-path are provided. Pre-create both files to avoid
    // startup failures with ENOENT across version variants.
    if fc_cfg.jailer.is_none() {
        if let Some(parent) = log_path.parent() {
            std::fs::create_dir_all(parent).map_err(VmmError::Io)?;
        }
        std::fs::File::create(&log_path).map_err(VmmError::Io)?;
        std::fs::File::create(&metrics_path).map_err(VmmError::Io)?;
    }

    // Spawn the Firecracker process (direct or via Jailer).
    let mut process = if let Some(ref jc) = fc_cfg.jailer {
        spawn_jailer(jc, fc_cfg, id).await?
    } else {
        spawn_direct(fc_cfg, id, &socket_path, &log_path, &metrics_path).await?
    };

    // Determine kernel, rootfs, and vsock paths.
    //
    // In jailer mode the files must exist inside the chroot, and paths passed
    // to the FC API are relative to the chroot root.  In direct mode the
    // host-absolute paths from the spec are used as-is.
    let (kernel_path, rootfs_path, vsock_fc_path, vsock_host_path, cow_handle) =
        if let Some(ref jc) = fc_cfg.jailer {
            // Jailer mode: stage kernel + rootfs into chroot.
            let base = jc.chroot_base_dir.as_deref().unwrap_or("/srv/jailer");
            let cr = chroot_root(&fc_cfg.binary, base, id);

            // Kernel is always copied (small, ~16MB).
            let k = stage_kernel_for_jailer(&cr, &spec.kernel, jc.uid, jc.gid).await?;

            // Rootfs: try dm-snapshot + mknod, fall back to full copy.
            let (r, cow) = match cow_manager.setup(id, &spec.rootfs).await {
                Ok(handle) => {
                    match stage_rootfs_device_for_jailer(&cr, &handle.dm_device, jc.uid, jc.gid)
                        .await
                    {
                        Ok(path) => (path, Some(handle)),
                        Err(e) => {
                            debug!(
                                sandbox_id = %id,
                                error = %e,
                                "mknod failed, falling back to rootfs copy"
                            );
                            cow_manager.teardown(&handle).await;
                            let path =
                                stage_rootfs_copy_for_jailer(&cr, &spec.rootfs, jc.uid, jc.gid)
                                    .await?;
                            (path, None)
                        }
                    }
                }
                Err(e) => {
                    debug!(
                        sandbox_id = %id,
                        error = %e,
                        "dm-snapshot unavailable, copying rootfs into chroot"
                    );
                    let path =
                        stage_rootfs_copy_for_jailer(&cr, &spec.rootfs, jc.uid, jc.gid).await?;
                    (path, None)
                }
            };

            let vsock_host = cr.join("run/firecracker.vsock");
            (k, r, "/run/firecracker.vsock".to_string(), vsock_host, cow)
        } else {
            // Direct mode: try dm-snapshot CoW, fall back to using rootfs directly.
            // When CoW is active, create a stable `{vm_dir}/rootfs.link` symlink
            // pointing at the dm device.  Firecracker records the symlink path
            // (not the ephemeral dm device name) in the vmstate, so a restored
            // sandbox can recreate a new dm-snapshot and retarget the symlink
            // transparently.
            let (rootfs, cow) = match cow_manager.setup(id, &spec.rootfs).await {
                Ok(handle) => match create_rootfs_symlink(vm_dir, &handle.dm_device) {
                    Ok(link) => (link, Some(handle)),
                    Err(e) => {
                        cow_manager.teardown(&handle).await;
                        return Err(e);
                    }
                },
                Err(e) => {
                    debug!(
                        sandbox_id = %id,
                        error = %e,
                        "dm-snapshot unavailable, using rootfs directly"
                    );
                    (spec.rootfs.clone(), None)
                }
            };
            let vsock_path = vm_dir.join("firecracker.vsock");
            (
                spec.kernel.clone(),
                rootfs,
                vsock_path.to_str().unwrap().to_owned(),
                vsock_path,
                cow,
            )
        };

    // Configure and boot the VM.
    let vcpu_count = NonZeroU64::new(spec.vcpus.max(1) as u64)
        .ok_or_else(|| VmmError::Config("vcpus must be > 0".into()))?;

    // Append static IP configuration to boot args so the kernel configures
    // eth0 before init runs.  The guest-side vm-agent parses this back via
    // `KernelIpParam::from_str` to derive the DNS nameserver.
    let boot_args = if let Some(net) = net_alloc {
        if spec.boot_args.contains("ip=") {
            spec.boot_args.clone()
        } else {
            let ip_param = KernelIpParam {
                client: net.ip_address,
                gateway: net.gateway,
                netmask: net.netmask(),
            };
            format!("{} {ip_param}", spec.boot_args)
        }
    } else {
        spec.boot_args.clone()
    };

    let mut builder = VmBuilder::new(process.socket_path())
        .boot_source(BootSource {
            kernel_image_path: kernel_path,
            boot_args: Some(boot_args),
            initrd_path: None,
        })
        .machine_config(fc_sdk::types::MachineConfiguration {
            vcpu_count,
            #[allow(clippy::cast_possible_wrap)]
            mem_size_mib: spec.memory_mib as i64,
            smt: false,
            // Enable dirty-page tracking so checkpointing is always available.
            track_dirty_pages: true,
            cpu_template: None,
            huge_pages: None,
        })
        .drive(Drive {
            drive_id: "rootfs".into(),
            path_on_host: Some(rootfs_path),
            is_root_device: true,
            is_read_only: Some(false),
            partuuid: None,
            cache_type: fc_sdk::types::DriveCacheType::Unsafe,
            rate_limiter: None,
            io_engine: fc_sdk::types::DriveIoEngine::Sync,
            socket: None,
        });

    if let Some(net) = net_alloc {
        builder = builder.network_interface(NetworkInterface {
            iface_id: "eth0".into(),
            guest_mac: Some(net.mac_address.clone()),
            host_dev_name: net.tap_name.clone(),
            rx_rate_limiter: None,
            tx_rate_limiter: None,
        });
    }

    // Configure vsock device so the guest agent can receive connections.
    // vsock_fc_path is the path FC uses inside its own filesystem view;
    // vsock_host_path is the host-absolute path used to connect from the host.
    builder = builder.vsock(Vsock {
        // CID 3 is the conventional guest CID; each Firecracker process is
        // isolated so the same CID is safe across concurrent sandboxes.
        guest_cid: 3,
        uds_path: vsock_fc_path,
        vsock_id: None,
    });

    let vm = match builder.start().await {
        Ok(v) => Arc::new(v),
        Err(e) => {
            // Clean up dm-snapshot if boot fails after setup.
            if let Some(ref handle) = cow_handle {
                // FC has likely opened the dm-snapshot block device by this
                // point.  Kill and wait before teardown so `dmsetup remove`
                // doesn't hit EBUSY and leak the dm/loop/COW resources.
                kill_and_reap_fc(&mut process).await;
                cow_manager.teardown(handle).await;
                // The rootfs.link symlink now points at a torn-down device.
                // Remove it so subsequent retries see a clean slate.  Only
                // applies to direct mode — jailer mode uses a chroot-internal
                // device node which is removed when the chroot is destroyed.
                if fc_cfg.jailer.is_none() {
                    let _ = std::fs::remove_file(vm_dir.join("rootfs.link"));
                }
            }
            return Err(VmmError::from(e));
        }
    };
    Ok((process, vm, vsock_host_path, cow_handle))
}

// =============================================================================
// Background task: remove a sandbox
// =============================================================================

/// Shared implementation for `remove_sandbox` and TTL expiry tasks.
#[allow(clippy::type_complexity)]
async fn remove_sandbox_impl(
    id: &str,
    _force: bool,
    instances: &Arc<RwLock<HashMap<SandboxId, Arc<Mutex<SandboxInstance>>>>>,
    network: &Arc<NetworkManager>,
    events_tx: &broadcast::Sender<SandboxEvent>,
    config: &Arc<VmmConfig>,
    cow_manager: &Arc<CowManager>,
) {
    let entry = instances.read().unwrap().get(id).cloned();
    let Some(arc) = entry else {
        return;
    };

    // Kill the Firecracker process and wait for it to exit before releasing
    // network resources. TAP destruction (ioctl TUNSETPERSIST clear / ip link
    // delete fallback) fails if the TAP fd is still held by a running process.
    let mut fc_process = {
        let mut inst = arc.lock().unwrap();
        if let Some(ref mut proc) = inst.process
            && let Some(pid) = proc.pid()
            && pid > 0
        {
            let _ = nix::sys::signal::kill(
                #[allow(clippy::cast_possible_wrap)]
                nix::unistd::Pid::from_raw(pid as i32),
                nix::sys::signal::Signal::SIGKILL,
            );
        }
        inst.process.take()
    };
    // Await process exit outside the lock. Use a timeout so cleanup proceeds
    // even if the process is stuck in uninterruptible sleep after SIGKILL.
    if let Some(ref mut proc) = fc_process {
        let _ = tokio::time::timeout(std::time::Duration::from_secs(5), proc.wait()).await;
    }

    // Teardown dm-snapshot CoW device (must happen after FC process exits
    // because Firecracker holds the block device open).
    {
        let cow_handle = arc.lock().unwrap().cow_handle.take();
        if let Some(ref handle) = cow_handle {
            cow_manager.teardown(handle).await;
        }
    }

    // Release network resources (destroys TAP via ioctl).
    {
        let inst = arc.lock().unwrap();
        if let Some(ref net) = inst.network {
            network.release(net);
        }
    }

    // Clean up the jailer chroot directory if applicable.
    if let Some(ref jc) = config.firecracker.jailer {
        let base = jc.chroot_base_dir.as_deref().unwrap_or("/srv/jailer");
        let chroot_dir = chroot_root(&config.firecracker.binary, base, id);
        // Remove {base}/{exec_name}/{id}/ (parent of "root/").
        if let Some(parent) = chroot_dir.parent()
            && let Err(e) = tokio::fs::remove_dir_all(parent).await
        {
            warn!(sandbox_id = %id, err = %e, "failed to remove jailer chroot dir");
        }
    }

    // Remove the sandbox working directory (sockets, logs, etc.).
    let vm_dir = PathBuf::from(&config.firecracker.data_dir)
        .join("sandboxes")
        .join(id);
    if let Err(e) = tokio::fs::remove_dir_all(&vm_dir).await
        && e.kind() != std::io::ErrorKind::NotFound
    {
        warn!(sandbox_id = %id, err = %e, "failed to remove sandbox dir");
    }

    // For restored sandboxes: also remove the original sandbox's vm_dir,
    // which we recreated during restore to host the vmstate-recorded
    // `rootfs.link` symlink and FC vsock socket.  Without this every
    // restore-and-remove cycle would leak one orphaned directory.
    let origin_dir = arc.lock().unwrap().restore_origin_dir.clone();
    if let Some(dir) = origin_dir
        && let Err(e) = tokio::fs::remove_dir_all(&dir).await
        && e.kind() != std::io::ErrorKind::NotFound
    {
        warn!(sandbox_id = %id, err = %e, "failed to remove restore origin dir");
    }

    instances.write().unwrap().remove(id);
    let _ = events_tx.send(SandboxEvent::new(id, "removed"));
}

// =============================================================================
// Conversion helpers
// =============================================================================

fn inst_to_info(inst: &SandboxInstance) -> SandboxInfo {
    SandboxInfo {
        id: inst.id.clone(),
        state: inst.state,
        labels: inst.labels.clone(),
        vcpus: inst.spec.vcpus,
        memory_mib: inst.spec.memory_mib,
        network: inst.network.as_ref().map(|n| SandboxNetworkInfo {
            ip_address: n.ip_address.to_string(),
            gateway: n.gateway.to_string(),
            tap_name: n.tap_name.clone(),
        }),
        created_at: inst.created_at,
        ready_at: inst.ready_at,
        last_exited_at: inst.last_exited_at,
        last_exit_code: inst.last_exit_code,
        error: inst.error.clone(),
    }
}