//! PID 1 of a job's microVM (docs/m0.md §2). It mounts the read-only

//! rootfs and the workspace, runs the job's argv as an unprivileged user

//! with exactly the job's environment, frames the output on the serial

//! console for the host, and powers off. It never decides anything the

//! host has not: the job comes from `/job.json` in the initramfs the host

//! built, and the host enforces the time limit.

//!

//! Frames, one per line on the console: `\x1eO ` stdout, `\x1eE ` stderr,

//! `\x1eX <code>` the exit code (128 + signal when killed), `\x1eF <text>`

//! a failure of init itself, `\x1eC <json>` a build cache pruned before

//! the guest powers off (oxygenflow#72): `{"target": ..., "pruned":

//! flow_init::prune::Pruned}`, sent before the exit code.

use flow_init::prune;

use serde::Deserialize;

use std::ffi::CString;

use std::io::{Read, Write};

use std::os::fd::FromRawFd;

use std::sync::{Arc, Mutex};

mod egress;

mod job;

use job::{dev_links, job_env, loopback_up, resolve_program, PATH};

#[derive(Deserialize)]

struct JobSpec {

    argv: Vec<String>,

    env: Vec<(String, String)>,

    #[serde(default)]

    mounts: Vec<MountSpec>,

    /// `network "loopback"`: bring up `lo` (the guest has no other device).

    #[serde(default)]

    loopback: bool,

    /// `network "egress"` (oxygenflow#55): `lo` comes up and a proxy on

    /// 127.0.0.1:3128 relays over vsock to the runner.

    #[serde(default)]

    egress: bool,

    /// Secrets to write before the job starts (spec §9): each at `path`

    /// under `/run/secrets`, 0400 and the job's user's, on a tmpfs.

    #[serde(default)]

    secrets: Vec<SecretSpec>,

}

#[derive(serde::Deserialize)]

struct SecretSpec {

    path: String,

    value: String,

}

/// An ephemeral mount from the host: a tmpfs of `size_mib`, or the disk at

/// `device`, mounted at `target` under the new root.

#[derive(serde::Deserialize)]

struct MountSpec {

    kind: String,

    target: String,

    size_mib: u32,

    #[serde(default)]

    device: Option<String>,

    /// A build cache's disk (oxygenflow#72): mounted `strictatime`, so

    /// what the job reads is seen, and after a passing job pruned to this

    /// percentage of its size before the guest powers off.

    #[serde(default)]

    prune: Option<u8>,

    /// Paths on the disk, relative to it, that pruning leaves: the job's

    /// outputs, which the runner reads after the guest stops.

    #[serde(default)]

    keep: Vec<String>,

}

fn main() {

    let console = Arc::new(Mutex::new(std::io::stdout()));

    let code = match run(&console) {

        Ok(code) => code,

        Err(e) => {

            frame(&console, b'F', e.as_bytes());

            125

        }

    };

    frame(&console, b'X', code.to_string().as_bytes());

    unsafe {

        libc::sync();

        // Firecracker ends the VM when the guest reboots (reboot=k).

        libc::reboot(libc::LINUX_REBOOT_CMD_RESTART);

    }

}

fn frame(console: &Mutex<std::io::Stdout>, kind: u8, text: &[u8]) {

    let mut out = console.lock().unwrap_or_else(|p| p.into_inner());

    let mut line = Vec::with_capacity(text.len() + 4);

    line.extend_from_slice(&[0x1e, kind, b' ']);

    line.extend_from_slice(text);

    line.push(b'\n');

    let _ = out.write_all(&line);

    let _ = out.flush();

}

fn cstr(s: &str) -> Result<CString, String> {

    CString::new(s).map_err(|_| format!("{s:?} contains NUL"))

}

