arcbox-vm 0.4.17

use super::boot::boot_sandbox;
use super::cleanup::{inst_to_info, remove_sandbox_impl};
use super::*;

impl SandboxManager {
    pub async fn create_sandbox(&self, mut spec: SandboxSpec) -> Result<(SandboxId, String)> {
        // Do not allocate any per-id resources until the startup orphan sweep
        // has run — otherwise a re-created same-id sandbox races it.
        self.await_reconcile().await;

        // 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());

        // Restrict caller-supplied ids to a safe charset (path components,
        // jailer --id, dm/TAP names). Auto-generated UUIDs pass unchanged.
        super::validate_id("sandbox id", &id)?;

        let vm_dir = PathBuf::from(&self.config.firecracker.data_dir)
            .join("sandboxes")
            .join(&id);

        // Reserve the id atomically BEFORE allocating any per-id resource, so
        // two concurrent creates of the same id can't both pass a uniqueness
        // check and race — the loser fails with AlreadyExists instead of
        // overwriting the map entry and leaking the winner's TAP/instance
        // (mirrors restore; unwound on any early return via Drop).
        let reservation = super::reserve_id(
            &self.instances,
            &id,
            SandboxInstance::new(id.clone(), spec.clone(), None, vm_dir.clone()),
        )?;

        // Allocate network resources (point-to-point TAP). The id is claimed,
        // so a concurrent create already failed at reserve_id above.
        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; release the TAP if this fails (the
        // reservation itself unwinds the placeholder on the early return).
        if let Err(e) = std::fs::create_dir_all(&vm_dir) {
            if let Some(net) = &net_alloc {
                self.network.release(net);
            }
            return Err(VmmError::Io(e));
        }

        // Populate the reserved instance and commit it. Keep a Weak to identify
        // this exact generation when its TTL timer fires (see expire_sandbox).
        let arc = reservation.instance();
        arc.lock().unwrap().network.clone_from(&net_alloc);
        let ttl_armed_for = Arc::downgrade(&arc);
        reservation.commit();

        // 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;
            let armed_for = ttl_armed_for;
            tokio::spawn(async move {
                tokio::time::sleep(Duration::from_secs(ttl as u64)).await;
                super::cleanup::expire_sandbox(
                    &id2, &armed_for, &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.
    ///
    /// Waits up to `timeout_seconds` (default 30 s) for an active workload
    /// to exit, asks the guest to shut down (Ctrl+Alt+Del reboots the guest,
    /// which exits Firecracker), and SIGKILLs Firecracker only if it
    /// outlives the remaining budget. All runtime resources (TAP + IP,
    /// dm-snapshot CoW, jailer chroot) are released on `Stopped`; only the
    /// inspectable record and the log directory survive until `Remove`.
    pub async fn stop_sandbox(&self, id: &SandboxId, timeout_seconds: u32) -> Result<()> {
        let budget = Duration::from_secs(u64::from(if timeout_seconds > 0 {
            timeout_seconds
        } else {
            30
        }));
        let deadline = tokio::time::Instant::now() + budget;

        let (was_running, 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(),
                    });
                }
            }
            let was_running = inst.state == SandboxState::Running;
            inst.state = SandboxState::Stopping;
            (was_running, inst.vm.as_ref().map(Arc::clone))
        };

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

        // Drain: give an active workload the budget to finish. The run/exec
        // watcher flips the state to Ready (over our Stopping) when the exit
        // chunk arrives, so poll for that transition.
        if was_running {
            while tokio::time::Instant::now() < deadline {
                let state = self.get_instance(id)?.lock().unwrap().state;
                if state != SandboxState::Stopping {
                    break;
                }
                tokio::time::sleep(Duration::from_millis(100)).await;
            }
            // Re-assert Stopping for observers regardless of drain outcome.
            self.get_instance(id)?.lock().unwrap().state = SandboxState::Stopping;
        }

        // Ask the guest to shut down. Ctrl+Alt+Del triggers a guest reboot,
        // which Firecracker turns into a VM exit. Errors are ignored — the
        // VM may already be gone.
        if let Some(vm) = vm_handle {
            let _ = tokio::time::timeout(Duration::from_secs(5), vm.send_ctrl_alt_del()).await;
        }

        // Wait for Firecracker to exit within the remaining budget; SIGKILL
        // as a fallback, then reap.
        let fc_process = self.get_instance(id)?.lock().unwrap().process.take();
        if let Some(mut proc) = fc_process {
            let remaining = deadline
                .checked_duration_since(tokio::time::Instant::now())
                .unwrap_or(Duration::from_secs(1))
                .max(Duration::from_secs(1));
            if tokio::time::timeout(remaining, proc.wait()).await.is_err() {
                warn!(sandbox_id = %id, "guest did not shut down in time; killing firecracker");
                super::boot::kill_and_reap_fc(&mut proc).await;
            }
        }

        // Release TAP/IP, CoW device, and chroot now that FC is gone; the
        // record itself stays inspectable until Remove.
        {
            let instance = self.get_instance(id)?;
            super::cleanup::release_runtime_resources(
                id,
                &instance,
                &self.network,
                &self.config,
                &self.cow_manager,
            )
            .await;
            let mut inst = instance.lock().unwrap();
            inst.state = SandboxState::Stopped;
            // Every reconcilable resource is gone; drop the crash record.
            super::reconcile::clear_state_record(&inst.vm_dir);
        }

        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> {
        // Snapshot the Arcs under the map read guard, then release it before
        // locking any instance. The manager's discipline is "never hold the
        // instances map lock while holding an instance lock"; taking both here
        // (as before) is the one place that could deadlock a future writer that
        // locks in the opposite order.
        let instances: Vec<_> = self.instances.read().unwrap().values().cloned().collect();
        instances
            .iter()
            .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()
    }

    pub(super) 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.
    pub(super) 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")))
    }

    /// Verify the sandbox is alive (Ready or Running) and return its vsock
    /// UDS path. Unlike [`Self::require_ready_vsock`], an in-flight workload
    /// does not block the operation — file I/O works alongside Run/Exec.
    pub(super) fn require_alive_vsock(&self, id: &SandboxId) -> Result<PathBuf> {
        let instance = self.get_instance(id)?;
        let 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.vsock_uds_path
            .clone()
            .ok_or_else(|| VmmError::Vsock(format!("sandbox {id} has no vsock configured")))
    }

    /// Read a file from inside an alive sandbox over the vsock file channel.
    pub async fn read_sandbox_file(&self, id: &SandboxId, path: &str) -> Result<Vec<u8>> {
        let uds = self.require_alive_vsock(id)?;
        crate::file_io::read_file(&uds, path).await
    }

    /// Write a file into an alive sandbox over the vsock file channel.
    pub async fn write_sandbox_file(
        &self,
        id: &SandboxId,
        path: &str,
        mode: u32,
        data: &[u8],
    ) -> Result<()> {
        let uds = self.require_alive_vsock(id)?;
        crate::file_io::write_file(&uds, path, mode, data).await
    }

    pub(super) 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(),
            })
    }
}