msb_krun 0.1.5

//! Sub-builders for VmBuilder nested configuration.

use std::os::fd::RawFd;
use std::path::{Path, PathBuf};

use vmm::resources::PortConfig;

#[cfg(not(any(feature = "tee", feature = "aws-nitro")))]
use crate::backends::fs::DynFileSystem;

#[cfg(feature = "net")]
use std::os::fd::OwnedFd;

#[cfg(feature = "net")]
use crate::backends::net::NetBackend;

//--------------------------------------------------------------------------------------------------
// Types: Machine Builder
//--------------------------------------------------------------------------------------------------

/// Builder for machine configuration (vCPUs, memory, etc.).
///
/// # Example
///
/// ```rust,no_run
/// # use msb_krun::VmBuilder;
/// VmBuilder::new()
///     .machine(|m| {
///         m.vcpus(4)
///             .memory_mib(2048)
///             .hyperthreading(true)
///             .nested_virt(true)
///     });
/// ```
#[derive(Debug, Clone)]
pub struct MachineBuilder {
    pub(crate) vcpus: u8,
    pub(crate) memory_mib: usize,
    pub(crate) hyperthreading: bool,
    pub(crate) nested_virt: bool,
}

//--------------------------------------------------------------------------------------------------
// Types: Kernel Builder
//--------------------------------------------------------------------------------------------------

/// Builder for kernel configuration.
///
/// # Example
///
/// ```rust,no_run
/// # use msb_krun::VmBuilder;
/// VmBuilder::new()
///     .kernel(|k| {
///         k.krunfw_path("/path/to/libkrunfw.dylib")
///             .cmdline("debug")
///     });
/// ```
#[derive(Debug, Clone, Default)]
pub struct KernelBuilder {
    pub(crate) cmdline: Option<String>,
    pub(crate) krunfw_path: Option<PathBuf>,
    pub(crate) init_path: Option<String>,
}

//--------------------------------------------------------------------------------------------------
// Types: Filesystem Builder
//--------------------------------------------------------------------------------------------------

/// Builder for filesystem configuration.
///
/// # Examples
///
/// Root filesystem only:
///
/// ```rust,no_run
/// # use msb_krun::VmBuilder;
/// VmBuilder::new()
///     .fs(|fs| fs.root("/path/to/rootfs"));
/// ```
///
/// Root filesystem with additional named mounts:
///
/// ```rust,no_run
/// # use msb_krun::VmBuilder;
/// VmBuilder::new()
///     .fs(|fs| fs.root("/path/to/rootfs"))
///     .fs(|fs| fs.tag("data").shm_size(1 << 30).path("/host/data"))
///     .fs(|fs| fs.tag("logs").path("/host/logs"));
/// ```
///
/// Custom filesystem backend:
///
/// ```rust,ignore
/// VmBuilder::new()
///     .fs(|fs| fs.tag("myfs").custom(Box::new(my_backend)));
/// ```
pub struct FsBuilder {
    pub(crate) configs: Vec<FsConfig>,
    current_tag: Option<String>,
    current_shm_size: Option<usize>,
}

/// Configuration for a single filesystem mount.
pub enum FsConfig {
    /// Path-based filesystem (passthrough).
    Path {
        tag: String,
        path: PathBuf,
        shm_size: Option<usize>,
    },
    /// Custom filesystem backend.
    #[cfg(not(any(feature = "tee", feature = "aws-nitro")))]
    Custom {
        tag: String,
        backend: Box<dyn DynFileSystem + Send + Sync>,
    },
}

//--------------------------------------------------------------------------------------------------
// Types: Network Builder
//--------------------------------------------------------------------------------------------------

/// Builder for network configuration.
///
/// # Example
///
/// ```rust,ignore
/// VmBuilder::new()
///     .net(|n| n.mac([0x52, 0x54, 0x00, 0x12, 0x34, 0x56]).custom(my_backend));
/// ```
#[cfg(feature = "net")]
pub struct NetBuilder {
    pub(crate) configs: Vec<NetConfig>,
    current_mac: Option<[u8; 6]>,
}