fn chown(path: &str, uid: libc::uid_t, gid: libc::gid_t) -> Result<(), String> {

    let p = cstr(path)?;

    if unsafe { libc::chown(p.as_ptr(), uid, gid) } != 0 {

        return Err(format!("chown {path}: {}", std::io::Error::last_os_error()));

    }

    Ok(())

}

/// Create `path` and its missing parents, each owned by the job's user.

fn make_dirs_owned(path: &str, uid: libc::uid_t, gid: libc::gid_t) -> Result<(), String> {

    let mut missing = Vec::new();

    let mut p = std::path::Path::new(path);

    while !p.exists() {

        missing.push(p.to_path_buf());

        p = match p.parent() {

            Some(parent) => parent,

            None => break,

        };

    }

    for d in missing.iter().rev() {

        std::fs::create_dir(d).map_err(|e| format!("mkdir {}: {e}", d.display()))?;

        chown(&d.to_string_lossy(), uid, gid)?;

    }

    Ok(())

}

fn mount(src: &str, target: &str, fstype: &str, flags: libc::c_ulong, data: &str) -> Result<(), String> {

    let (s, t, f, d) = (cstr(src)?, cstr(target)?, cstr(fstype)?, cstr(data)?);

    let rc = unsafe { libc::mount(s.as_ptr(), t.as_ptr(), f.as_ptr(), flags, d.as_ptr() as *const libc::c_void) };

    if rc != 0 {

        return Err(format!("mount {src} on {target}: {}", std::io::Error::last_os_error()));

    }

    Ok(())

}

