supermachine/vmm/

resources.rs

//! Typed VM configuration resources.
//!
//! This is the first reusable config boundary lifted out of the supermachine-worker
//! command-line harness. Keep this free of CLI parsing and host process state
//! so it can grow into the general Supermachine VM-library input surface.

use std::fmt;

pub const DEFAULT_CMDLINE: &str = "console=ttyAMA0 reboot=t panic=-1";
pub const DEFAULT_MEMORY_MIB: usize = 256;
pub const DEFAULT_VCPUS: u32 = 1;
pub const THROUGHPUT_PROFILE_VCPUS: u32 = 4;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum VmProfile {
    Latency,
    Throughput,
}

impl VmProfile {
    pub fn parse(s: &str) -> Option<Self> {
        match s {
            "latency" | "low-latency" => Some(Self::Latency),
            "throughput" => Some(Self::Throughput),
            _ => None,
        }
    }

    pub fn default_vcpus(self) -> u32 {
        match self {
            Self::Latency => DEFAULT_VCPUS,
            Self::Throughput => THROUGHPUT_PROFILE_VCPUS,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct VmResources {
    pub kernel_path: Option<String>,
    pub initrd_path: Option<String>,
    pub cmdline: String,
    pub memory_mib: usize,
    /// Read-only block devices (squashfs layers, delta squashfs).
    /// Attached as virtio-blk before any volumes.
    pub block_devices: Vec<String>,
    /// Read-write block devices for `--volume` persistent volumes.
    /// Attached as virtio-blk after `block_devices`. Each is a host
    /// file the guest mounts as a writable filesystem.
    pub volumes: Vec<VolumeSpec>,
    /// virtio-fs DAX mounts. Each entry attaches a host directory
    /// to the guest, served via FUSE-over-virtio with DAX
    /// (zero-copy mmap'd reads, page-cache-shared across VMs).
    pub mounts: Vec<MountSpec>,
    pub vcpus: u32,
    pub restore_from: Option<String>,
    pub cow_restore: bool,
    pub snapshot: SnapshotResources,
    pub endpoints: EndpointResources,
    /// If `Some(N)`, the runner asks the guest's virtio-balloon
    /// driver to inflate by `N` 4 KiB pages after restore (or
    /// after kernel boot in the no-restore path). The host then
    /// reclaims those pages via `madvise(MADV_FREE)`, dropping
    /// per-worker idle RSS. Driven by metadata.json's
    /// `balloon_target_pages` field on the snapshot, plumbed
    /// through `--balloon-target-pages N` on the worker CLI.
    pub balloon_target_pages: Option<u32>,
    /// TSI control-channel auth token (32 bytes). When `Some`,
    /// the vsock muxer enforces that every TSI control DGRAM
    /// carries this prefix; the kernel patch 0014 + the
    /// `supermachine.tsi_token=<hex>` cmdline parameter produce
    /// matching prefixes. Per-snapshot at bake time, re-injected
    /// at restore. `None` means legacy (pre-0.6.0) snapshots
    /// that don't have a baked token; the muxer accepts all
    /// control DGRAMs in that case, matching the kernel's
    /// "auth disabled" send path.
    pub tsi_token: Option<[u8; 32]>,
}

/// A `--volume HOST_FILE:GUEST_PATH` entry. The host file is the
/// canonical store; on first use init-oci formats it ext4. Once
/// formatted, the same file persists across runs and across
/// snapshot restores. Snapshots don't capture the volume contents
/// (they live on the host); they only capture the mapping.
///
/// init-oci auto-mounts every declared volume at its `guest_path`
/// before the guest workload starts. When `guest_path` nests under
/// a virtio-fs [`MountSpec`] target (e.g. mount at `/workspace`,
/// volume at `/workspace/node_modules`), set `guest_path` on BOTH
/// — 0.7.28+ init-oci auto-mounts virtio-fs binds first, volumes
/// second, so the layering matches Docker / Apple `container`
/// conventions and child volumes correctly overlay parent binds.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct VolumeSpec {
    /// Path to the host file backing this volume. Created (sparse,
    /// `size_bytes`) if missing.
    pub host_path: String,
    /// Mount point inside the guest (e.g. `/var/lib/postgres`).
    pub guest_path: String,
    /// Size of the volume in bytes. Default 1 GiB. Files are
    /// sparse — actual disk usage matches what the guest writes.
    pub size_bytes: u64,
}

impl VolumeSpec {
    pub const DEFAULT_SIZE_BYTES: u64 = 1024 * 1024 * 1024;

    pub fn new(host_path: impl Into<String>, guest_path: impl Into<String>) -> Self {
        Self {
            host_path: host_path.into(),
            guest_path: guest_path.into(),
            size_bytes: Self::DEFAULT_SIZE_BYTES,
        }
    }

    pub fn with_size_bytes(mut self, size_bytes: u64) -> Self {
        self.size_bytes = size_bytes;
        self
    }
}

/// Per-mount symlink policy. Picks the right shape of symlink
/// support for the trust level of what's running in the VM.
///
/// All three modes use `O_NOFOLLOW` on host-side path ops; the policy
/// gates (a) whether the guest can create symlinks at all and
/// (b) whether existing host symlinks pointing OUTSIDE the mount
/// root are visible/traversable.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum SymlinkPolicy {
    /// Guest cannot create symlinks (FUSE_SYMLINK → EPERM) and cannot
    /// create hard links (FUSE_LINK → EPERM). Existing host symlinks
    /// whose canonical target leaves the mount root are rejected with
    /// EACCES at LOOKUP. Use for paranoid mounts where you want pure
    /// file content with no metadata surprises.
    Deny,
    /// Guest can create symlinks; targets are stored as opaque bytes
    /// (the host never resolves them — POSIX symlink(2) semantics).
    /// Existing external host symlinks are rejected at LOOKUP under
    /// the same rule as `Deny`. This is the safe-multi-tenant default:
    /// npm/pnpm/yarn all work, but a hostile guest can't trick the
    /// host into walking outside the mount root via a planted
    /// `/escape -> /etc/shadow` symlink.
    #[default]
    Opaque,
    /// Guest can create symlinks; external host symlinks (those
    /// resolving outside the canonical mount root) are followed
    /// unconditionally. Use ONLY for trusted single-tenant workloads
    /// where the mount tree may legitimately reference absolute host
    /// paths (e.g. `~/.cache` from a dev tree). This is the equivalent
    /// of the pre-0.5.5 `allow_external_symlinks: true`.
    Follow,
}

/// A virtio-fs DAX mount: expose a host directory to the guest at
/// `guest_tag`. init-oci auto-mounts the share at `guest_path`
/// (absolute path inside the guest) before the workload starts and
/// before any volume auto-mounts — Docker / Apple `container`
/// convention. Nested volume `guest_path`s overlay correctly.
///
/// Reads + writes flow through virtio-fs. When the guest mmaps a
/// file, FUSE_SETUPMAPPING routes the file's host pages into the
/// VM's DAX window (validated by spike 22) — zero-copy reads, host
/// page-cache shared across VMs that mount the same path.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MountSpec {
    /// Host directory to expose. Must exist; must be a directory.
    pub host_path: String,
    /// virtiofs tag the guest mounts by. Max 35 UTF-8 bytes (one
    /// reserved for the trailing NUL in the device config space).
    pub guest_tag: String,
    /// Absolute guest mount point. init-oci auto-mounts the share
    /// here before the workload starts and before any volume
    /// auto-mounts. Required — no "userspace will mount it" mode.
    pub guest_path: String,
    /// Read-only mode. When true the FUSE backend rejects WRITE +
    /// CREATE + UNLINK / MKDIR / RMDIR. (Currently always read-only
    /// at the backend level — write path lands in a follow-up
    /// slice; this field is the future API surface.)
    pub read_only: bool,
    /// Symlink policy. See [`SymlinkPolicy`] for the three modes:
    /// `Deny` (no symlinks at all), `Opaque` (default — guest may
    /// create, host stores targets verbatim, external host symlinks
    /// rejected), `Follow` (legacy `allow_external_symlinks: true`
    /// behaviour).
    pub symlinks: SymlinkPolicy,
}

