supermachine 0.4.4

//! Spike 22 command-line harness.
//!
//! The VM implementation lives in the library crate. This binary parses CLI
//! flags, applies process-wide harness hooks, and calls `vmm::runner`.

use supermachine as vmm;

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn arg_value(args: &[String], i: usize, flag: &str) -> String {
    args.get(i + 1).cloned().unwrap_or_else(|| {
        eprintln!("{flag}: missing value");
        std::process::exit(2);
    })
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn parse_arg<T>(args: &[String], i: usize, flag: &str) -> T
where
    T: std::str::FromStr,
    T::Err: std::fmt::Display,
{
    let value = arg_value(args, i, flag);
    value.parse().unwrap_or_else(|e| {
        eprintln!("{flag}: invalid value {value:?}: {e}");
        std::process::exit(2);
    })
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn main() {
    env_logger::init();

    use vmm::vmm::resources::{VmProfile, VmResources, DEFAULT_CMDLINE, DEFAULT_MEMORY_MIB};
    use vmm::vmm::runner::{self, RunOptions};

    // P-core hint for vCPU 0 (the main thread enters dispatch_vcpu).
    // Secondary vCPU threads call this themselves in run_secondary.
    #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
    vmm::vmm::worker::pin_vcpu_thread_to_pcore();

    let args: Vec<String> = std::env::args().collect();
    let mut kernel_path: Option<String> = None;
    let mut initrd_path: Option<String> = None;
    let mut cmdline = String::from(DEFAULT_CMDLINE);
    let mut memory_mib: usize = DEFAULT_MEMORY_MIB;
    let mut blk_paths: Vec<String> = Vec::new();
    // `--volume HOST_FILE:GUEST_PATH` — writable attachments,
    // opened RW and ordered after RO layers in /dev/vd*. The
    // matching guest paths are written to /.supermachine-volumes
    // by the bake pipeline; init-oci consumes them post-pivot.
    let mut volumes: Vec<vmm::vmm::resources::VolumeSpec> = Vec::new();
    let mut snapshot_after_ms: Option<u64> = None;
    let mut snapshot_at: Option<u64> = None;
    let mut snapshot_on_listener: bool = false;
    let mut snapshot_out: Option<String> = None;
    let mut restore_from: Option<String> = None;
    let mut cow_restore: bool = false;
    let mut quiesce_ms: u64 = 0;
    let mut log_sink: Option<String> = None;
    let mut n_vcpus: u32 = 1;
    let mut profile: Option<VmProfile> = None;
    let mut vcpus_explicit = false;
    let mut tls_listen: Option<String> = None;
    let mut tls_vm_port: Option<u32> = None;
    let mut tls_cert: Option<String> = None;
    let mut tls_key: Option<String> = None;
    let mut env_pairs: Vec<(String, String)> = Vec::new();
    let mut env_file: Option<String> = None;
    let mut egress_policy: Option<String> = None;
    let mut vsock_mux: Option<String> = None;
    let mut vsock_mux_handoff: Option<String> = None;
    let mut vsock_exec: Option<String> = None;
    let mut vsock_exec_guest_port: Option<u32> = None;
    let mut http_port: Option<String> = None;
    let mut pool_worker: Option<String> = None;

    let mut i = 1;
    while i < args.len() {
        match args[i].as_str() {
            "--kernel" => {
                kernel_path = Some(arg_value(&args, i, "--kernel"));
                i += 2;
            }
            "--initramfs" => {
                initrd_path = Some(arg_value(&args, i, "--initramfs"));
                i += 2;
            }
            "--cmdline" => {
                cmdline = arg_value(&args, i, "--cmdline");
                i += 2;
            }
            "--memory" => {
                memory_mib = parse_arg(&args, i, "--memory");
                i += 2;
            }
            "--virtio-blk" => {
                blk_paths.push(arg_value(&args, i, "--virtio-blk"));
                i += 2;
            }
            "--volume" => {
                let raw = arg_value(&args, i, "--volume");
                let (host, guest) = raw.split_once(':').unwrap_or_else(|| {
                    eprintln!("--volume expects HOST:GUEST, got {raw:?}");
                    std::process::exit(2);
                });
                volumes.push(vmm::vmm::resources::VolumeSpec::new(host, guest));
                i += 2;
            }
            "--vcpus" => {
                n_vcpus = parse_arg(&args, i, "--vcpus");
                vcpus_explicit = true;
                i += 2;
            }
            "--profile" => {
                let value = arg_value(&args, i, "--profile");
                profile = VmProfile::parse(&value).or_else(|| {
                    eprintln!("--profile: expected latency or throughput, got {value:?}");
                    std::process::exit(2);
                });
                i += 2;
            }
            "--snapshot-after-ms" => {
                snapshot_after_ms = Some(parse_arg(&args, i, "--snapshot-after-ms"));
                i += 2;
            }
            "--snapshot-at" => {
                snapshot_at = Some(parse_arg(&args, i, "--snapshot-at"));
                i += 2;
            }
            "--snapshot-on-listener" => {
                snapshot_on_listener = true;
                i += 1;
            }
            "--snapshot-out" => {
                snapshot_out = Some(arg_value(&args, i, "--snapshot-out"));
                i += 2;
            }
            "--restore-from" => {
                restore_from = Some(arg_value(&args, i, "--restore-from"));
                i += 2;
            }
            "--cow-restore" => {
                cow_restore = true;
                i += 1;
            }
            "--quiesce-ms" => {
                quiesce_ms = parse_arg(&args, i, "--quiesce-ms");
                i += 2;
            }
            "--log-sink" => {
                log_sink = Some(arg_value(&args, i, "--log-sink"));
                i += 2;
            }
            "--tls-listen" => {
                tls_listen = Some(arg_value(&args, i, "--tls-listen"));
                i += 2;
            }
            "--tls-vm-port" => {
                tls_vm_port = Some(parse_arg(&args, i, "--tls-vm-port"));
                i += 2;
            }
            "--tls-cert" => {
                tls_cert = Some(arg_value(&args, i, "--tls-cert"));
                i += 2;
            }
            "--tls-key" => {
                tls_key = Some(arg_value(&args, i, "--tls-key"));
                i += 2;
            }
            "--env" => {
                let value = arg_value(&args, i, "--env");
                if let Some((k, v)) = value.split_once('=') {
                    env_pairs.push((k.to_string(), v.to_string()));
                } else {
                    eprintln!("--env: expected K=V, got {value:?}");
                    std::process::exit(2);
                }
                i += 2;
            }
            "--env-file" => {
                env_file = Some(arg_value(&args, i, "--env-file"));
                i += 2;
            }
            "--egress-policy" => {
                egress_policy = Some(arg_value(&args, i, "--egress-policy"));
                i += 2;
            }
            "--vsock-mux" => {
                vsock_mux = Some(arg_value(&args, i, "--vsock-mux"));
                i += 2;
            }
            "--vsock-mux-handoff" => {
                vsock_mux_handoff = Some(arg_value(&args, i, "--vsock-mux-handoff"));
                i += 2;
            }
            "--vsock-exec" => {
                vsock_exec = Some(arg_value(&args, i, "--vsock-exec"));
                i += 2;
            }
            "--vsock-exec-guest-port" => {
                let v = arg_value(&args, i, "--vsock-exec-guest-port");
                vsock_exec_guest_port = Some(v.parse().unwrap_or_else(|e| {
                    eprintln!("--vsock-exec-guest-port: {e}");
                    std::process::exit(2);
                }));
                i += 2;
            }
            "--http-port" => {
                http_port = Some(arg_value(&args, i, "--http-port"));
                i += 2;
            }
            "--pool-worker" => {
                pool_worker = Some(arg_value(&args, i, "--pool-worker"));
                i += 2;
            }
            _ => {
                eprintln!("unknown arg: {}", args[i]);
                std::process::exit(2);
            }
        }
    }

    let tls_cfg = match (tls_listen, tls_cert, tls_key) {
        (Some(l), Some(c), Some(k)) => Some(vmm::vmm::tls::TlsConfig {
            listen_addr: l,
            vm_port: tls_vm_port,
            cert_path: c,
            key_path: k,
        }),
        (None, None, None) => None,
        _ => {
            eprintln!("--tls-listen / --tls-cert / --tls-key must all be set together (--tls-vm-port optional)");
            std::process::exit(2);
        }
    };

    if let Some(p) = log_sink.as_deref() {
        vmm::devices::serial::set_log_sink(p).unwrap_or_else(|e| {
            eprintln!("--log-sink: {e}");
            std::process::exit(2);
        });
    }
    // Enable line-marker detection (heartbeat counter + "parking
    // PID 1") for any snapshot mode. Listener-only bakes use the
    // parked-marker as the early non-service fallback so they
    // don't sit through the full --snapshot-after-ms timeout.
    vmm::devices::serial::set_heartbeat_detection(
        snapshot_at.is_some() || snapshot_on_listener || snapshot_after_ms.is_some(),
    );

    if let Some(json) = build_env_payload(&env_pairs, env_file.as_deref()) {
        eprintln!(
            "  env JSON: {} bytes (served on AF_VSOCK port 1026)",
            json.len()
        );
        vmm::devices::virtio::vsock::muxer::set_env_json(json);
    }

    if let Some(p) = egress_policy.as_deref() {
        eprintln!("  egress policy: {p}");
        vmm::vmm::egress_policy::set(p);
    }

    let mut pool_restore_path: Option<String> = None;
    let mut pool_sock: Option<std::os::unix::net::UnixStream> = None;
    if let Some(sock_path) = pool_worker.as_deref() {
        use std::io::{BufRead, Write};
        let mut sock = std::os::unix::net::UnixStream::connect(sock_path).unwrap_or_else(|e| {
            eprintln!("--pool-worker connect {sock_path}: {e}");
            std::process::exit(1);
        });
        sock.write_all(b"READY\n").ok();
        eprintln!("  pool-worker connected to {sock_path}");
        let mut reader = std::io::BufReader::new(sock.try_clone().unwrap_or_else(|e| {
            eprintln!("--pool-worker clone {sock_path}: {e}");
            std::process::exit(1);
        }));
        let mut line = String::new();
        if reader.read_line(&mut line).is_err() {
            eprintln!("  pool-worker: supervisor closed before RESTORE");
            std::process::exit(0);
        }
        let cmd = line.trim();
        let Some(rest) = cmd.strip_prefix("RESTORE ") else {
            eprintln!("  pool-worker: expected RESTORE, got {cmd:?}");
            std::process::exit(1);
        };
        let mut parts = rest.split_ascii_whitespace();
        let base = parts.next().unwrap_or("").to_string();
        for kv in parts {
            if let Some(v) = kv.strip_prefix("egress_policy=") {
                vmm::vmm::egress_policy::set(v);
            }
        }
        pool_restore_path = Some(base);
        pool_sock = Some(sock);
        cow_restore = true;
    }
    let restore_from = pool_restore_path.or(restore_from);

    if kernel_path.is_none() && restore_from.is_none() {
        runner::run_proof_of_life().unwrap_or_else(|e| {
            eprintln!("HVF proof-of-life failed: {e}");
            std::process::exit(1);
        });
        return;
    }

    let mut resources = match (kernel_path, initrd_path) {
        (Some(kernel), Some(initramfs)) => VmResources::for_kernel(kernel, initramfs),
        (Some(kernel), None) => VmResources::new().with_kernel_path(kernel),
        (None, Some(initramfs)) => VmResources::new().with_initramfs(initramfs),
        (None, None) => VmResources::new(),
    }
    .with_cmdline(cmdline)
    .with_memory_mib(memory_mib)
    .with_vcpus(n_vcpus)
    .with_cow_restore(cow_restore)
    .with_quiesce_ms(quiesce_ms);
    if let Some(path) = restore_from {
        resources = resources.with_restore(path);
    }
    for path in blk_paths {
        resources = resources.with_block_device(path);
    }
    for volume in volumes {
        resources = resources.with_volume(volume);
    }
    if let Some(after_ms) = snapshot_after_ms {
        if let Some(out_path) = snapshot_out.as_deref() {
            resources = resources.with_snapshot_after_ms(after_ms, out_path);
        } else {
            resources.snapshot.after_ms = Some(after_ms);
        }
    }
    if let Some(at_heartbeat) = snapshot_at {
        if let Some(out_path) = snapshot_out.as_deref() {
            resources = resources.with_snapshot_at_heartbeat(at_heartbeat, out_path);
        } else {
            resources.snapshot.at_heartbeat = Some(at_heartbeat);
        }
    }
    if snapshot_on_listener {
        if let Some(out_path) = snapshot_out.as_deref() {
            resources = resources.with_snapshot_on_listener(out_path);
        } else {
            resources.snapshot.on_listener = true;
        }
    }
    if resources.snapshot.out_path.is_none() {
        resources.snapshot.out_path = snapshot_out;
    }
    if let Some(path) = vsock_mux {
        resources = resources.with_vsock_mux(path);
    }
    if let Some(path) = vsock_mux_handoff {
        resources = resources.with_vsock_mux_handoff(path);
    }
    if let Some(path) = vsock_exec {
        resources = resources.with_vsock_exec(path);
    }
    if let Some(port) = vsock_exec_guest_port {
        resources = resources.with_vsock_exec_guest_port(port);
    }
    if let Some(port) = http_port {
        resources = resources.with_http_port(port);
    }
    if let Some(profile) = profile {
        if !vcpus_explicit {
            resources.apply_profile_defaults(profile);
        }
    }

    runner::run(
        &resources,
        RunOptions {
            tls: tls_cfg,
            pool_sock,
            pool_worker: None,
            experimental_skip_warm_gic_restore: std::env::var_os("SUPERMACHINE_SKIP_WARM_GIC_RESTORE")
                .is_some(),
        },
    )
    .unwrap_or_else(|e| {
        eprintln!("VM run failed: {e}");
        std::process::exit(2);
    });
}

fn build_env_payload(pairs: &[(String, String)], file: Option<&str>) -> Option<String> {
    if pairs.is_empty() && file.is_none() {
        return None;
    }
    if let Some(path) = file {
        return std::fs::read_to_string(path).ok();
    }
    let mut out = String::from(r#"{"env":{"#);
    for (i, (k, v)) in pairs.iter().enumerate() {
        if i > 0 {
            out.push(',');
        }
        out.push('"');
        json_escape_into(&mut out, k);
        out.push_str(r#"":"#);
        out.push('"');
        json_escape_into(&mut out, v);
        out.push('"');
    }
    out.push_str(r#"},"secrets":{}}"#);
    Some(out)
}

fn json_escape_into(out: &mut String, s: &str) {
    for c in s.chars() {
        match c {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if (c as u32) < 0x20 => {
                use std::fmt::Write;
                let _ = write!(out, "\\u{:04x}", c as u32);
            }
            c => out.push(c),
        }
    }
}

#[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
fn main() {
    eprintln!("supermachine only runs on macOS aarch64");
}