//
// Syd: rock-solid application kernel
// benches/sys/exec.rs: exec microbenchmarks
//
// Copyright (c) 2024 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0

// This benchmark tests various ways of calling `exec` on Linux:
//   1) execve("/dev/null", [], [])
//   2) execve("/bin/true", [], [])
//   3) execve("/bin/true", NULL, NULL)  // "ROP exec" style
//   4) execveat(fd-to-bin-true, [], [], AT_EMPTY_PATH)
//   5) execveat(memfd-to-bin-true, [], [], AT_EMPTY_PATH)
//
// For #5, we open a memfd in the initialization phase, copy `/bin/true` into
// it, then run `execveat` from that in the benchmark.

use std::{
    ffi::CString,
    fs::File,
    io::{Read as _, Seek},
    os::fd::{AsRawFd, FromRawFd, RawFd},
    ptr,
};

use brunch::{benches, Bench};
use libc::{_exit, c_char, c_int, execve, fork, waitpid, SYS_execve, SYS_execveat, AT_EMPTY_PATH};
use nix::{
    errno::Errno,
    fcntl::{open, OFlag},
    sys::{
        memfd::{memfd_create, MemFdCreateFlag},
        stat::Mode,
    },
    unistd::{close, dup, lseek, write as nix_write},
};

/// Fork a child, and in the child, call the provided function (which performs
/// `execve` or `execveat`). The parent waits for the child to exit.
fn do_fork_exec(exec_fn: impl FnOnce() + Send + 'static) {
    unsafe {
        let pid = fork();
        if pid == 0 {
            // Child
            exec_fn();
            // If we get here, exec failed. Exit with error.
            _exit(127);
        } else if pid < 0 {
            panic!("fork() failed: {:?}", Errno::last());
        } else {
            // Parent: wait for child
            let mut status: c_int = 0;
            let w = waitpid(pid, &mut status, 0);
            if w < 0 {
                panic!("waitpid() failed: {:?}", Errno::last());
            }
            // We won't deeply check the exit code, but normally 0 if success, 127 if failure.
        }
    }
}

/// Benchmark 0: execve("/dev/null", [], [])
fn bench_execve_dev_null() {
    do_fork_exec(|| unsafe {
        // We'll call execve with empty argv/env.
        // /dev/null is not a valid ELF, likely it fails with ENOEXEC or EACCES.
        let path_c = CString::new("/dev/null").unwrap();
        let argv: [*const c_char; 1] = [ptr::null()];
        let envp: [*const c_char; 1] = [ptr::null()];
        execve(path_c.as_ptr(), argv.as_ptr(), envp.as_ptr());
        // If we get here, it failed. We'll just _exit(127) above.
    });
}

/// Benchmark 1: execve("/bin/true", [], [])
fn bench_execve_bin_true() {
    do_fork_exec(|| unsafe {
        let path_c = CString::new("/bin/true").unwrap();
        let argv: [*const c_char; 1] = [ptr::null()];
        let envp: [*const c_char; 1] = [ptr::null()];
        execve(path_c.as_ptr(), argv.as_ptr(), envp.as_ptr());
    });
}

/// Benchmark 2: execve("/bin/true", NULL, NULL) -- "ROP exec" style
/// This passes literal NULL for `argv` and `envp`.
fn bench_execve_bin_true_null_null() {
    do_fork_exec(|| unsafe {
        let path_c = CString::new("/bin/true").unwrap();
        // We pass actual NULL pointers for argv/envp.
        libc::syscall(
            SYS_execve,
            path_c.as_ptr(),
            ptr::null::<*const c_char>(),
            ptr::null::<*const c_char>(),
        );
    });
}

/// Benchmark 3: execveat(fd-of-/bin/true, [], [], AT_EMPTY_PATH)
fn bench_execveat_bin_true_fd(fd: RawFd) {
    do_fork_exec(move || unsafe {
        let c_empty = CString::new("").unwrap();
        let argv: [*const c_char; 1] = [ptr::null()];
        let envp: [*const c_char; 1] = [ptr::null()];
        // execveat(fd, "", [], [], AT_EMPTY_PATH)
        libc::syscall(
            SYS_execveat,
            fd,
            c_empty.as_ptr(),
            argv.as_ptr(),
            envp.as_ptr(),
            AT_EMPTY_PATH,
        );
    });
}