/// Configuration for a single network device.
#[cfg(feature = "net")]
pub enum NetConfig {
    /// Unixgram backend from a pre-opened fd.
    UnixgramFd { mac: Option<[u8; 6]>, fd: OwnedFd },
    /// Unixgram backend connecting to a socket path.
    UnixgramPath {
        mac: Option<[u8; 6]>,
        path: PathBuf,
        send_vfkit_magic: bool,
    },
    /// Unixstream backend from a pre-opened fd.
    UnixstreamFd { mac: Option<[u8; 6]>, fd: OwnedFd },
    /// Unixstream backend connecting to a socket path.
    UnixstreamPath { mac: Option<[u8; 6]>, path: PathBuf },
    /// TAP backend (Linux only).
    #[cfg(target_os = "linux")]
    Tap { mac: Option<[u8; 6]>, name: String },
    /// Custom network backend.
    Custom {
        mac: Option<[u8; 6]>,
        backend: Box<dyn NetBackend + Send>,
    },
}

//--------------------------------------------------------------------------------------------------
// Types: Console Builder
//--------------------------------------------------------------------------------------------------

/// Builder for console/output configuration.
///
/// # Example
///
/// ```rust,no_run
/// # use msb_krun::VmBuilder;
/// VmBuilder::new()
///     .console(|c| c.output("/tmp/vm.log"));
/// ```
///
/// With the `gpu` and `snd` features:
///
/// ```rust,ignore
/// VmBuilder::new()
///     .console(|c| {
///         c.output("/tmp/vm.log")
///             .sound(true)
///             .gpu_virgl_flags(0x1)
///             .gpu_shm_size(1 << 28)
///     });
/// ```
#[derive(Debug, Clone, Default)]
pub struct ConsoleBuilder {
    pub(crate) output: Option<PathBuf>,
    pub(crate) ports: Vec<PortConfig>,
    pub(crate) disable_implicit: bool,
    #[cfg(feature = "snd")]
    pub(crate) sound: bool,
    #[cfg(feature = "gpu")]
    pub(crate) gpu_virgl_flags: Option<u32>,
    #[cfg(feature = "gpu")]
    pub(crate) gpu_shm_size: Option<usize>,
}

//--------------------------------------------------------------------------------------------------
// Types: Exec Builder
//--------------------------------------------------------------------------------------------------

/// Builder for execution configuration.
///
/// # Examples
///
/// Setting environment variables one at a time with `.env()`:
///
/// ```rust,no_run
/// # use msb_krun::VmBuilder;
/// VmBuilder::new()
///     .exec(|e| {
///         e.path("/bin/myapp")
///             .args(["--flag", "value"])
///             .env("HOME", "/root")
///             .env("LANG", "en_US.UTF-8")
///             .workdir("/app")
///             .rlimit("NOFILE", 1024, 4096)
///     });
/// ```
///
/// Setting environment variables in bulk with `.envs()`:
///
/// ```rust,no_run
/// # use msb_krun::VmBuilder;
/// VmBuilder::new()
///     .exec(|e| {
///         e.path("/bin/myapp")
///             .envs([("HOME", "/root"), ("LANG", "en_US.UTF-8")])
///     });
/// ```
#[derive(Debug, Clone, Default)]
pub struct ExecBuilder {
    pub(crate) path: Option<String>,
    pub(crate) args: Vec<String>,
    pub(crate) env: Vec<(String, String)>,
    pub(crate) workdir: Option<String>,
    pub(crate) rlimits: Vec<(String, u64, u64)>,
}

//--------------------------------------------------------------------------------------------------
// Types: Disk Builder
//--------------------------------------------------------------------------------------------------

/// Supported disk image formats.
#[cfg(feature = "blk")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DiskImageFormat {
    Raw,
    Qcow2,
    Vmdk,
}

/// Builder for block device configuration.
///
/// # Example
///
/// ```rust,no_run
/// # use msb_krun::VmBuilder;
/// VmBuilder::new()
///     .disk(|d| d.path("/path/to/disk.img").read_only(true));
/// ```
#[cfg(feature = "blk")]
#[derive(Debug, Clone)]
pub struct DiskBuilder {
    pub(crate) configs: Vec<DiskConfig>,
    current_path: Option<PathBuf>,
    current_read_only: bool,
    current_format: DiskImageFormat,
}

/// Configuration for a single block device.
#[cfg(feature = "blk")]
#[derive(Debug, Clone)]
pub struct DiskConfig {
    pub path: PathBuf,
    pub read_only: bool,
    pub format: DiskImageFormat,
}

