//! ClawMates microVM guest agent — pid 1 inside a Firecracker microVM.
//!
//! Runs as `init=/usr/local/bin/fcagent`'s exec target and answers the host over
//! **vsock** (port 9001), never the serial console: feeding a guest over stdin
//! races its startup and arrives half-consumed. The console stays a log.
//!
//! # Why this is a static Rust binary and not the python script it replaces
//!
//! The python version worked only because Firecracker's CI Ubuntu image happens
//! to ship python3. **None of our own images do** — `agent-base` has neither
//! python nor git, `agent-terminal` has git but no python — so the agent could
//! never have run in a real mission rootfs. An agent that dictates what must be
//! installed in the image has the dependency backwards. This is a
//! `x86_64-unknown-linux-musl` static binary: it needs nothing from the rootfs
//! it is dropped into.
//!
//! # Wire protocol (unchanged from the python agent, deliberately)
//!
//! One request per connection: a 4-byte big-endian length followed by JSON, and
//! the reply framed the same way. The length prefix is the point — a reply
//! larger than a socket buffer arrives in pieces, and reading "whatever was
//! available" would parse a truncated object as a complete one.
//!
//! Ops: `ping`, `exec`, `put`, `get`. `crates/bins/clawmates-node/src/microvm.rs`
//! and `crates/cm-api/src/microvm_client.rs` speak this and needed no change.

use std::io::{Read, Write};
use std::os::unix::process::CommandExt;
use std::path::Path;
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};

use base64::Engine;
use serde_json::{json, Value};

const PORT: u32 = 9001;
/// Cap on a single request. A hostile or broken host must not be able to make
/// pid 1 allocate without bound and get the VM OOM-killed.
const MAX_REQUEST: u32 = 512 * 1024 * 1024;

const B64: base64::engine::general_purpose::GeneralPurpose =
    base64::engine::general_purpose::STANDARD;

fn main() {
    // The mounts the init script would otherwise do. Done here so the agent
    // works whether it is exec'd from a shell init or used as `init=` directly:
    // /proc missing makes every process-inspecting tool in the guest lie.
    for (fstype, target) in [
        ("proc", "/proc"),
        ("sysfs", "/sys"),
        ("devtmpfs", "/dev"),
        ("tmpfs", "/tmp"),
    ] {
        if !Path::new(target).join(".").exists() {
            let _ = std::fs::create_dir_all(target);
        }
        let _ = Command::new("mount")
            .args(["-t", fstype, fstype, target])
            .status();
    }

    let listener = match vsock::VsockListener::bind_with_cid_port(libc_vmaddr_cid_any(), PORT) {
        Ok(l) => l,
        Err(e) => {
            // Printed to the console, which is where the host's boot check
            // looks. Exiting pid 1 panics the kernel, which is the honest
            // outcome: a VM whose agent cannot listen is unusable, and it must
            // not sit there looking booted.
            eprintln!("FC-AGENT-FATAL could not bind vsock port {PORT}: {e}");
            std::process::exit(1);
        }
    };

    // The host greps the console for this before it tries to connect.
    println!("FC-AGENT-LISTENING port={PORT}");
    let _ = std::io::stdout().flush();

    for conn in listener.incoming() {
        match conn {
            Ok(mut s) => {
                if let Err(e) = serve_one(&mut s) {
                    // A bad request must never kill the agent — the VM would
                    // look booted and answer nothing, the worst of both.
                    eprintln!("FC-AGENT-ERROR {e}");
                }
            }
            Err(e) => eprintln!("FC-AGENT-ERROR accept: {e}"),
        }
    }
}

/// `VMADDR_CID_ANY` — bind for any host CID.
fn libc_vmaddr_cid_any() -> u32 {
    u32::MAX
}

