//! forktest -- checks fork, execve, wait4, vfork and fd inheritance.
//!
//!   forktest          run all the checks (exit code = failed checks)
//!   forktest bomb     fork until the kernel says no (EAGAIN), then clean
//!                     up: the process table is bounded, nothing leaks
//!   forktest args ... (used by the exec check) verify argv, envp and the
//!                     auxiliary vector the kernel built for us
#![no_std]
#![no_main]

extern crate alloc;

#[macro_use]
#[path = "ulib.rs"]
mod ulib;

use alloc::vec::Vec;
use ulib::*;

const SELF: &str = "/mnt/disk0/bin/forktest";

fn main() -> i64 {
    match arg(1) {
        Some("bomb") => bomb(),
        Some("args") => check_args(),
        _ => all(),
    }
}

fn all() -> i64 {
    let mut c = Checks::new("forktest");
    let me = getpid();

    // 1. The basic dance: fork returns twice.
    let pid = fork();
    if pid == 0 {
        // The child: a new pid, and its parent is the process that forked.
        exit(if getppid() == me && getpid() != me { 7 } else { 1 });
    }
    c.check(pid > 0, "fork returns the child's pid in the parent");
    let r = wait4(pid, 0);
    c.check(matches!(r, Ok((p, s)) if p == pid && wifexited(s) && wexitstatus(s) == 7), "child saw getppid() == parent; wait4 gives its exit code 7");

    // 2. Many children, reaped in any order with wait4(-1).
    let mut sum = 0;
    for i in 1..=10 {
        if fork() == 0 {
            exit(i);
        }
    }
    let mut reaped = 0;
    while let Ok((_, status)) = wait4(-1, 0) {
        sum += wexitstatus(status);
        reaped += 1;
    }
    c.check(reaped == 10 && sum == 55, "10 children, wait4(-1) reaps them all (exit codes add up to 55)");
    c.check(wait4(-1, 0) == Err(-ECHILD), "no children left: ECHILD");

    // 3. execve: the child becomes another program (here: forktest itself
    // in "args" mode, which checks what it was given).
    let pid = fork();
    if pid == 0 {
        let err = execve(SELF, &["forktest", "args", "two words", ""], &["FOO=bar", "EMPTY="]);
        exit(100 + (-err).min(50)); // execve only returns on failure
    }
    let code = wait(pid);
    c.check(code == 0, "execve: argv, envp and auxv arrive the Linux way (see the [forktest args] lines)");
    let pid = fork();
    if pid == 0 {
        let err = execve("/mnt/disk0/bin/no-such-program", &["x"], &[]);
        exit(if err == -ENOENT { 0 } else { 1 });
    }
    c.check(wait(pid) == 0, "execve of a missing file fails with ENOENT (and the caller keeps running)");

    // 4. vfork: the parent waits until the child has exec'd.
    let pid = vfork();
    if pid == 0 {
        execve("/mnt/disk0/bin/child", &["child", "5"], &[]);
        exit(99);
    }
    c.check(wait(pid) == 5, "vfork + execve child 5 -> exit code 5");

    // 5. fds are inherited: the child writes into our pipe via dup2.
    let (r, w) = pipe2(O_CLOEXEC).unwrap_or((-1, -1));
    c.check(fcntl(r, F_GETFD, 0) == FD_CLOEXEC as i64, "pipe2(O_CLOEXEC) sets FD_CLOEXEC");
    let pid = fork();
    if pid == 0 {
        dup2(w, 1); // stdout = the pipe (dup2 clears close-on-exec on fd 1)
        execve("/mnt/disk0/bin/child", &["child", "pipe"], &[]);
        exit(99);
    }
    close(w); // else the read below would never see end-of-file
    let mut out = Vec::new();
    let mut buf = [0u8; 128];
    loop {
        let n = read(r, &mut buf);
        if n <= 0 {
            break;
        }
        out.extend_from_slice(&buf[..n as usize]);
    }
    close(r);
    c.check(wait(pid) == 0 && out.starts_with(b"hello through a pipe"), "fork + dup2 + execve: the child's stdout comes through our pipe");

    // 6. A killed child: the status says which signal.
    let pid = fork();
    if pid == 0 {
        unsafe { core::ptr::read_volatile(0x10 as *const u8) }; // NULL-ish pointer
        exit(0);
    }
    let r = wait4(pid, 0);
    c.check(matches!(r, Ok((_, s)) if wifsignaled(s) && wtermsig(s) == SIGSEGV), "a child that touches address 0x10 dies of SIGSEGV");
    c.done()
}

