//
// Syd: rock-solid unikernel
// src/syd-run.rs: Run a program inside a syd container (requires Linux-5.8 or newer).
//
// Copyright (c) 2024 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0-or-later

use std::{
    os::{
        fd::{FromRawFd, OwnedFd},
        unix::ffi::OsStrExt,
    },
    process::{Command, ExitCode},
};

use nix::{
    errno::Errno,
    libc::{pid_t, syscall, SYS_pidfd_open},
    sched::{setns, CloneFlags},
};
use procfs::process::Process;
use syd::{config::SYD_SH, 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 mut args = std::env::args().skip(1);

    let pid = match args.next().map(|arg| arg.parse::<pid_t>()) {
        Some(Ok(pid)) => pid,
        _ => {
            help();
            return Ok(ExitCode::FAILURE);
        }
    };
    let namespaces = match nsget(pid) {
        Ok(namespaces) => namespaces,
        Err(errno) => {
            eprintln!("syd-run: nsget: {errno}");
            return Ok(ExitCode::FAILURE);
        }
    };

    if !namespaces.is_empty() {
        if let Err(errno) = nsenter(pid, namespaces) {
            eprintln!("syd-run: nsenter: {errno}");
            return Ok(ExitCode::FAILURE);
        }
    }
    eprintln!("syd-run: {namespaces:#?}");

    let mut cmd = match Command::new(args.next().unwrap_or(SYD_SH.to_string()))
        .args(args)
        .spawn()
    {
        Ok(cmd) => cmd,
        Err(error) => {
            eprintln!("syd-run: spawn: {error}");
            return Ok(ExitCode::FAILURE);
        }
    };

    Ok(match cmd.wait() {
        Ok(status) => {
            if let Some(code) = status.code() {
                ExitCode::from(code as u8)
            } else {
                ExitCode::FAILURE
            }
        }
        Err(error) => {
            eprintln!("syd-run: wait: {error}");
            ExitCode::FAILURE
        }
    })
}

fn help() {
    println!("Usage: syd-run pid [<program> [<argument>...]]");
    println!("Run a program inside a syd container (requires Linux-5.8 or newer).");
}

fn nsenter(pid: pid_t, namespaces: CloneFlags) -> Result<(), Errno> {
    setns(pidfd_open(pid)?, namespaces)
}

fn nsget(pid: pid_t) -> SydResult<CloneFlags> {
    let current_proc = Process::myself()?;
    let current_namespaces = current_proc.namespaces()?;

    let target_proc = Process::new(pid)?;
    let target_namespaces = target_proc.namespaces()?.0;

    let mut flags = CloneFlags::empty();

    for (name, target_ns) in target_namespaces {
        if let Some(current_ns) = current_namespaces.0.get(&name) {
            if target_ns.identifier != current_ns.identifier {
                flags |= match name.as_bytes() {
                    b"cgroup" => CloneFlags::CLONE_NEWCGROUP,
                    b"ipc" => CloneFlags::CLONE_NEWIPC,
                    b"mnt" => CloneFlags::CLONE_NEWNS,
                    b"net" => CloneFlags::CLONE_NEWNET,
                    b"pid" => CloneFlags::CLONE_NEWPID,
                    b"user" => CloneFlags::CLONE_NEWUSER,
                    b"uts" => CloneFlags::CLONE_NEWUTS,
                    _ => continue,
                };
            }
        }
    }

    Ok(flags)
}

fn pidfd_open(pid: pid_t) -> Result<OwnedFd, Errno> {
    // SAFETY: Open a PID file descriptor
    match unsafe { syscall(SYS_pidfd_open, pid, 0) } {
        -1 => Err(Errno::last()),
        fd => Ok(unsafe { OwnedFd::from_raw_fd(fd as i32) }),
    }
}