fn serve_one(s: &mut vsock::VsockStream) -> Result<(), String> {
    let mut len = [0u8; 4];
    s.read_exact(&mut len)
        .map_err(|e| format!("read length: {e}"))?;
    let len = u32::from_be_bytes(len);
    if len > MAX_REQUEST {
        // Answer rather than hang up: a caller that sent something absurd needs
        // to be told, not left waiting for a reply that will never come.
        return reply(s, &json!({ "ok": false, "error": format!("request of {len} bytes exceeds the {MAX_REQUEST} cap") }));
    }
    let mut buf = vec![0u8; len as usize];
    s.read_exact(&mut buf)
        .map_err(|e| format!("read body: {e}"))?;

    let resp = match serde_json::from_slice::<Value>(&buf) {
        Ok(req) => handle(&req),
        Err(e) => json!({ "ok": false, "error": format!("undecodable request: {e}") }),
    };
    reply(s, &resp)
}

fn reply(s: &mut vsock::VsockStream, v: &Value) -> Result<(), String> {
    let body = serde_json::to_vec(v).map_err(|e| format!("encode reply: {e}"))?;
    s.write_all(&(body.len() as u32).to_be_bytes())
        .map_err(|e| format!("write length: {e}"))?;
    s.write_all(&body)
        .map_err(|e| format!("write body: {e}"))?;
    s.flush().map_err(|e| format!("flush: {e}"))
}

fn handle(req: &Value) -> Value {
    let op = req.get("op").and_then(Value::as_str).unwrap_or_default();
    match op {
        "ping" => json!({ "ok": true, "pid": std::process::id() }),
        "exec" => op_exec(req),
        "put" => op_put(req),
        "get" => op_get(req),
        other => json!({ "ok": false, "error": format!("unknown op: {other}") }),
    }
}

/// Extra environment for the command, on top of the image's own.
///
/// This is how credentials reach the agent CLI. An env var rather than a file
/// because the per-VM rootfs is destroyed with the VM but an env var never
/// touches the guest disk at all — it exists only in the process's environment
/// for the length of one exec.
///
/// **Every problem here fails the exec.** The tempting alternative — skip the
/// entry we could not use and run anyway — produces a `claude -p` with no
/// credential, and that does not error: it hangs. A phase stuck at `running`
/// for ten minutes with nothing in the logs is exactly what a missing token
/// looked like on the container path, so a request we cannot honour in full is
/// refused with a reason instead.
///
/// Errors name the key and never the value: the value is the secret, and an
/// error string travels back over the wire and into logs.
fn env_pairs(req: &Value) -> Result<Vec<(String, String)>, String> {
    let Some(env) = req.get("env") else {
        return Ok(Vec::new());
    };
    // `null` means "nothing to add" — that is what a caller with no credentials
    // serialises. Anything else that is not an object is a caller bug.
    if env.is_null() {
        return Ok(Vec::new());
    }
    let Some(map) = env.as_object() else {
        return Err("exec env must be an object of name → string".into());
    };
    let mut out = Vec::with_capacity(map.len());
    for (k, v) in map {
        let Some(val) = v.as_str() else {
            return Err(format!("exec env {k}: value must be a string"));
        };
        // `putenv` semantics: a name containing '=' would be parsed as part of
        // the value, silently defining a different variable than the one asked
        // for. A NUL truncates at the C boundary, for the same class of reason.
        if k.is_empty() {
            return Err("exec env has an empty variable name".into());
        }
        if k.contains('=') || k.contains('\0') {
            return Err(format!("exec env {k:?}: name may not contain '=' or NUL"));
        }
        if val.contains('\0') {
            return Err(format!("exec env {k}: value may not contain NUL"));
        }
        out.push((k.clone(), val.to_string()));
    }
    Ok(out)
}