/// Benchmark 4: execveat(memfd-of-/bin/true, [], [], AT_EMPTY_PATH)
fn bench_execveat_memfd_bin_true_fd(memfd_fd: RawFd) {
    do_fork_exec(move || unsafe {
        let c_empty = CString::new("").unwrap();
        let argv: [*const c_char; 1] = [ptr::null()];
        let envp: [*const c_char; 1] = [ptr::null()];
        libc::syscall(
            SYS_execveat,
            memfd_fd,
            c_empty.as_ptr(),
            argv.as_ptr(),
            envp.as_ptr(),
            AT_EMPTY_PATH,
        );
    });
}

/// Copy `/bin/true` into a memfd, returning its fd.
fn prepare_memfd_with_bin_true() -> RawFd {
    // Step 1: open /bin/true in normal mode.
    let bintrue_fd =
        open("/bin/true", OFlag::O_RDONLY, Mode::empty()).expect("Failed to open /bin/true");
    let bintrue_file = unsafe { File::from_raw_fd(bintrue_fd) };

    // Step 2: create memfd
    let name = CString::new("memfd_bin_true").unwrap();
    let memfd = memfd_create(&name, MemFdCreateFlag::MFD_CLOEXEC).expect("Failed to create memfd");
    let memfd_fd = memfd.as_raw_fd();

    // Step 3: copy /bin/true into memfd
    let mut buf = [0u8; 4096];
    let mut total_file = bintrue_file;
    total_file
        .seek(std::io::SeekFrom::Start(0))
        .expect("seek /bin/true failed");
    loop {
        let n = total_file.read(&mut buf).unwrap();
        if n == 0 {
            break;
        }
        let written = nix_write(&memfd, &buf[..n]).expect("write to memfd failed");
        if written < n {
            panic!("short write to memfd?! wrote {}", written);
        }
    }

    // Rewind memfd
    let _ = lseek(memfd_fd, 0, nix::unistd::Whence::SeekSet);

    // We won't close bintrue_fd because we used from_raw_fd. We'll let it drop.
    // The memfd we keep open; we just return its FD.
    // But we must not drop `memfd` or it closes. We'll dup it:
    let new_fd = dup(memfd_fd).expect("dup memfd failed");
    // Once we dup, the OwnedFd can drop safely.
    new_fd
}

fn main() {
    // We'll open /bin/true as well for the execveat fd scenario. That file must remain open.
    // We also create a memfd with /bin/true inside it.
    let bintrue_fd = open("/bin/true", OFlag::O_RDONLY, Mode::empty())
        .expect("Failed to open /bin/true for execveat");
    let memfd_fd = prepare_memfd_with_bin_true();

    benches!(
        inline:

        // 0) execve("/dev/null", [], [])
        Bench::new("Execve(/dev/null, [], [])").run(|| {
            bench_execve_dev_null();
        }),

        // 1) execve("/bin/true", [], [])
        Bench::new("Execve(/bin/true, [], [])").run(|| {
            bench_execve_bin_true();
        }),

        // 2) execve("/bin/true", NULL, NULL)
        Bench::new("Execve(/bin/true, NULL, NULL)").run(|| {
            bench_execve_bin_true_null_null();
        }),

        // 3) execveat(fd-to-bin-true, [], [], AT_EMPTY_PATH)
        Bench::new("Execveat(fd-of-/bin/true, [], [], AT_EMPTY_PATH)").run(|| {
            bench_execveat_bin_true_fd(bintrue_fd);
        }),

        // 4) execveat(memfd-to-bin-true, [], [], AT_EMPTY_PATH)
        Bench::new("Execveat(memfd-of-/bin/true, [], [], AT_EMPTY_PATH)").run(|| {
            bench_execveat_memfd_bin_true_fd(memfd_fd);
        }),
    );

    // The FDs remain open through the entire run. We'll close at the end.
    let _ = close(bintrue_fd);
    let _ = close(memfd_fd);
}
