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

use std::process::ExitCode;

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

fn main() -> SydResult<ExitCode> {
    use lexopt::prelude::*;

    syd::set_sigpipe_dfl()?;

    // Parse CLI options.
    let mut opt_mdwe = false;
    let mut opt_scmp = false;
    let mut opt_libp = None;

    let mut parser = lexopt::Parser::from_env();
    while let Some(arg) = parser.next()? {
        match arg {
            Short('h') => {
                help();
                return Ok(ExitCode::SUCCESS);
            }
            Short('m') => opt_mdwe = true,
            Short('s') => opt_scmp = true,
            Value(lib) => opt_libp = Some(XPathBuf::from(lib)),
            _ => return Err(arg.unexpected().into()),
        }
    }

    let lib = if let Some(lib) = opt_libp {
        lib
    } else {
        help();
        return Ok(ExitCode::FAILURE);
    };

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

    if opt_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 opt_scmp {
        let mut ctx = ScmpFilterContext::new(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).");
}
