//
// Syd: rock-solid application kernel
// src/syd-load.rs: Run a command under Memory-Deny-Write-Execute protections
//
// Copyright (c) 2024, 2025 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0

use std::{
    env,
    ffi::OsString,
    os::unix::process::CommandExt,
    process::{Command, ExitCode},
};

use nix::errno::Errno;
use syd::{
    config::*,
    err::SydResult,
    libseccomp::{ScmpAction, ScmpFilterContext, ScmpSyscall},
    proc::proc_mmap_min_addr,
    scmp_cmp,
};

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

    syd::set_sigpipe_dfl()?;

    // Configure syd::proc.
    syd::config::proc_init()?;

    // Parse CLI options.
    //
    // Note, option parsing is POSIXly correct:
    // POSIX recommends that no more options are parsed after the first
    // positional argument. The other arguments are then all treated as
    // positional arguments.
    // See: https://pubs.opengroup.org/onlinepubs/9699919799/basedefs/V1_chap12.html#tag_12_02
    let mut opt_mdwe = false;
    let mut opt_scmp = false;
    let mut opt_cmd = env::var_os(ENV_SH).unwrap_or(OsString::from(SYD_SH));
    let mut opt_arg = Vec::new();

    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(prog) => {
                opt_cmd = prog;
                opt_arg.extend(parser.raw_args()?);
            }
            _ => return Err(arg.unexpected().into()),
        }
    }

    if !opt_mdwe && !opt_scmp {
        // Default is to enable both.
        opt_mdwe = true;
        opt_scmp = true;
    }

    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);
        }
    }

    if opt_scmp {
        let mut ctx = ScmpFilterContext::new(ScmpAction::Allow)?;
        // 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)?;

        // Restriction -1: Prevent mmap(addr<${mmap_min_addr}, MAP_FIXED).
        const MAP_FIXED: u64 = nix::libc::MAP_FIXED as u64;
        const MAP_FIXED_NOREPLACE: u64 = nix::libc::MAP_FIXED_NOREPLACE as u64;
        let mmap_min_addr = proc_mmap_min_addr().unwrap_or(4096);
        for sysname in ["mmap", "mmap2"] {
            #[allow(clippy::disallowed_methods)]
            let syscall = ScmpSyscall::from_name(sysname).unwrap();
            ctx.add_rule_conditional(
                ScmpAction::KillProcess,
                syscall,
                &[
                    scmp_cmp!($arg0 < mmap_min_addr),
                    scmp_cmp!($arg3 & MAP_FIXED == MAP_FIXED),
                ],
            )?;
            ctx.add_rule_conditional(
                ScmpAction::KillProcess,
                syscall,
                &[
                    scmp_cmp!($arg0 < mmap_min_addr),
                    scmp_cmp!($arg3 & MAP_FIXED_NOREPLACE == 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 W: u64 = nix::libc::PROT_WRITE as u64;
        const X: u64 = nix::libc::PROT_EXEC as u64;
        const WX: u64 = W | X;
        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.
            #[allow(clippy::disallowed_methods)]
            let syscall = ScmpSyscall::from_name(sysname).unwrap();
            ctx.add_rule_conditional(
                ScmpAction::KillProcess,
                syscall,
                &[scmp_cmp!($arg2 & WX == WX)],
            )?;

            // Prevent executable shared memory.
            ctx.add_rule_conditional(
                ScmpAction::KillProcess,
                syscall,
                &[scmp_cmp!($arg2 & X == X), scmp_cmp!($arg3 & MAP_S == MAP_S)],
            )?;
        }

        for sysname in ["mprotect", "pkey_mprotect"] {
            #[allow(clippy::disallowed_methods)]
            let syscall = ScmpSyscall::from_name(sysname).unwrap();
            ctx.add_rule_conditional(
                ScmpAction::KillProcess,
                syscall,
                &[scmp_cmp!($arg2 & X == X)],
            )?;
        }

        #[allow(clippy::disallowed_methods)]
        ctx.add_rule_conditional(
            ScmpAction::KillProcess,
            ScmpSyscall::from_name("shmat").unwrap(),
            &[scmp_cmp!($arg2 & SHM_X == SHM_X)],
        )?;

        ctx.load()?;
    }

    Ok(ExitCode::from(
        127 + Command::new(opt_cmd)
            .args(opt_arg)
            .exec()
            .raw_os_error()
            .unwrap_or(0) as u8,
    ))
}

fn help() {
    println!("Usage: syd-mdwe [-hms] {{command [args..]}}");
    println!("Run a command under Memory-Deny-Write-Execute protections.");
    println!("Use -m to enable protections using prctl(2) PR_SET_MDWE (default).");
    println!("Use -s to enable protections using seccomp(2) (use with -m to enable both).");
}