fn run(console: &Arc<Mutex<std::io::Stdout>>) -> Result<i32, String> {

    let job: JobSpec = serde_json::from_slice(&std::fs::read("/job.json").map_err(|e| format!("read /job.json: {e}"))?)

        .map_err(|e| format!("parse /job.json: {e}"))?;

    if job.argv.is_empty() {

        return Err("empty argv".into());

    }

    let nosuid = libc::MS_NOSUID | libc::MS_NODEV;

    std::fs::create_dir_all("/dev").ok();

    mount("devtmpfs", "/dev", "devtmpfs", libc::MS_NOSUID, "")?;

    std::fs::create_dir_all("/newroot").map_err(|e| e.to_string())?;

    mount("/dev/vda", "/newroot", "ext4", libc::MS_RDONLY | libc::MS_NODEV, "")?;

    mount("/dev/vdb", "/newroot/src", "ext4", nosuid, "")?;

    mount("tmpfs", "/newroot/tmp", "tmpfs", nosuid, "mode=1777")?;

    mount("proc", "/newroot/proc", "proc", nosuid | libc::MS_NOEXEC, "")?;

    mount("devtmpfs", "/newroot/dev", "devtmpfs", libc::MS_NOSUID, "")?;

    dev_links("/newroot/dev")?;

    if std::path::Path::new("/newroot/run").is_dir() {

        mount("tmpfs", "/newroot/run", "tmpfs", nosuid, "mode=755")?;

    }

    // The job runs as whoever owns the workspace (the host's builder), never root.

    let (uid, gid) = {

        let mut st: libc::stat = unsafe { std::mem::zeroed() };

        let p = cstr("/newroot/src")?;

        if unsafe { libc::stat(p.as_ptr(), &mut st) } != 0 {

            return Err("stat /src".into());

        }

        if st.st_uid == 0 {

            (1000, 1000)

        } else {

            (st.st_uid, st.st_gid)

        }

    };

    // Ephemeral mounts: mount points are made (and owned by the job's user)

    // under the writable /src or /tmp, then each is mounted and given to it.

    for m in &job.mounts {

        let target = format!("/newroot{}", m.target);

        make_dirs_owned(&target, uid, gid)?;

        match (m.kind.as_str(), &m.device) {

            ("tmpfs", _) => mount("tmpfs", &target, "tmpfs", nosuid, &format!("size={}m,uid={uid},gid={gid},mode=755", m.size_mib))?,

            ("disk", Some(dev)) => {

                let atime = if m.prune.is_some() { libc::MS_STRICTATIME | libc::MS_LAZYTIME } else { 0 };

                mount(dev, &target, "ext4", nosuid | atime, "")?;

                chown(&target, uid, gid)?;

            }

            (kind, _) => return Err(format!("unknown mount kind {kind}")),

        }

    }

    // Secrets: an in-memory mount of their own, only the job's user's,

    // each file 0400; the job file they came from is outside the new root.

    if !job.secrets.is_empty() {

        if !std::path::Path::new("/newroot/run").is_dir() {

            return Err("the image has no /run, where secrets are mounted".into());

        }

        let dir = "/newroot/run/secrets";

        std::fs::create_dir_all(dir).map_err(|e| format!("mkdir {dir}: {e}"))?;

        mount("tmpfs", dir, "tmpfs", nosuid | libc::MS_NOEXEC, &format!("size=1m,uid={uid},gid={gid},mode=700"))?;

        for s in &job.secrets {

            let Some(rel) = s.path.strip_prefix("/run/secrets/") else {

                return Err(format!("secret path {:?} is not under /run/secrets/", s.path));

            };

            let full = format!("{dir}/{rel}");

            if let Some(parent) = std::path::Path::new(&full).parent().and_then(|p| p.to_str()) {

                make_dirs_owned(parent, uid, gid)?;

            }

            std::fs::write(&full, s.value.as_bytes()).map_err(|e| format!("write {}: {e}", s.path))?;

            chown(&full, uid, gid)?;

            std::fs::set_permissions(&full, std::os::unix::fs::PermissionsExt::from_mode(0o400))

                .map_err(|e| format!("chmod {}: {e}", s.path))?;

        }

    }

    if job.loopback || job.egress {

        loopback_up()?;

    }

    if job.egress {

        egress::start()?;

    }

    let root = cstr("/newroot")?;

    if unsafe { libc::chroot(root.as_ptr()) } != 0 || std::env::set_current_dir("/src").is_err() {

        return Err(format!("enter the rootfs: {}", std::io::Error::last_os_error()));

    }

    // What the job reads or writes on a cache disk from now on is what

    // this run used.

    let started = job.mounts.iter().any(|m| m.prune.is_some()).then(prune::start_clock);

    let argv: Vec<CString> = job.argv.iter().map(|a| cstr(a)).collect::<Result<_, _>>()?;

    // glibc's execvpe searches the caller's PATH, and PID 1 has none of its

    // own, so the program is found here, in the job's PATH, before forking.

    let vars = job_env(std::fs::read_to_string("/etc/flow/env").ok().as_deref(), &job.env);

    let path = vars.get("PATH").map(String::as_str).unwrap_or(PATH);

    let program = cstr(&resolve_program(&job.argv[0], path).ok_or_else(|| format!("cannot run {:?}: not found in {path}", job.argv[0]))?)?;

    let env: Vec<CString> = vars.iter().map(|(k, v)| cstr(&format!("{k}={v}"))).collect::<Result<_, _>>()?;

    let mut out = [0; 2];

    let mut err = [0; 2];

    if unsafe { libc::pipe2(out.as_mut_ptr(), libc::O_CLOEXEC) } != 0 || unsafe { libc::pipe2(err.as_mut_ptr(), libc::O_CLOEXEC) } != 0 {

        return Err("pipe".into());

    }

    let pid = unsafe { libc::fork() };

    if pid < 0 {

        return Err("fork".into());

    }

    if pid == 0 {

        // Child: no way back to root, then the program itself (execvpe: no shell).

        unsafe {

            let null = libc::open(c"/dev/null".as_ptr(), libc::O_RDONLY);

            libc::dup2(null, 0);

            libc::dup2(out[1], 1);

            libc::dup2(err[1], 2);

            libc::setsid();

            if libc::setgroups(0, std::ptr::null()) != 0 || libc::setgid(gid) != 0 || libc::setuid(uid) != 0 {

                libc::_exit(126);

            }

            let mut a: Vec<*const libc::c_char> = argv.iter().map(|c| c.as_ptr()).collect();

            a.push(std::ptr::null());

            let mut e: Vec<*const libc::c_char> = env.iter().map(|c| c.as_ptr()).collect();

            e.push(std::ptr::null());

            libc::execve(program.as_ptr(), a.as_ptr(), e.as_ptr());

            let msg = format!("flow: cannot run {:?}: {}\n", job.argv[0], std::io::Error::last_os_error());

            libc::write(2, msg.as_ptr() as *const libc::c_void, msg.len());

            libc::_exit(127);

        }

    }

    unsafe {

        libc::close(out[1]);

        libc::close(err[1]);

    }

    let readers: Vec<_> = [(out[0], b'O'), (err[0], b'E')]

        .into_iter()

        .map(|(fd, kind)| {

            let console = console.clone();

            std::thread::spawn(move || pump(unsafe { std::fs::File::from_raw_fd(fd) }, kind, &console))

        })

        .collect();

    let mut status = 0;

    // Reap everything (we are PID 1) until the job itself exits.

    loop {

        let r = unsafe { libc::waitpid(-1, &mut status, 0) };

        if r == pid || r < 0 {

            break;

        }

    }

    for r in readers {

        let _ = r.join();

    }

    let code = if libc::WIFEXITED(status) { libc::WEXITSTATUS(status) } else { 128 + libc::WTERMSIG(status) };

    // A failing job's disks are thrown away: only a passing one's are

    // worth pruning.

    if let (Some(started), 0) = (started, code) {

        prune_caches(console, &job.mounts, started);

    }

    Ok(code)

}