impl MountSpec {
    pub fn new(
        host_path: impl Into<String>,
        guest_tag: impl Into<String>,
        guest_path: impl Into<String>,
    ) -> Self {
        Self {
            host_path: host_path.into(),
            guest_tag: guest_tag.into(),
            guest_path: guest_path.into(),
            read_only: false,
            symlinks: SymlinkPolicy::default(),
        }
    }

    pub fn read_only(mut self) -> Self {
        self.read_only = true;
        self
    }

    /// Set the symlink policy. See [`SymlinkPolicy`].
    pub fn with_symlinks(mut self, p: SymlinkPolicy) -> Self {
        self.symlinks = p;
        self
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct SnapshotResources {
    pub after_ms: Option<u64>,
    pub at_heartbeat: Option<u64>,
    pub on_listener: bool,
    /// Trigger snapshot on the per-VM
    /// [`crate::devices::serial::SerialState::pre_exec_ready`] atomic —
    /// init-oci has just printed "workload-pre-exec" and is in a 100 ms
    /// nanosleep, giving us a stable WFI window to capture in. Used
    /// by the always-pipelined-skip-warm `.build()` path: snapshots
    /// guest BEFORE the workload runs, so each restore re-execs
    /// the workload fresh. Saves 50-150 ms vs `on_listener`.
    /// Mutually exclusive with `on_listener` in practice — when both
    /// are set, pre-exec wins (it fires earlier in the boot timeline).
    pub on_pre_exec: bool,
    pub quiesce_ms: u64,
    pub out_path: Option<String>,
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct EndpointResources {
    pub vsock_mux: Option<String>,
    pub http_port: Option<String>,
    /// SCM_RIGHTS handoff acceptor path. The router connects here
    /// and passes accepted client TCP fds (plus a small prefix of
    /// already-consumed bytes) so the worker bridges directly to
    /// the guest without the router process being on the data path.
    pub vsock_mux_handoff: Option<String>,
    /// `<vsock_mux>-exec.sock` path. Each accepted unix-socket
    /// client gets bridged to a *native* AF_VSOCK connection to
    /// the guest's exec agent on `vsock_exec_guest_port` (default
    /// 1028 — see [`DEFAULT_EXEC_GUEST_PORT`]). Used by
    /// `supermachine exec` and `Vm::exec`.
    pub vsock_exec: Option<String>,
    /// Native AF_VSOCK port the guest agent listens on. None ⇒
    /// [`DEFAULT_EXEC_GUEST_PORT`].
    pub vsock_exec_guest_port: Option<u32>,
}

/// Default native AF_VSOCK port for the in-guest exec agent. Picked
/// to sit above the host-direction reserved range (env service is
/// `VSOCK_ENV_PORT=1026`, leaving room for future host-side
/// listeners) and below typical TSI vm_port allocations.
pub const DEFAULT_EXEC_GUEST_PORT: u32 = 1028;

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ResourceError {
    MissingKernel,
    ZeroMemory,
    ZeroVcpus,
    SnapshotTriggerWithoutOutput,
}

impl VmResources {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn for_kernel(kernel_path: impl Into<String>, initrd_path: impl Into<String>) -> Self {
        Self::new()
            .with_kernel_path(kernel_path)
            .with_initramfs(initrd_path)
    }

    pub fn from_snapshot(path: impl Into<String>) -> Self {
        Self::new().with_restore(path)
    }

    pub fn with_kernel_path(mut self, path: impl Into<String>) -> Self {
        self.kernel_path = Some(path.into());
        self
    }

    pub fn with_initramfs(mut self, path: impl Into<String>) -> Self {
        self.initrd_path = Some(path.into());
        self
    }

    pub fn with_cmdline(mut self, cmdline: impl Into<String>) -> Self {
        self.cmdline = cmdline.into();
        self
    }

    pub fn with_memory_mib(mut self, memory_mib: usize) -> Self {
        self.memory_mib = memory_mib;
        self
    }

    pub fn with_profile(mut self, profile: VmProfile) -> Self {
        self.apply_profile_defaults(profile);
        self
    }

    pub fn with_vcpus(mut self, vcpus: u32) -> Self {
        self.vcpus = vcpus;
        self
    }

    pub fn with_block_device(mut self, path: impl Into<String>) -> Self {
        self.block_devices.push(path.into());
        self
    }

    /// Attach a writable volume. See [`VolumeSpec`].
    pub fn with_volume(mut self, volume: VolumeSpec) -> Self {
        self.volumes.push(volume);
        self
    }

    /// Attach a virtio-fs DAX mount. See [`MountSpec`].
    pub fn with_mount(mut self, mount: MountSpec) -> Self {
        self.mounts.push(mount);
        self
    }

    pub fn with_restore(mut self, path: impl Into<String>) -> Self {
        self.restore_from = Some(path.into());
        self
    }

    pub fn with_cow_restore(mut self, enabled: bool) -> Self {
        self.cow_restore = enabled;
        self
    }

    pub fn with_snapshot_after_ms(mut self, after_ms: u64, out_path: impl Into<String>) -> Self {
        self.snapshot.after_ms = Some(after_ms);
        self.snapshot.out_path = Some(out_path.into());
        self
    }

    pub fn with_snapshot_at_heartbeat(
        mut self,
        at_heartbeat: u64,
        out_path: impl Into<String>,
    ) -> Self {
        self.snapshot.at_heartbeat = Some(at_heartbeat);
        self.snapshot.out_path = Some(out_path.into());
        self
    }

    pub fn with_snapshot_on_listener(mut self, out_path: impl Into<String>) -> Self {
        self.snapshot.on_listener = true;
        self.snapshot.out_path = Some(out_path.into());
        self
    }

    /// Trigger snapshot on init-oci's "workload-pre-exec" marker
    /// (see [`SnapshotResources::on_pre_exec`]). For the always-
    /// pipelined-skip-warm bake; with_warmup / service-image bakes
    /// keep `on_listener` instead.
    pub fn with_snapshot_on_pre_exec(mut self, out_path: impl Into<String>) -> Self {
        self.snapshot.on_pre_exec = true;
        self.snapshot.out_path = Some(out_path.into());
        self
    }

    pub fn with_quiesce_ms(mut self, quiesce_ms: u64) -> Self {
        self.snapshot.quiesce_ms = quiesce_ms;
        self
    }

    pub fn with_vsock_mux(mut self, path: impl Into<String>) -> Self {
        self.endpoints.vsock_mux = Some(path.into());
        self
    }

    pub fn with_http_port(mut self, port: impl Into<String>) -> Self {
        self.endpoints.http_port = Some(port.into());
        self
    }

    pub fn with_vsock_mux_handoff(mut self, path: impl Into<String>) -> Self {
        self.endpoints.vsock_mux_handoff = Some(path.into());
        self
    }

    /// Set the path of the `<vsock_mux>-exec.sock` frontend. See
    /// [`EndpointResources::vsock_exec`].
    pub fn with_vsock_exec(mut self, path: impl Into<String>) -> Self {
        self.endpoints.vsock_exec = Some(path.into());
        self
    }

    /// Override the default guest-side AF_VSOCK port the exec
    /// agent listens on. Match this to whatever your guest agent
    /// binds. Default is [`DEFAULT_EXEC_GUEST_PORT`].
    pub fn with_vsock_exec_guest_port(mut self, port: u32) -> Self {
        self.endpoints.vsock_exec_guest_port = Some(port);
        self
    }

    /// Set the TSI control-channel auth token (see
    /// [`VmResources::tsi_token`]). Pass `None` to disable
    /// enforcement; pass `Some([…])` to require every TSI control
    /// DGRAM to carry this 32-byte prefix.
    pub fn with_tsi_token(mut self, token: Option<[u8; 32]>) -> Self {
        self.tsi_token = token;
        self
    }

    pub fn memory_bytes(&self) -> usize {
        self.memory_mib * 1024 * 1024
    }

    pub fn is_restore(&self) -> bool {
        self.restore_from.is_some()
    }

    pub fn apply_profile_defaults(&mut self, profile: VmProfile) {
        self.vcpus = profile.default_vcpus();
    }

    pub fn validate_for_run(&self) -> Result<(), ResourceError> {
        if self.memory_mib == 0 {
            return Err(ResourceError::ZeroMemory);
        }
        if self.vcpus == 0 {
            return Err(ResourceError::ZeroVcpus);
        }
        if self.kernel_path.is_none() && self.restore_from.is_none() {
            return Err(ResourceError::MissingKernel);
        }
        let wants_snapshot = self.snapshot.after_ms.is_some()
            || self.snapshot.at_heartbeat.is_some()
            || self.snapshot.on_listener
            || self.snapshot.on_pre_exec;
        if wants_snapshot && self.snapshot.out_path.is_none() {
            return Err(ResourceError::SnapshotTriggerWithoutOutput);
        }
        Ok(())
    }
}

impl Default for VmResources {
    fn default() -> Self {
        Self {
            kernel_path: None,
            initrd_path: None,
            cmdline: DEFAULT_CMDLINE.to_string(),
            memory_mib: DEFAULT_MEMORY_MIB,
            block_devices: Vec::new(),
            volumes: Vec::new(),
            mounts: Vec::new(),
            vcpus: DEFAULT_VCPUS,
            restore_from: None,
            cow_restore: false,
            snapshot: SnapshotResources::default(),
            endpoints: EndpointResources::default(),
            balloon_target_pages: None,
            tsi_token: None,
        }
    }
}

impl fmt::Display for ResourceError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ResourceError::MissingKernel => {
                write!(f, "kernel path or restore snapshot is required")
            }
            ResourceError::ZeroMemory => write!(f, "memory must be greater than zero"),
            ResourceError::ZeroVcpus => write!(f, "vCPU count must be greater than zero"),
            ResourceError::SnapshotTriggerWithoutOutput => {
                write!(f, "snapshot trigger requires snapshot output path")
            }
        }
    }
}

