//
// Syd: rock-solid unikernel
// src/syd-load.rs: Load a dynamic library like syd for testing
//
// Copyright (c) 2024 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0-or-later

use std::process::ExitCode;

use getargs::{Opt, Options};
use libseccomp::{scmp_cmp, ScmpAction, ScmpFilterContext, ScmpSyscall};
use nix::{
    errno::Errno,
    sys::wait::{waitpid, WaitStatus},
    unistd::{fork, ForkResult},
};
use syd::err::SydResult;

// Set global allocator to mimalloc.
#[global_allocator]
static GLOBAL: mimalloc2_rust::GlobalMiMalloc = mimalloc2_rust::GlobalMiMalloc;

fn main() -> SydResult<ExitCode> {
    syd::set_sigpipe_dfl()?;

    // Configure mimalloc.
    syd::configure_mimalloc();

    let args = std::env::args().skip(1).collect::<Vec<_>>();
    let args: &'static Vec<String> = Box::leak(Box::new(args));
    let mut opts = Options::new(args.iter().map(String::as_str));
    let mut mdwe = false;
    let mut scmp = false;

    while let Some(opt) = opts.next_opt()? {
        match opt {
            Opt::Short('h') => {
                help();
                return Ok(ExitCode::SUCCESS);
            }
            Opt::Short('m') => {
                mdwe = true;
            }
            Opt::Short('s') => {
                scmp = true;
            }
            Opt::Short(opt) => {
                eprintln!("Unknown option: -{opt}!");
                help();
                return Ok(ExitCode::FAILURE);
            }
            Opt::Long(opt) => {
                eprintln!("Unknown option: --{opt}!");
                help();
                return Ok(ExitCode::FAILURE);
            }
        }
    }

    let lib = if let Some(lib) = opts.positionals().next() {
        lib
    } else {
        help();
        return Ok(ExitCode::FAILURE);
    };

    if !lib.ends_with(".so") {
        eprintln!("syd-load: invalid library extension, expected \".so\"!");
        return Ok(ExitCode::FAILURE);
    }

    if mdwe {
        const PR_SET_MDWE: nix::libc::c_int = 65;
        const PR_MDWE_REFUSE_EXEC_GAIN: nix::libc::c_ulong = 1;

        // SAFETY: In libc, we trust.
        if unsafe { nix::libc::prctl(PR_SET_MDWE, PR_MDWE_REFUSE_EXEC_GAIN, 0, 0, 0) } != 0 {
            eprintln!("Failed to set mdwe-refuse-exec-gain: {}!", Errno::last());
            return Ok(ExitCode::FAILURE);
        }
    }

    #[allow(clippy::disallowed_methods)]
    if scmp {
        let mut ctx = ScmpFilterContext::new_filter(ScmpAction::Allow).expect("scmp_new_filter");
        // We don't want ECANCELED, we want actual errnos.
        let _ = ctx.set_api_sysrawrc(true);
        // We kill for bad system call and bad arch.
        let _ = ctx.set_act_badarch(ScmpAction::KillProcess);
        // Use a binary tree sorted by syscall number.
        let _ = ctx.set_ctl_optimize(2);

        syd::seccomp_add_architectures(&mut ctx).expect("scmp_add_arch");

        // Restriction -1: Prevent mmap(NULL, MAP_FIXED).
        const MAP_FIXED: u64 = nix::libc::MAP_FIXED as u64;
        const MAP_FIXED_NOREPLACE: u64 = nix::libc::MAP_FIXED_NOREPLACE as u64;
        for sysname in ["mmap", "mmap2"] {
            let syscall = ScmpSyscall::from_name(sysname).unwrap();
            ctx.add_rule_conditional(
                ScmpAction::Errno(nix::libc::EACCES),
                syscall,
                &[
                    scmp_cmp!($arg0 == 0),
                    scmp_cmp!($arg3 & MAP_FIXED == MAP_FIXED),
                ],
            )
            .expect("scmp_MAP_FIXED");
            ctx.add_rule_conditional(
                ScmpAction::Errno(nix::libc::EACCES),
                syscall,
                &[
                    scmp_cmp!($arg0 == 0),
                    scmp_cmp!($arg3 & MAP_FIXED_NOREPLACE == MAP_FIXED_NOREPLACE),
                ],
            )
            .expect("scmp_MAP_FIXED_NOREPLACE");
        }

        // Restriction 0: Prohibit attempts to create memory mappings
        // that are writable and executable at the same time, or to
        // change existing memory mappings to become executable, or
        // mapping shared memory segments as executable.
        const R: u64 = nix::libc::PROT_READ as u64;
        const X: u64 = nix::libc::PROT_EXEC as u64;
        const W: u64 = nix::libc::PROT_WRITE as u64;
        #[cfg(target_arch = "aarch64")]
        const B: u64 = 10u64; // PROT_BTI
        #[cfg(not(target_arch = "aarch64"))]
        const B: u64 = 0u64;
        const RX: u64 = R | X;
        const WX: u64 = W | X;
        const BRX: u64 = B | RX;
        const SHM_X: u64 = nix::libc::SHM_EXEC as u64;
        const MAP_S: u64 = nix::libc::MAP_SHARED as u64;
        for sysname in ["mmap", "mmap2"] {
            // Prevent writable and executable memory.
            let syscall = ScmpSyscall::from_name(sysname).unwrap();
            ctx.add_rule_conditional(
                ScmpAction::Errno(nix::libc::EACCES),
                syscall,
                &[scmp_cmp!($arg2 & WX == WX)],
            )
            .expect("scmp_WX");

            // Prevent executable shared memory.
            ctx.add_rule_conditional(
                ScmpAction::Errno(nix::libc::EACCES),
                syscall,
                &[scmp_cmp!($arg2 & X == X), scmp_cmp!($arg3 & MAP_S == MAP_S)],
            )
            .expect("scmp_X_SHARED");
        }

        for sysname in ["mprotect", "pkey_mprotect"] {
            let syscall = ScmpSyscall::from_name(sysname).unwrap();
            // SAFETY: We allow PROT_READ|PROT_EXEC.
            // SAFETY: We allow PROT_BTI|PROT_EXEC on aarch64.
            ctx.add_rule_conditional(
                ScmpAction::Errno(nix::libc::EACCES),
                syscall,
                &[scmp_cmp!($arg2 & BRX == X)],
            )
            .expect("scmp_mprotect_X");
        }

        ctx.add_rule_conditional(
            ScmpAction::Errno(nix::libc::EACCES),
            ScmpSyscall::from_name("shmat").unwrap(),
            &[scmp_cmp!($arg2 & SHM_X == SHM_X)],
        )
        .expect("scmp_SHM_X");

        ctx.load().expect("scmp");
    }

    // SAFETY: Yes of course dlopening a library is unsafe.
    // Use this tool for testing purposes.
    let lib = match unsafe { libloading::os::unix::Library::new(lib) } {
        Ok(lib) => lib,
        Err(err) => {
            eprintln!("syd-load: {err:?}");
            return Ok(ExitCode::FAILURE);
        }
    };

    // To mimick, syd's behaviour we fork here and
    // execute the syd_main function in the child.
    // SAFETY: Fork & FFI is unsafe, use with care.
    match unsafe { fork() } {
        Ok(ForkResult::Parent { child, .. }) => Ok(match waitpid(child, None) {
            Ok(WaitStatus::Exited(_, code)) => ExitCode::from(code as u8),
            Ok(WaitStatus::Signaled(_, signal, _)) => ExitCode::from(128 + signal as u8),
            Ok(status) => unreachable!("BUG: invalid waitstatus: {status:?}"),
            Err(errno) => ExitCode::from(errno as i32 as u8),
        }),
        Ok(ForkResult::Child) => {
            // SAFETY: See above.
            let fun: libloading::os::unix::Symbol<unsafe extern "C" fn() -> i32> =
                match unsafe { lib.get(b"syd_main") } {
                    Ok(fun) => fun,
                    Err(err) => {
                        eprintln!("syd-load: {err}");
                        return Ok(ExitCode::FAILURE);
                    }
                };

            // SAFETY: See above.
            Ok(ExitCode::from(unsafe { fun() } as u8))
        }
        Err(errno) => {
            eprintln!("syd-load: fork failed: {errno}!");
            Ok(ExitCode::FAILURE)
        }
    }
}

fn help() {
    println!("Usage: syd-load [-hms] {{library.so}}");
    println!("Given a dynamic library, loads it and runs the function \"syd_main\" in it.");
    println!("The library is loaded in the parent and executed in the child like syd does.");
    println!("Use -m to enable MDWE protections using prctl(2) PR_SET_MDWE.");
    println!("Use -s to enable MDWE protections using seccomp(2).");
}
