//! 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.

use serde::Deserialize;
use std::ffi::CString;
use std::io::{Read, Write};
use std::os::fd::FromRawFd;
use std::sync::{Arc, Mutex};

#[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,
}

/// 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>,
}

const PATH: &str = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin";

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(())
}

/// The file `program` names: itself when it holds a `/`, otherwise the first
/// executable regular file of that name in `path`'s directories.
fn resolve_program(program: &str, path: &str) -> Option<String> {
    use std::os::unix::fs::PermissionsExt;
    if program.is_empty() {
        return None;
    }
    if program.contains('/') {
        return Some(program.to_string());
    }
    path.split(':')
        .filter(|d| !d.is_empty())
        .map(|d| format!("{d}/{program}"))
        .find(|p| std::fs::metadata(p).is_ok_and(|m| m.is_file() && m.permissions().mode() & 0o111 != 0))
}

/// A job's environment: `PATH` and `HOME=/tmp` by default, then the
/// image's own (`/etc/flow/env`, written by `flow rootfs import --env`),
/// then the job's (its sandbox's and its own, already merged by the host).
/// Later ones win, so a pipeline can always set `PATH`.
fn job_env(image: Option<&str>, job: &[(String, String)]) -> std::collections::BTreeMap<String, String> {
    let mut env = std::collections::BTreeMap::new();
    env.insert("PATH".to_string(), PATH.to_string());
    env.insert("HOME".to_string(), "/tmp".to_string());
    for line in image.unwrap_or("").lines() {
        if let Some((k, v)) = line.split_once('=') {
            if !k.is_empty() {
                env.insert(k.to_string(), v.to_string());
            }
        }
    }
    for (k, v) in job {
        env.insert(k.clone(), v.clone());
    }
    env
}

/// The links every Linux system has under /dev and devtmpfs does not
/// make: `/dev/fd` (bash's process substitution, `<(…)`) and the standard
/// streams, all through the job's own /proc.
fn dev_links(dev: &str) -> Result<(), String> {
    for (name, target) in DEV_LINKS {
        let link = format!("{dev}/{name}");
        if std::fs::symlink_metadata(&link).is_err() {
            std::os::unix::fs::symlink(target, &link).map_err(|e| format!("link {link}: {e}"))?;
        }
    }
    Ok(())
}

const DEV_LINKS: [(&str, &str); 4] =
    [("fd", "/proc/self/fd"), ("stdin", "/proc/self/fd/0"), ("stdout", "/proc/self/fd/1"), ("stderr", "/proc/self/fd/2")];

/// Set `lo` up; the kernel gives it 127.0.0.1 and ::1 itself.
fn loopback_up() -> Result<(), String> {
    let fd = unsafe { libc::socket(libc::AF_INET, libc::SOCK_DGRAM | libc::SOCK_CLOEXEC, 0) };
    if fd < 0 {
        return Err(format!("socket for lo: {}", std::io::Error::last_os_error()));
    }
    let mut req: libc::ifreq = unsafe { std::mem::zeroed() };
    for (i, b) in b"lo".iter().enumerate() {
        req.ifr_name[i] = *b as libc::c_char;
    }
    let ok = unsafe {
        libc::ioctl(fd, libc::SIOCGIFFLAGS, &mut req) == 0 && {
            req.ifr_ifru.ifru_flags |= libc::IFF_UP as libc::c_short;
            libc::ioctl(fd, libc::SIOCSIFFLAGS, &req) == 0
        }
    };
    let err = std::io::Error::last_os_error();
    unsafe { libc::close(fd) };
    if !ok {
        return Err(format!("bring up lo: {err}"));
    }
    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)) => {
                mount(dev, &target, "ext4", nosuid, "")?;
                chown(&target, uid, gid)?;
            }
            (kind, _) => return Err(format!("unknown mount kind {kind}")),
        }
    }
    if job.loopback {
        loopback_up()?;
    }
    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()));
    }

    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();
    }
    Ok(if libc::WIFEXITED(status) { libc::WEXITSTATUS(status) } else { 128 + libc::WTERMSIG(status) })
}

/// 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);
    }
}

#[cfg(test)]
mod tests {
    use super::{dev_links, job_env, resolve_program};

    #[test]
    fn the_image_environment_sits_under_the_jobs() {
        let image = "PATH=/opt/llvm-mos/bin:/usr/bin:/bin\nX16EMU=/opt/x16emu\nnot a line\n";
        let env = job_env(Some(image), &[("X16EMU".into(), "/job".into()), ("LANG".into(), "C".into())]);
        assert_eq!(env["PATH"], "/opt/llvm-mos/bin:/usr/bin:/bin", "the image's PATH replaces the default");
        assert_eq!(env["X16EMU"], "/job", "the job's value wins");
        assert_eq!(env["HOME"], "/tmp");
        assert_eq!(env["LANG"], "C");
        assert_eq!(job_env(None, &[])["PATH"], super::PATH);
    }

    #[test]
    fn dev_gets_the_usual_links_once() {
        let dir = std::env::temp_dir().join(format!("flow-init-dev-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let d = dir.to_str().unwrap();
        dev_links(d).unwrap();
        dev_links(d).unwrap();
        assert_eq!(std::fs::read_link(dir.join("fd")).unwrap(), std::path::Path::new("/proc/self/fd"));
        assert_eq!(std::fs::read_link(dir.join("stdout")).unwrap(), std::path::Path::new("/proc/self/fd/1"));
        // Through the host's /proc, the link works as in a job.
        assert!(dir.join("fd").join("0").exists() || !std::path::Path::new("/proc/self/fd/0").exists());
        std::fs::remove_dir_all(&dir).unwrap();
    }

    use std::os::unix::fs::PermissionsExt;

    #[test]
    fn programs_are_found_in_the_jobs_path() {
        let dir = std::env::temp_dir().join(format!("flow-init-test-{}", std::process::id()));
        let (a, b) = (dir.join("a"), dir.join("b"));
        std::fs::create_dir_all(&a).unwrap();
        std::fs::create_dir_all(&b).unwrap();
        std::fs::write(a.join("tool"), "not executable").unwrap();
        std::fs::write(b.join("tool"), "#!/bin/sh\n").unwrap();
        std::fs::set_permissions(b.join("tool"), std::fs::Permissions::from_mode(0o755)).unwrap();
        std::fs::create_dir_all(a.join("dir")).unwrap();
        let path = format!("{}::{}", a.display(), b.display());
        assert_eq!(resolve_program("tool", &path), Some(format!("{}/tool", b.display())), "skips a file that is not executable");
        assert_eq!(resolve_program("dir", &path), None, "and directories");
        assert_eq!(resolve_program("missing", &path), None);
        assert_eq!(resolve_program("./x", &path).as_deref(), Some("./x"), "a path is used as given");
        assert_eq!(resolve_program("", &path), None);
        std::fs::remove_dir_all(&dir).unwrap();
    }
}