//--------------------------------------------------------------------------------------------------
// Methods: Machine Builder
//--------------------------------------------------------------------------------------------------

impl MachineBuilder {
    /// Create a new machine builder with defaults.
    pub fn new() -> Self {
        Self {
            vcpus: 1,
            memory_mib: 512,
            hyperthreading: false,
            nested_virt: false,
        }
    }

    /// Set the number of virtual CPUs.
    pub fn vcpus(mut self, count: u8) -> Self {
        self.vcpus = count;
        self
    }

    /// Set the memory size in MiB.
    pub fn memory_mib(mut self, mib: usize) -> Self {
        self.memory_mib = mib;
        self
    }

    /// Enable or disable hyperthreading.
    pub fn hyperthreading(mut self, enabled: bool) -> Self {
        self.hyperthreading = enabled;
        self
    }

    /// Enable or disable nested virtualization.
    pub fn nested_virt(mut self, enabled: bool) -> Self {
        self.nested_virt = enabled;
        self
    }
}

impl Default for MachineBuilder {
    fn default() -> Self {
        Self::new()
    }
}

//--------------------------------------------------------------------------------------------------
// Methods: Kernel Builder
//--------------------------------------------------------------------------------------------------

impl KernelBuilder {
    /// Create a new kernel builder.
    pub fn new() -> Self {
        Self::default()
    }

    /// Append to the kernel command line.
    pub fn cmdline(mut self, cmdline: &str) -> Self {
        if let Some(ref mut existing) = self.cmdline {
            existing.push(' ');
            existing.push_str(cmdline);
        } else {
            self.cmdline = Some(cmdline.to_string());
        }
        self
    }

    /// Set an explicit path to the libkrunfw shared library.
    ///
    /// When not set, the OS dynamic linker's default search path is used.
    pub fn krunfw_path(mut self, path: impl AsRef<Path>) -> Self {
        self.krunfw_path = Some(path.as_ref().to_path_buf());
        self
    }

    /// Set the path to the init binary inside the guest.
    ///
    /// This controls the kernel `init=` parameter. When not set, defaults
    /// to `/init.krun`.
    pub fn init_path(mut self, path: impl AsRef<Path>) -> Self {
        self.init_path = Some(path.as_ref().to_string_lossy().to_string());
        self
    }
}

//--------------------------------------------------------------------------------------------------
// Methods: Filesystem Builder
//--------------------------------------------------------------------------------------------------

impl FsBuilder {
    /// Create a new filesystem builder.
    pub fn new() -> Self {
        Self {
            configs: Vec::new(),
            current_tag: None,
            current_shm_size: None,
        }
    }

    /// Set the root filesystem path.
    ///
    /// Uses the virtiofs tag `/dev/root`, matching the kernel's expected root device name.
    pub fn root(mut self, path: impl AsRef<Path>) -> Self {
        self.configs.push(FsConfig::Path {
            tag: "/dev/root".to_string(),
            path: path.as_ref().to_path_buf(),
            shm_size: None,
        });
        self
    }

    /// Set the tag for the next mount.
    pub fn tag(mut self, tag: &str) -> Self {
        self.current_tag = Some(tag.to_string());
        self
    }

    /// Set the path for a filesystem mount.
    ///
    /// If `tag()` was called previously, uses that tag; otherwise generates one.
    pub fn path(mut self, path: impl AsRef<Path>) -> Self {
        let tag = self
            .current_tag
            .take()
            .unwrap_or_else(|| format!("fs{}", self.configs.len()));
        let shm_size = self.current_shm_size.take();

        self.configs.push(FsConfig::Path {
            tag,
            path: path.as_ref().to_path_buf(),
            shm_size,
        });
        self
    }

    /// Set the DAX shared memory size for the next mount.
    pub fn shm_size(mut self, size: usize) -> Self {
        self.current_shm_size = Some(size);
        self
    }

    /// Use a custom filesystem backend.
    #[cfg(not(any(feature = "tee", feature = "aws-nitro")))]
    pub fn custom(mut self, backend: Box<dyn DynFileSystem + Send + Sync>) -> Self {
        let tag = self
            .current_tag
            .take()
            .unwrap_or_else(|| format!("fs{}", self.configs.len()));

        self.configs.push(FsConfig::Custom { tag, backend });
        self
    }
}