fn op_exec(req: &Value) -> Value {
    let cmd = req.get("cmd").and_then(Value::as_str).unwrap_or_default();
    if cmd.is_empty() {
        return json!({ "ok": false, "error": "exec needs a cmd" });
    }
    let cwd = req.get("cwd").and_then(Value::as_str).unwrap_or("/");
    let secs = req.get("timeout").and_then(Value::as_u64).unwrap_or(3600);
    let env = match env_pairs(req) {
        Ok(v) => v,
        Err(e) => return json!({ "ok": false, "error": e }),
    };

    // The image's ENV was written to /etc/profile.d by the rootfs builder;
    // `sh -c` does not read it, so source it here — otherwise a CLI that relies
    // on `ENV PATH` behaves differently in the VM than in the container, which
    // is exactly the drift the builder extracted that file to prevent.
    //
    // The `if [ -f ]` guard is load-bearing. `. missing-file` makes a
    // NON-INTERACTIVE POSIX shell exit immediately with status 1, so the naive
    // `. env.sh 2>/dev/null; cmd` returned rc=1 without running `cmd` at all on
    // any rootfs lacking that file — every exec silently failing while looking
    // like an ordinary non-zero exit. Caught by the exit-7 unit test.
    const ENV_FILE: &str = "/etc/profile.d/00-image-env.sh";
    let sourced = format!("if [ -f {ENV_FILE} ]; then . {ENV_FILE}; fi\n{cmd}");
    let mut c = Command::new("/bin/sh");
    c.arg("-c")
        .arg(&sourced)
        .envs(env)
        .current_dir(if Path::new(cwd).is_dir() { cwd } else { "/" })
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        // A new process group so a command that spawns background children can
        // be killed wholesale. Without it a stray daemon keeps the run alive and
        // the host's timeout is the only thing that ends it.
        .process_group(0);

    let mut child = match c.spawn() {
        Ok(ch) => ch,
        Err(e) => return json!({ "ok": false, "error": format!("spawn: {e}") }),
    };
    let pid = child.id() as i32;

    // std has no wait-with-timeout, so poll. The output pipes are read after
    // the wait, which is safe here because a command producing more than a pipe
    // buffer of output while we are not draining it would deadlock — so the
    // deadline is enforced by killing the group, and the pipes are drained by
    // `wait_with_output` immediately after.
    let deadline = Instant::now() + Duration::from_secs(secs);
    let timed_out = loop {
        match child.try_wait() {
            Ok(Some(_)) => break false,
            Ok(None) => {}
            Err(e) => return json!({ "ok": false, "error": format!("wait: {e}") }),
        }
        if Instant::now() >= deadline {
            kill_group(pid);
            break true;
        }
        std::thread::sleep(Duration::from_millis(20));
    };

    let out = match child.wait_with_output() {
        Ok(o) => o,
        Err(e) => return json!({ "ok": false, "error": format!("collect output: {e}") }),
    };
    if timed_out {
        // Reported as ok:false, not as rc=124: "we stopped it" is a different
        // fact from "it exited non-zero", and the caller must be able to tell.
        return json!({
            "ok": false,
            "error": format!("command exceeded its {secs}s budget and was killed"),
            "stdout": String::from_utf8_lossy(&out.stdout),
            "stderr": String::from_utf8_lossy(&out.stderr),
        });
    }
    json!({
        "ok": true,
        // A signalled process has no exit code; report the conventional
        // 128+signal rather than silently claiming success.
        "rc": exit_code(&out.status),
        "stdout": String::from_utf8_lossy(&out.stdout),
        "stderr": String::from_utf8_lossy(&out.stderr),
    })
}

fn exit_code(status: &std::process::ExitStatus) -> i32 {
    use std::os::unix::process::ExitStatusExt;
    status
        .code()
        .unwrap_or_else(|| 128 + status.signal().unwrap_or(0))
}

fn kill_group(pid: i32) {
    let _ = Command::new("kill")
        .args(["-9", "--", &format!("-{pid}")])
        .status();
}

fn op_put(req: &Value) -> Value {
    let dest = req.get("dest").and_then(Value::as_str).unwrap_or_default();
    if dest.is_empty() {
        return json!({ "ok": false, "error": "put needs a dest" });
    }
    let b64 = req.get("tar_b64").and_then(Value::as_str).unwrap_or_default();
    let raw = match B64.decode(b64) {
        Ok(r) => r,
        Err(e) => return json!({ "ok": false, "error": format!("undecodable archive: {e}") }),
    };
    if let Err(e) = std::fs::create_dir_all(dest) {
        return json!({ "ok": false, "error": format!("mkdir {dest}: {e}") });
    }
    let mut ar = tar::Archive::new(&raw[..]);
    ar.set_overwrite(true);
    // Ownership from the host archive is meaningless in here and re-applying it
    // is how the container path grew a uid split. The guest is root; let it own
    // what it is given.
    ar.set_preserve_permissions(false);
    match ar.unpack(dest) {
        Ok(()) => json!({ "ok": true, "dest": dest, "bytes": raw.len() }),
        Err(e) => json!({ "ok": false, "error": format!("unpack into {dest}: {e}") }),
    }
}