impl std::error::Error for ResourceError {}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn defaults_match_supermachine_cli() {
        let resources = VmResources::default();
        assert_eq!(resources.cmdline, DEFAULT_CMDLINE);
        assert_eq!(resources.memory_mib, 256);
        assert_eq!(resources.vcpus, 1);
        assert_eq!(resources.memory_bytes(), 256 * 1024 * 1024);
    }

    #[test]
    fn profile_defaults_are_stable() {
        assert_eq!(VmProfile::parse("latency"), Some(VmProfile::Latency));
        assert_eq!(VmProfile::parse("low-latency"), Some(VmProfile::Latency));
        assert_eq!(VmProfile::parse("throughput"), Some(VmProfile::Throughput));
        assert_eq!(VmProfile::parse("unknown"), None);
        assert_eq!(VmProfile::Latency.default_vcpus(), 1);
        assert_eq!(VmProfile::Throughput.default_vcpus(), 4);
    }

    #[test]
    fn applies_profile_defaults_to_resources() {
        let mut resources = VmResources::default();
        resources.apply_profile_defaults(VmProfile::Throughput);
        assert_eq!(resources.vcpus, 4);
        resources.apply_profile_defaults(VmProfile::Latency);
        assert_eq!(resources.vcpus, 1);
    }

    #[test]
    fn convenience_constructors_cover_kernel_and_restore() {
        let kernel = VmResources::for_kernel("kernel", "initrd");
        assert_eq!(kernel.kernel_path.as_deref(), Some("kernel"));
        assert_eq!(kernel.initrd_path.as_deref(), Some("initrd"));
        assert!(!kernel.is_restore());

        let restore = VmResources::from_snapshot("snap.sm");
        assert_eq!(restore.restore_from.as_deref(), Some("snap.sm"));
        assert!(restore.is_restore());
    }

    #[test]
    fn builder_style_methods_cover_common_library_config() {
        let resources = VmResources::new()
            .with_kernel_path("kernel")
            .with_initramfs("initrd")
            .with_cmdline("console=ttyS0")
            .with_memory_mib(512)
            .with_profile(VmProfile::Throughput)
            .with_vcpus(2)
            .with_block_device("rootfs.squashfs")
            .with_snapshot_at_heartbeat(1, "snap.sm")
            .with_quiesce_ms(7)
            .with_vsock_mux("/tmp/vsock.sock")
            .with_http_port("8080");

        assert_eq!(resources.kernel_path.as_deref(), Some("kernel"));
        assert_eq!(resources.initrd_path.as_deref(), Some("initrd"));
        assert_eq!(resources.cmdline, "console=ttyS0");
        assert_eq!(resources.memory_mib, 512);
        assert_eq!(resources.vcpus, 2);
        assert_eq!(resources.block_devices, vec!["rootfs.squashfs"]);
        assert_eq!(resources.snapshot.at_heartbeat, Some(1));
        assert_eq!(resources.snapshot.quiesce_ms, 7);
        assert_eq!(resources.snapshot.out_path.as_deref(), Some("snap.sm"));
        assert_eq!(
            resources.endpoints.vsock_mux.as_deref(),
            Some("/tmp/vsock.sock")
        );
        assert_eq!(resources.endpoints.http_port.as_deref(), Some("8080"));
    }

    #[test]
    fn builder_style_restore_config_is_valid_without_kernel() {
        let resources = VmResources::new()
            .with_restore("snap.sm")
            .with_cow_restore(true);

        assert!(resources.is_restore());
        assert!(resources.cow_restore);
        assert_eq!(resources.validate_for_run(), Ok(()));
    }

    #[test]
    fn validates_kernel_or_restore() {
        let mut resources = VmResources::default();
        assert_eq!(
            resources.validate_for_run(),
            Err(ResourceError::MissingKernel)
        );

        resources.kernel_path = Some("vmlinux".to_string());
        assert_eq!(resources.validate_for_run(), Ok(()));

        resources.kernel_path = None;
        resources.restore_from = Some("snap.sm".to_string());
        assert_eq!(resources.validate_for_run(), Ok(()));
    }

    // === MountSpec / VolumeSpec — auto-mount API ===

    #[test]
    fn mount_spec_constructs_with_required_guest_path() {
        let m = MountSpec::new("/host/x", "tag", "/workspace");
        assert_eq!(m.host_path, "/host/x");
        assert_eq!(m.guest_tag, "tag");
        assert_eq!(m.guest_path, "/workspace");
        assert_eq!(m.symlinks, SymlinkPolicy::Opaque);
        assert!(!m.read_only);
    }

    #[test]
    fn mount_spec_symlinks_composes_with_construction() {
        let m = MountSpec::new("/h", "t", "/w").with_symlinks(SymlinkPolicy::Deny);
        assert_eq!(m.guest_path, "/w");
        assert_eq!(m.symlinks, SymlinkPolicy::Deny);
    }

    #[test]
    fn mount_spec_read_only_composes_with_construction() {
        let m = MountSpec::new("/h", "t", "/w").read_only();
        assert_eq!(m.guest_path, "/w");
        assert!(m.read_only);
    }

    #[test]
    fn mount_spec_constructor_accepts_string_types() {
        // Builder uses `impl Into<String>` for all three positional
        // args — &str and String should both work.
        let m_str = MountSpec::new("/h", "t", "/abs");
        let m_string = MountSpec::new(String::from("/h"), String::from("t"), String::from("/abs"));
        assert_eq!(m_str.guest_path, m_string.guest_path);
        assert_eq!(m_str.host_path, m_string.host_path);
        assert_eq!(m_str.guest_tag, m_string.guest_tag);
    }

    #[test]
    fn volume_spec_basic_construction() {
        let v = VolumeSpec::new("/host/nm.img", "/workspace/node_modules");
        assert_eq!(v.host_path, "/host/nm.img");
        assert_eq!(v.guest_path, "/workspace/node_modules");
        assert_eq!(v.size_bytes, VolumeSpec::DEFAULT_SIZE_BYTES);
    }

    #[test]
    fn volume_spec_with_size_bytes() {
        let v = VolumeSpec::new("/h", "/g").with_size_bytes(4 * 1024 * 1024 * 1024);
        assert_eq!(v.size_bytes, 4 * 1024 * 1024 * 1024);
    }
}