impl Default for FsBuilder {
    fn default() -> Self {
        Self::new()
    }
}

//--------------------------------------------------------------------------------------------------
// Methods: Network Builder
//--------------------------------------------------------------------------------------------------

#[cfg(feature = "net")]
impl NetBuilder {
    /// Create a new network builder.
    pub fn new() -> Self {
        Self {
            configs: Vec::new(),
            current_mac: None,
        }
    }

    /// Set the MAC address for the next network device.
    pub fn mac(mut self, mac: [u8; 6]) -> Self {
        self.current_mac = Some(mac);
        self
    }

    /// Attach a unixgram network backend from a pre-opened fd.
    pub fn unixgram(mut self, fd: OwnedFd) -> Self {
        let mac = self.current_mac.take();
        self.configs.push(NetConfig::UnixgramFd { mac, fd });
        self
    }

    /// Attach a unixgram network backend connecting to a socket path.
    pub fn unixgram_path(mut self, path: impl AsRef<Path>, send_vfkit_magic: bool) -> Self {
        let mac = self.current_mac.take();
        self.configs.push(NetConfig::UnixgramPath {
            mac,
            path: path.as_ref().to_path_buf(),
            send_vfkit_magic,
        });
        self
    }

    /// Attach a unixstream network backend from a pre-opened fd.
    pub fn unixstream(mut self, fd: OwnedFd) -> Self {
        let mac = self.current_mac.take();
        self.configs.push(NetConfig::UnixstreamFd { mac, fd });
        self
    }

    /// Attach a unixstream network backend connecting to a socket path.
    pub fn unixstream_path(mut self, path: impl AsRef<Path>) -> Self {
        let mac = self.current_mac.take();
        self.configs.push(NetConfig::UnixstreamPath {
            mac,
            path: path.as_ref().to_path_buf(),
        });
        self
    }

    /// Attach a TAP network backend.
    #[cfg(target_os = "linux")]
    pub fn tap(mut self, name: impl Into<String>) -> Self {
        let mac = self.current_mac.take();
        self.configs.push(NetConfig::Tap {
            mac,
            name: name.into(),
        });
        self
    }

    /// Use a custom network backend.
    pub fn custom(mut self, backend: Box<dyn NetBackend + Send>) -> Self {
        let mac = self.current_mac.take();
        self.configs.push(NetConfig::Custom { mac, backend });
        self
    }
}

#[cfg(feature = "net")]
impl Default for NetBuilder {
    fn default() -> Self {
        Self::new()
    }
}

//--------------------------------------------------------------------------------------------------
// Methods: Console Builder
//--------------------------------------------------------------------------------------------------

impl ConsoleBuilder {
    /// Create a new console builder.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the path to send console output.
    pub fn output(mut self, path: impl AsRef<Path>) -> Self {
        self.output = Some(path.as_ref().to_path_buf());
        self
    }

    /// Enable the virtio-snd device.
    #[cfg(feature = "snd")]
    pub fn sound(mut self, enabled: bool) -> Self {
        self.sound = enabled;
        self
    }

    /// Set GPU virgl flags.
    #[cfg(feature = "gpu")]
    pub fn gpu_virgl_flags(mut self, flags: u32) -> Self {
        self.gpu_virgl_flags = Some(flags);
        self
    }

    /// Set GPU shared memory size.
    #[cfg(feature = "gpu")]
    pub fn gpu_shm_size(mut self, size: usize) -> Self {
        self.gpu_shm_size = Some(size);
        self
    }

    /// Add a bidirectional I/O port to the console device.
    ///
    /// Creates a named port accessible in the guest via `/sys/class/virtio-ports/<name>`.
    /// The host reads from `input_fd` and writes to `output_fd`. Pass the same FD for both
    /// when using a bidirectional socket.
    pub fn port(mut self, name: &str, input_fd: RawFd, output_fd: RawFd) -> Self {
        self.ports.push(PortConfig::InOut {
            name: name.to_string(),
            input_fd,
            output_fd,
        });
        self
    }