fn op_get(req: &Value) -> Value {
    let path = req.get("path").and_then(Value::as_str).unwrap_or_default();
    if path.is_empty() {
        return json!({ "ok": false, "error": "get needs a path" });
    }
    let p = Path::new(path);
    if !p.exists() {
        // A missing path is an error, NOT an empty archive — an empty tar looks
        // exactly like a run that produced nothing.
        return json!({ "ok": false, "error": format!("no such path: {path}") });
    }
    let name = p
        .file_name()
        .map(|s| s.to_string_lossy().to_string())
        .unwrap_or_else(|| "root".to_string());

    let mut b = tar::Builder::new(Vec::new());
    // Do not follow symlinks: a link pointing outside the collected tree would
    // otherwise be dereferenced and its target smuggled back to the host.
    b.follow_symlinks(false);
    let added = if p.is_dir() {
        b.append_dir_all(&name, p)
    } else {
        b.append_path_with_name(p, &name)
    };
    if let Err(e) = added {
        return json!({ "ok": false, "error": format!("archive {path}: {e}") });
    }
    match b.into_inner() {
        Ok(bytes) => json!({ "ok": true, "tar_b64": B64.encode(&bytes), "bytes": bytes.len() }),
        Err(e) => json!({ "ok": false, "error": format!("finish archive for {path}: {e}") }),
    }
}

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

    /// The credential has to actually reach the command. This is the whole
    /// point of the op, and the failure it prevents is silent: a `claude -p`
    /// with no token hangs rather than erroring.
    #[test]
    fn injected_env_reaches_the_command() {
        let r = op_exec(&json!({
            "op": "exec",
            "cmd": "printf %s \"$CLAUDE_CODE_OAUTH_TOKEN\"",
            "env": { "CLAUDE_CODE_OAUTH_TOKEN": "sk-test-value" },
            "timeout": 30,
        }));
        assert_eq!(r["rc"], json!(0));
        assert_eq!(r["stdout"], json!("sk-test-value"));
    }

    /// And it must survive the profile.d sourcing that runs first — a
    /// credential set on the process and then clobbered by the shell would
    /// look identical to one that never arrived.
    #[test]
    fn injected_env_survives_the_image_env_file() {
        let r = op_exec(&json!({
            "op": "exec",
            "cmd": "printf %s \"$INJECTED_PROBE\"",
            "env": { "INJECTED_PROBE": "still-here" },
            "timeout": 30,
        }));
        assert_eq!(r["stdout"], json!("still-here"));
    }

    /// No env is the ordinary case and must not be an error.
    #[test]
    fn absent_or_null_env_is_not_an_error() {
        for req in [
            json!({ "op": "exec", "cmd": "true", "timeout": 30 }),
            json!({ "op": "exec", "cmd": "true", "env": null, "timeout": 30 }),
            json!({ "op": "exec", "cmd": "true", "env": {}, "timeout": 30 }),
        ] {
            assert_eq!(op_exec(&req)["rc"], json!(0), "{req}");
        }
    }

    /// An env entry we cannot honour fails the whole exec rather than being
    /// dropped. Running without the credential is the outcome this refuses:
    /// it does not error, it hangs, which is far harder to diagnose than a
    /// rejected request.
    #[test]
    fn an_unusable_env_entry_fails_the_exec_instead_of_being_skipped() {
        let cases = [
            json!({ "A=B": "x" }),
            json!({ "": "x" }),
            json!({ "TOKEN": 42 }),
            json!({ "TOKEN": null }),
        ];
        for env in cases {
            let r = op_exec(&json!({
                "op": "exec", "cmd": "true", "env": env.clone(), "timeout": 30,
            }));
            assert_eq!(r["ok"], json!(false), "env {env} should be refused");
            assert!(r["rc"].is_null(), "nothing ran, so there is no rc: {r}");
        }
        // A non-object env is a caller bug, not an empty map.
        let r = op_exec(&json!({ "op": "exec", "cmd": "true", "env": "TOKEN=x" }));
        assert_eq!(r["ok"], json!(false));
    }

    /// An error about a credential must not quote the credential: it travels
    /// back over the wire and into the server's logs.
    #[test]
    fn an_env_error_never_echoes_the_value() {
        let r = op_exec(&json!({
            "op": "exec", "cmd": "true", "timeout": 30,
            "env": { "A=B": "super-secret-token" },
        }));
        let err = r["error"].as_str().unwrap_or_default();
        assert!(!err.contains("super-secret-token"), "leaked the value: {err}");
        assert!(err.contains("A=B"), "should name the key: {err}");
    }

    #[test]
    fn an_unknown_op_is_reported_not_ignored() {
        let r = handle(&json!({ "op": "teleport" }));
        assert_eq!(r["ok"], json!(false));
        assert!(r["error"].as_str().unwrap().contains("teleport"));
    }

    #[test]
    fn ping_answers() {
        assert_eq!(handle(&json!({ "op": "ping" }))["ok"], json!(true));
    }

    /// A missing path must be an error, not an empty archive: an empty tar is
    /// indistinguishable from a run that produced nothing.
    #[test]
    fn getting_a_missing_path_is_an_error() {
        let r = op_get(&json!({ "op": "get", "path": "/definitely/not/here" }));
        assert_eq!(r["ok"], json!(false));
        assert!(r["tar_b64"].is_null(), "no archive may be returned");
    }

    /// A command that ran and failed reports `rc`; one we killed reports
    /// `ok:false`. Collapsing the two would make a timeout look like a build
    /// failure and vice versa.
    #[test]
    fn a_failing_command_reports_rc_and_a_killed_one_does_not() {
        let r = op_exec(&json!({ "op": "exec", "cmd": "exit 7", "timeout": 30 }));
        assert_eq!(r["ok"], json!(true), "it ran, so ok is true");
        assert_eq!(r["rc"], json!(7));

        let r = op_exec(&json!({ "op": "exec", "cmd": "sleep 30", "timeout": 1 }));
        assert_eq!(r["ok"], json!(false), "we killed it, so ok is false");
        assert!(r["rc"].is_null(), "a killed command has no exit code");
        assert!(r["error"].as_str().unwrap().contains("budget"));
    }

    #[test]
    fn exec_needs_a_command() {
        assert_eq!(op_exec(&json!({ "op": "exec" }))["ok"], json!(false));
    }

    /// A tar must round-trip through put and get.
    #[test]
    fn a_tar_round_trips_through_put_and_get() {
        let tmp = std::env::temp_dir().join(format!("fcagent-test-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&tmp);

        let mut b = tar::Builder::new(Vec::new());
        let body = b"ROUND-TRIP-OK\n";
        let mut h = tar::Header::new_gnu();
        h.set_path("marker.txt").unwrap();
        h.set_size(body.len() as u64);
        h.set_mode(0o644);
        h.set_entry_type(tar::EntryType::Regular);
        h.set_cksum();
        b.append(&h, &body[..]).unwrap();
        let archive = b.into_inner().unwrap();

        let r = op_put(&json!({
            "op": "put",
            "dest": tmp.display().to_string(),
            "tar_b64": B64.encode(&archive),
        }));
        assert_eq!(r["ok"], json!(true), "put failed: {r}");
        assert_eq!(
            std::fs::read_to_string(tmp.join("marker.txt")).unwrap(),
            "ROUND-TRIP-OK\n"
        );

        let r = op_get(&json!({ "op": "get", "path": tmp.display().to_string() }));
        assert_eq!(r["ok"], json!(true), "get failed: {r}");
        let bytes = B64.decode(r["tar_b64"].as_str().unwrap()).unwrap();
        let mut ar = tar::Archive::new(&bytes[..]);
        let found = ar
            .entries()
            .unwrap()
            .filter_map(Result::ok)
            .any(|e| e.path().map(|p| p.ends_with("marker.txt")).unwrap_or(false));
        assert!(found, "the collected archive must contain marker.txt");

        let _ = std::fs::remove_dir_all(&tmp);
    }
}