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

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

use getargs::{Opt, Options};
use libseccomp::{scmp_cmp, ScmpAction, ScmpFilterContext, ScmpSyscall};
use nix::errno::Errno;
use syd::{
    err::{SydError, SydResult},
    proc::proc_mmap_min_addr,
};

// 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 positionals: Vec<&str> = opts.positionals().collect();
    if positionals.is_empty() {
        help();
        return Ok(ExitCode::FAILURE);
    }

    if !mdwe && !scmp {
        // Default is to enable both.
        mdwe = true;
        scmp = true;
    }

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

    if scmp {
        let mut ctx = ScmpFilterContext::new_filter(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()?;
    }

    let (command, args) = positionals
        .split_first()
        .ok_or::<SydError>(Errno::EINVAL.into())?;
    Ok(ExitCode::from(
        127 + Command::new(command)
            .args(args)
            .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).");
}