    /// Add a TTY port to the console device.
    ///
    /// Creates a named port accessible in the guest via `/sys/class/virtio-ports/<name>`.
    /// The `tty_fd` must be a valid terminal file descriptor. Terminal raw mode is configured
    /// automatically.
    pub fn port_tty(mut self, name: &str, tty_fd: RawFd) -> Self {
        self.ports.push(PortConfig::Tty {
            name: name.to_string(),
            tty_fd,
        });
        self
    }

    /// Disable the implicit console device.
    ///
    /// By default libkrun creates an implicit console that reads from `STDIN_FILENO`.
    /// Call this to suppress that console when using only explicit ports.
    pub fn disable_implicit(mut self) -> Self {
        self.disable_implicit = true;
        self
    }
}

//--------------------------------------------------------------------------------------------------
// Methods: Exec Builder
//--------------------------------------------------------------------------------------------------

impl ExecBuilder {
    /// Create a new exec builder.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the executable path.
    pub fn path(mut self, path: impl AsRef<Path>) -> Self {
        self.path = Some(path.as_ref().to_string_lossy().to_string());
        self
    }

    /// Set command line arguments.
    pub fn args<I, S>(mut self, args: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        self.args = args.into_iter().map(|s| s.as_ref().to_string()).collect();
        self
    }

    /// Add a single environment variable.
    pub fn env(mut self, key: &str, value: &str) -> Self {
        self.env.push((key.to_string(), value.to_string()));
        self
    }

    /// Add multiple environment variables.
    pub fn envs<I, K, V>(mut self, envs: I) -> Self
    where
        I: IntoIterator<Item = (K, V)>,
        K: AsRef<str>,
        V: AsRef<str>,
    {
        self.env.extend(
            envs.into_iter()
                .map(|(k, v)| (k.as_ref().to_string(), v.as_ref().to_string())),
        );
        self
    }

    /// Set the working directory.
    pub fn workdir(mut self, path: impl AsRef<Path>) -> Self {
        self.workdir = Some(path.as_ref().to_string_lossy().to_string());
        self
    }

    /// Set a resource limit.
    pub fn rlimit(mut self, resource: &str, soft: u64, hard: u64) -> Self {
        self.rlimits.push((resource.to_string(), soft, hard));
        self
    }
}

//--------------------------------------------------------------------------------------------------
// Methods: Disk Builder
//--------------------------------------------------------------------------------------------------

#[cfg(feature = "blk")]
impl DiskBuilder {
    /// Create a new disk builder.
    pub fn new() -> Self {
        Self {
            configs: Vec::new(),
            current_path: None,
            current_read_only: false,
            current_format: DiskImageFormat::Raw,
        }
    }

    /// Set the disk image format for the current disk.
    pub fn format(mut self, format: DiskImageFormat) -> Self {
        self.current_format = format;
        self
    }

    /// Set the path for a block device.
    pub fn path(mut self, path: impl AsRef<Path>) -> Self {
        // Finalize any pending config
        if let Some(pending_path) = self.current_path.take() {
            self.configs.push(DiskConfig {
                path: pending_path,
                read_only: self.current_read_only,
                format: self.current_format,
            });
            self.current_read_only = false;
            self.current_format = DiskImageFormat::Raw;
        }

        self.current_path = Some(path.as_ref().to_path_buf());
        self
    }

    /// Set read-only mode for the current disk.
    pub fn read_only(mut self, ro: bool) -> Self {
        self.current_read_only = ro;
        self
    }

    /// Finalize the builder (called internally).
    pub(crate) fn finalize(mut self) -> Self {
        if let Some(path) = self.current_path.take() {
            self.configs.push(DiskConfig {
                path,
                read_only: self.current_read_only,
                format: self.current_format,
            });
        }
        self
    }
}

//--------------------------------------------------------------------------------------------------
// Trait Implementations: Disk Builder
//--------------------------------------------------------------------------------------------------

#[cfg(feature = "blk")]
impl Default for DiskBuilder {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(feature = "blk")]
impl From<DiskImageFormat> for devices::virtio::block::ImageType {
    fn from(format: DiskImageFormat) -> Self {
        match format {
            DiskImageFormat::Raw => devices::virtio::block::ImageType::Raw,
            DiskImageFormat::Qcow2 => devices::virtio::block::ImageType::Qcow2,
            DiskImageFormat::Vmdk => devices::virtio::block::ImageType::Vmdk,
        }
    }
}