//
// Syd: rock-solid application kernel
// src/utils/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::{
    sys::wait::{waitpid, WaitStatus},
    unistd::{fork, ForkResult},
};
use syd::{
    confine::{confine_mdwe, confine_scmp_wx},
    path::XPathBuf,
};

// Set global allocator to mimalloc.
#[cfg(all(not(feature = "prof"), target_pointer_width = "64"))]
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;

// Set global allocator to tcmalloc if profiling is enabled.
#[cfg(feature = "prof")]
#[global_allocator]
static GLOBAL: tcmalloc::TCMalloc = tcmalloc::TCMalloc;

syd::main! {
    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 {
        if let Err(errno) = confine_mdwe(false) {
            eprintln!("prctl failed to set Memory-Deny-Write-Execute: {errno}!");
            return Err(errno.into());
        }
    }

    if opt_scmp {
        if let Err(error) = confine_scmp_wx() {
            eprintln!("seccomp failed to set W^X restrictions: {error}!");
            return Err(error);
        }
    }

    // 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 mimic, 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).");
}
