lazaret/crates/lazaret-host-coordinator/src/sandbox_init.rs

use std::ffi::{OsStr, OsString};
use std::io;
use std::os::unix::ffi::OsStrExt as _;

use lazaret_linux::process::{TargetCommand, exec_target};

pub const PRIVATE_SANDBOX_INIT_ROLE: &str = "--lazaret-private-sandbox-init-v1";
pub const PRIVATE_SESSION_EXEC_ROLE: &str = "--lazaret-private-session-exec-v1";
pub const PRIVATE_SESSION_IDENTITY_ROLE: &str = "--lazaret-private-session-identity-v1";

#[must_use]
pub fn is_private_sandbox_init_role(arguments: &[OsString]) -> bool {
    arguments.len() == 1 && arguments[0] == OsStr::new(PRIVATE_SANDBOX_INIT_ROLE)
}

#[must_use]
pub fn dispatch_private_sandbox_init(arguments: &[OsString]) -> Option<io::Result<()>> {
    if is_private_sandbox_init_role(arguments) {
        loop {
            std::thread::park();
        }
    }
    if arguments
        .first()
        .is_some_and(|argument| argument == PRIVATE_SESSION_EXEC_ROLE)
    {
        return Some(run_session_exec(&arguments[1..]));
    }
    if arguments
        .first()
        .is_some_and(|argument| argument == PRIVATE_SESSION_IDENTITY_ROLE)
    {
        return Some(report_session_identity(&arguments[1..]));
    }
    None
}

fn run_session_exec(arguments: &[OsString]) -> io::Result<()> {
    let (separator, argv) = arguments
        .split_first()
        .filter(|(separator, argv)| {
            separator == &OsStr::new("--")
                && !argv.is_empty()
                && !argv[0].is_empty()
                && argv
                    .iter()
                    .all(|argument| !argument.as_bytes().contains(&0))
        })
        .ok_or_else(|| io::Error::other("session exec requires -- and a nonempty argv"))?;
    let _ = separator;
    rustix::process::setsid()?;
    let target = argv[1..]
        .iter()
        .cloned()
        .fold(TargetCommand::new(argv[0].clone()), TargetCommand::argument);
    Err(exec_target(target))
}

fn report_session_identity(arguments: &[OsString]) -> io::Result<()> {
    match arguments {
        [pid] => report_process_identity(parse_positive_pid(pid)?),
        [mode, group] if mode == "--group" => report_group_presence(parse_positive_pid(group)?),
        _ => Err(io::Error::other("session identity query is invalid")),
    }
}

fn parse_positive_pid(value: &OsStr) -> io::Result<u32> {
    value
        .to_str()
        .and_then(|value| value.parse::<u32>().ok())
        .filter(|value| *value > 0)
        .ok_or_else(|| io::Error::other("session identity number is invalid"))
}

fn report_process_identity(pid: u32) -> io::Result<()> {
    let Some((observed_pid, start_ticks, group)) = read_process_identity(pid)? else {
        println!("absent");
        return Ok(());
    };
    println!("{observed_pid} {start_ticks} {group}");
    Ok(())
}

fn report_group_presence(group: u32) -> io::Result<()> {
    println!(
        "{}",
        if group_is_present_except(group, std::process::id())? {
            "present"
        } else {
            "absent"
        }
    );
    Ok(())
}

fn group_is_present_except(group: u32, observer: u32) -> io::Result<bool> {
    for entry in std::fs::read_dir("/proc")? {
        let entry = entry?;
        let Some(pid) = entry
            .file_name()
            .to_str()
            .and_then(|value| value.parse::<u32>().ok())
        else {
            continue;
        };
        if pid == observer {
            continue;
        }
        if read_process_identity(pid)?.is_some_and(|(_, _, observed_group)| observed_group == group)
        {
            return Ok(true);
        }
    }
    Ok(false)
}

fn read_process_identity(pid: u32) -> io::Result<Option<(u32, u64, u32)>> {
    let stat = match std::fs::read_to_string(format!("/proc/{pid}/stat")) {
        Ok(stat) => stat,
        Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(None),
        Err(error) => return Err(error),
    };
    let close = stat
        .rfind(')')
        .ok_or_else(|| io::Error::other("guest process stat is malformed"))?;
    let observed_pid = stat[..close]
        .split_once(' ')
        .and_then(|(pid, _)| pid.parse::<u32>().ok())
        .ok_or_else(|| io::Error::other("guest process stat PID is malformed"))?;
    let fields = stat[close + 1..].split_whitespace().collect::<Vec<_>>();
    let group = fields
        .get(2)
        .and_then(|value| value.parse::<u32>().ok())
        .ok_or_else(|| io::Error::other("guest process stat group is malformed"))?;
    let start_ticks = fields
        .get(19)
        .and_then(|value| value.parse::<u64>().ok())
        .ok_or_else(|| io::Error::other("guest process stat start is malformed"))?;
    Ok(Some((observed_pid, start_ticks, group)))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn private_role_dispatch_is_exact_and_closed() {
        assert!(is_private_sandbox_init_role(&[OsString::from(
            PRIVATE_SANDBOX_INIT_ROLE
        )]));
        assert!(!is_private_sandbox_init_role(&[]));
        assert!(!is_private_sandbox_init_role(&[
            OsString::from(PRIVATE_SANDBOX_INIT_ROLE),
            OsString::from("extra"),
        ]));
        assert!(!is_private_sandbox_init_role(&[OsString::from(
            "--private-sandbox-init"
        )]));
    }

    #[test]
    fn private_session_roles_are_shape_separated() {
        assert!(dispatch_private_sandbox_init(&[]).is_none());
        assert!(
            dispatch_private_sandbox_init(&[
                OsString::from(PRIVATE_SESSION_EXEC_ROLE),
                OsString::from("missing-separator"),
            ])
            .unwrap()
            .is_err()
        );
        assert!(
            dispatch_private_sandbox_init(&[
                OsString::from(PRIVATE_SESSION_IDENTITY_ROLE),
                OsString::from("not-a-pid"),
            ])
            .unwrap()
            .is_err()
        );
    }

    #[test]
    fn group_presence_is_independent_of_a_designated_leader() {
        let (_, _, group) = read_process_identity(std::process::id()).unwrap().unwrap();
        assert!(group_is_present_except(group, u32::MAX).unwrap());
        assert!(!group_is_present_except(u32::MAX, u32::MAX).unwrap());
    }
}