/// The exec'd side of check 3.
fn check_args() -> i64 {
    let mut c = Checks::new("forktest args");
    c.check(argc() == 4, "argc == 4");
    c.check(arg(0) == Some("forktest") && arg(2) == Some("two words") && arg(3) == Some(""), "argv[0..4] as passed (incl. an empty one)");
    c.check(env(0) == Some("FOO=bar") && env(1) == Some("EMPTY=") && env(2).is_none(), "envp exactly as passed");
    c.check(getenv("FOO") == Some("bar"), "getenv(\"FOO\") == \"bar\"");
    c.check(getauxval(AT_PAGESZ) == Some(4096), "AT_PAGESZ = 4096");
    let entry = getauxval(AT_ENTRY).unwrap_or(0);
    c.check(entry != 0 && entry == ulib_start(), "AT_ENTRY = the address of _start");
    // AT_PHDR points at our own ELF program headers: the first loadable
    // one (type 1 = PT_LOAD) should contain the entry point.
    let (phdr, phent, phnum) = (getauxval(AT_PHDR).unwrap_or(0), getauxval(AT_PHENT).unwrap_or(0), getauxval(AT_PHNUM).unwrap_or(0));
    let mut found = false;
    if phdr != 0 && phent == 56 {
        for i in 0..phnum {
            let h = (phdr + i * phent) as *const u64;
            let (ty, vaddr, memsz) = unsafe { (*(h as *const u32), *h.add(2), *h.add(5)) };
            found |= ty == 1 && (vaddr..vaddr + memsz).contains(&entry);
        }
    }
    c.check(found, "AT_PHDR/AT_PHENT/AT_PHNUM: our program headers, a PT_LOAD contains the entry");
    let random = getauxval(AT_RANDOM).map(|p| unsafe { *(p as *const [u8; 16]) });
    c.check(random.is_some_and(|r| r != [0; 16]), "AT_RANDOM: 16 random bytes");
    let execfn = getauxval(AT_EXECFN).map(|p| unsafe { cstr(p as *const u8) });
    c.check(execfn == Some(SELF), "AT_EXECFN = the path we were run as");
    let sp = &c as *const _ as u64;
    c.check(sp > 0xff_0000_0000 && sp < 0x100_0000_0000, "the stack is at the top of user space");
    c.done()
}

fn ulib_start() -> u64 {
    unsafe extern "C" {
        fn _start();
    }
    _start as usize as u64
}

/// Fork until fork fails. Every child just waits (pause) for a signal.
/// Then kill them all and reap them; the kernel tests check that no memory
/// leaked and that the system still works afterwards.
fn bomb() -> i64 {
    let mut pids = Vec::new();
    let error = loop {
        let pid = fork();
        if pid == 0 {
            loop {
                pause(); // until SIGKILL
            }
        }
        if pid < 0 {
            break pid;
        }
        pids.push(pid);
        if pids.len() > 1000 {
            break 0; // unbounded? stop anyway
        }
    };
    println!("[forktest bomb] {} children, then fork failed with {} ({})", pids.len(), error, strerror(error));
    for &pid in &pids {
        kill(pid, SIGKILL);
    }
    let mut reaped = 0;
    while let Ok((_, status)) = wait4(-1, 0) {
        if wifsignaled(status) && wtermsig(status) == SIGKILL {
            reaped += 1;
        }
    }
    println!("[forktest bomb] killed and reaped {}", reaped);
    if error == -EAGAIN && reaped == pids.len() && !pids.is_empty() { 0 } else { 1 }
}