/// Prune each cache disk and say so to the host. Whatever the job left

/// running stops first, so nothing changes the disks meanwhile.

fn prune_caches(console: &Mutex<std::io::Stdout>, mounts: &[MountSpec], started: std::time::SystemTime) {

    unsafe {

        libc::kill(-1, libc::SIGKILL);

        let mut status = 0;

        while libc::waitpid(-1, &mut status, 0) > 0 {}

    }

    for m in mounts {

        let Some(percent) = m.prune else { continue };

        let root = std::path::Path::new(&m.target);

        let pruned = match prune::usage(root) {

            Ok((used, size)) => {

                let mut p = prune::prune(root, started, &m.keep, size, used, percent).unwrap_or_else(|e| prune::Pruned {

                    size,

                    percent,

                    before: used,

                    after: used,

                    unpruned: Some(format!("pruning failed: {e}")),

                    ..Default::default()

                });

                // What is left as the file system counts it.

                if p.freed > 0 {

                    unsafe { libc::sync() };

                    if let Ok((now, _)) = prune::usage(root) {

                        p.after = now;

                    }

                }

                p

            }

            Err(e) => prune::Pruned { percent, unpruned: Some(format!("statvfs {}: {e}", m.target)), ..Default::default() },

        };

        let msg = serde_json::json!({ "target": m.target, "pruned": pruned });

        frame(console, b'C', msg.to_string().as_bytes());

    }

}

/// Frame a stream line by line; a long line is cut at 16 KiB.

fn pump(mut f: std::fs::File, kind: u8, console: &Mutex<std::io::Stdout>) {

    let mut buf = [0u8; 8192];

    let mut line: Vec<u8> = Vec::new();

    loop {

        let n = match f.read(&mut buf) {

            Ok(0) | Err(_) => break,

            Ok(n) => n,

        };

        for &b in &buf[..n] {

            if b == b'\n' || line.len() >= 16 * 1024 {

                frame(console, kind, &line);

                line.clear();

                if b == b'\n' {

                    continue;

                }

            }

            line.push(b);

        }

    }

    if !line.is_empty() {

        frame(console, kind, &line);

    }

}