lazaret/crates/lazaret-linux/src/process.rs

use std::ffi::{OsStr, OsString};
use std::fs::File;
use std::io::{self, Read, Write};
use std::os::fd::{AsFd, AsRawFd, BorrowedFd, OwnedFd};
use std::process::Stdio;
use std::sync::{Mutex, OnceLock};

use lazaret_core::identity::{
    BootId, CoordinatorId, DomainId, GenerationId, HostPidLocator, PidfsDeviceId, PidfsInodeId,
    PidfsProcessAuthorityV1, PidfsProcessIdentityV1, ProcessStartTimeTicks,
};
use lazaret_core::records::{
    HostProcessOwnershipSlot, ProcessOwnershipCommitAck, StagedProcessOwnership,
};
use rustix::event::{PollFd, PollFlags, Timespec, poll};
use rustix::fs::{OFlags, fcntl_getfl, fcntl_setfl, fstat, fstatfs};
use rustix::io::{FdFlags, fcntl_setfd};
use rustix::pipe::{PipeFlags, pipe_with};
use rustix::process::{Pid, PidfdFlags, Signal, WaitId, WaitIdOptions};

pub const PID_FS_MAGIC: i64 = 0x5049_4446;
const BARRIER_RELEASE: &[u8; 32] = b"LAZARET-CAUSAL-BARRIER-RELEASE!!";
pub const PRIVATE_PROCESS_ROLE: &str = "--lazaret-private-process-role";
const WRAPPER_ROLE: &str = "launch-wrapper";
const MAPPED_WRAPPER_ROLE: &str = "mapped-launch-wrapper";
const SIGCHLD_PROBE_ROLE: &str = "sigchld-probe";
static CHILD_BROKER: OnceLock<ChildBroker> = OnceLock::new();

#[derive(Debug)]
pub struct TargetCommand {
    program: OsString,
    arguments: Vec<OsString>,
    parent_edge: ParentEdge,
    executable: Option<OwnedFd>,
    stdin: Option<File>,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum TargetStdinAccess {
    ReadOnly,
    WriteOnly,
    ReadWrite,
}

impl TargetStdinAccess {
    fn from_flags(flags: OFlags) -> io::Result<Self> {
        if flags.contains(OFlags::PATH) {
            return Err(io::Error::other(
                "target stdin descriptor has no read or write authority",
            ));
        }
        match flags & OFlags::ACCMODE {
            flags if flags == OFlags::RDONLY => Ok(Self::ReadOnly),
            flags if flags == OFlags::WRONLY => Ok(Self::WriteOnly),
            flags if flags == OFlags::RDWR => Ok(Self::ReadWrite),
            _ => Err(io::Error::other(
                "target stdin descriptor has an unsupported access mode",
            )),
        }
    }

    fn from_descriptor(descriptor: i32) -> io::Result<Self> {
        let metadata = std::fs::read_to_string(format!("/proc/self/fdinfo/{descriptor}"))?;
        let flags = metadata
            .lines()
            .find_map(|line| line.strip_prefix("flags:"))
            .map(str::trim)
            .ok_or_else(|| io::Error::other("target stdin descriptor flags are missing"))
            .and_then(|flags| {
                u32::from_str_radix(flags, 8)
                    .map(OFlags::from_bits_retain)
                    .map_err(|_| io::Error::other("target stdin descriptor flags are invalid"))
            })?;
        Self::from_flags(flags)
    }

    fn reopen(self, descriptor: i32) -> io::Result<File> {
        let mut options = File::options();
        match self {
            Self::ReadOnly => {
                options.read(true);
            }
            Self::WriteOnly => {
                options.write(true);
            }
            Self::ReadWrite => {
                options.read(true).write(true);
            }
        }
        options.open(format!("/proc/self/fd/{descriptor}"))
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ParentEdge {
    ParentDeathKill,
    LifetimePipe,
}

impl TargetCommand {
    #[must_use]
    pub fn new(program: impl Into<OsString>) -> Self {
        Self {
            program: program.into(),
            arguments: Vec::new(),
            parent_edge: ParentEdge::ParentDeathKill,
            executable: None,
            stdin: None,
        }
    }

    /// Constructs an argv-only target bound to the exact inherited executable descriptor.
    ///
    /// # Errors
    /// Returns an error when the descriptor cannot be made inheritable by the launch wrapper.
    pub fn from_executable(
        executable: OwnedFd,
        arguments: impl IntoIterator<Item = OsString>,
    ) -> io::Result<Self> {
        fcntl_setfd(executable.as_fd(), FdFlags::empty())?;
        Ok(Self {
            program: OsString::from(format!("/proc/self/fd/{}", executable.as_raw_fd())),
            arguments: arguments.into_iter().collect(),
            parent_edge: ParentEdge::ParentDeathKill,
            executable: Some(executable),
            stdin: None,
        })
    }

    #[must_use]
    pub fn argument(mut self, argument: impl Into<OsString>) -> Self {
        self.arguments.push(argument.into());
        self
    }

    #[must_use]
    pub fn with_lifetime_pipe_parent_edge(mut self) -> Self {
        self.parent_edge = ParentEdge::LifetimePipe;
        self
    }

    #[must_use]
    pub fn with_parent_death_kill_edge(mut self) -> Self {
        self.parent_edge = ParentEdge::ParentDeathKill;
        self
    }

    /// Attaches the target's standard input after the private wrapper consumes its barrier.
    ///
    /// # Errors
    /// Returns an error when the descriptor cannot be inherited by the wrapper.
    pub fn with_target_stdin(mut self, stdin: File) -> io::Result<Self> {
        TargetStdinAccess::from_flags(fcntl_getfl(&stdin)?)?;
        fcntl_setfd(stdin.as_fd(), FdFlags::empty())?;
        self.stdin = Some(stdin);
        Ok(self)
    }

    #[must_use]
    pub const fn parent_edge(&self) -> ParentEdge {
        self.parent_edge
    }

    #[must_use]
    pub fn program(&self) -> &OsStr {
        &self.program
    }

    #[must_use]
    pub fn arguments(&self) -> &[OsString] {
        &self.arguments
    }

    #[allow(clippy::disallowed_types)]
    fn append_to(&self, command: &mut std::process::Command) {
        let _ = &self.executable;
        command
            .arg("--parent-action")
            .arg(match self.parent_edge {
                ParentEdge::ParentDeathKill => "kill",
                ParentEdge::LifetimePipe => "lifetime-pipe",
            })
            .args(
                self.stdin
                    .as_ref()
                    .map(|stdin| {
                        [
                            OsString::from("--target-stdin-fd"),
                            OsString::from(stdin.as_raw_fd().to_string()),
                        ]
                    })
                    .into_iter()
                    .flatten(),
            )
            .arg("--")
            .arg(&self.program)
            .args(&self.arguments);
    }
}

pub struct ChildStandardIo {
    stdout: Option<Stdio>,
    stderr: Option<Stdio>,
    policy: ChildStandardIoPolicy,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ChildStandardIoPolicy {
    Inherit,
    PrivateControlWithLifetimePipe,
    Silenced,
    Terminal,
}

impl ChildStandardIo {
    #[must_use]
    pub const fn inherit() -> Self {
        Self {
            stdout: None,
            stderr: None,
            policy: ChildStandardIoPolicy::Inherit,
        }
    }

    #[must_use]
    pub fn private_control_with_lifetime_pipe(
        control: impl Into<Stdio>,
        lifetime_pipe: impl Into<Stdio>,
    ) -> Self {
        Self {
            stdout: Some(control.into()),
            stderr: Some(lifetime_pipe.into()),
            policy: ChildStandardIoPolicy::PrivateControlWithLifetimePipe,
        }
    }

    #[must_use]
    pub fn silenced() -> Self {
        Self {
            stdout: Some(Stdio::null()),
            stderr: Some(Stdio::null()),
            policy: ChildStandardIoPolicy::Silenced,
        }
    }

    /// # Errors
    /// Returns an error when the output descriptor cannot be duplicated for merged terminal I/O.
    pub fn terminal(output: File) -> io::Result<Self> {
        let error = output.try_clone()?;
        Ok(Self {
            stdout: Some(Stdio::from(output)),
            stderr: Some(Stdio::from(error)),
            policy: ChildStandardIoPolicy::Terminal,
        })
    }

    #[must_use]
    pub const fn policy(&self) -> ChildStandardIoPolicy {
        self.policy
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct ProcessOwnershipContext {
    pub domain_id: DomainId,
    pub generation_id: GenerationId,
    pub coordinator_id: CoordinatorId,
    pub slot: HostProcessOwnershipSlot,
}

pub struct AcknowledgedChildAcquisition<'target> {
    pub expected_parent: PidfsProcessIdentityV1,
    pub ownership: ProcessOwnershipContext,
    pub target: &'target TargetCommand,
    pub standard_io: ChildStandardIo,
}

#[derive(Debug)]
pub struct RetainedProcess {
    pidfd: OwnedFd,
    identity: PidfsProcessIdentityV1,
}

#[derive(Debug)]
pub enum PidfsRecoveryMatch {
    BootMismatch,
    Matched(RetainedProcess),
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum RecoveryActionResult<T> {
    BootReconciled(T),
    ProcessAction(T),
}

#[derive(Debug)]
pub struct CapturedProcess {
    process: RetainedProcess,
    stdout: File,
    stderr: File,
}

pub struct BootstrapChild {
    process: RetainedProcess,
    status: File,
}

impl BootstrapChild {
    /// Reads the private bootstrap acquisition status byte.
    ///
    /// # Errors
    /// Returns an error when the child closes or corrupts its status channel.
    pub fn acquisition_status(&mut self) -> io::Result<u8> {
        let mut status = [0_u8; 1];
        self.status.read_exact(&mut status)?;
        Ok(status[0])
    }

    /// # Errors
    /// Returns an error when polling the exact child pidfd fails.
    pub fn has_exited(&self) -> io::Result<bool> {
        self.process.has_exited()
    }

    /// Reaps the exact retained child through its pidfd.
    ///
    /// # Errors
    /// Returns an error when serialized exact reaping fails.
    pub fn reap(self) -> io::Result<ProcessExit> {
        child_broker().reap_status(self.process)
    }
}

impl CapturedProcess {
    #[must_use]
    pub const fn process(&self) -> &RetainedProcess {
        &self.process
    }

    pub fn streams(&mut self) -> (&mut File, &mut File) {
        (&mut self.stdout, &mut self.stderr)
    }

    #[must_use]
    pub fn into_process(self) -> RetainedProcess {
        self.process
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct ProcessExit {
    pub exit_code: Option<i32>,
    pub terminating_signal: Option<i32>,
}

impl RetainedProcess {
    #[must_use]
    pub const fn identity(&self) -> PidfsProcessIdentityV1 {
        self.identity
    }

    /// # Errors
    /// Returns an error when the exact retained-pidfd action fails.
    #[allow(clippy::disallowed_methods)]
    pub fn send_signal(&self, signal: Signal) -> io::Result<()> {
        rustix::process::pidfd_send_signal(self.pidfd.as_fd(), signal).map_err(Into::into)
    }

    /// # Errors
    /// Returns an error when polling the retained pidfd fails.
    pub fn has_exited(&self) -> io::Result<bool> {
        let mut descriptors = [PollFd::new(&self.pidfd, PollFlags::IN)];
        poll(&mut descriptors, Some(&Timespec::default()))?;
        Ok(descriptors[0]
            .revents()
            .intersects(PollFlags::IN | PollFlags::HUP))
    }
}

impl AsFd for RetainedProcess {
    fn as_fd(&self) -> BorrowedFd<'_> {
        self.pidfd.as_fd()
    }
}

#[derive(Debug)]
pub struct ChildBroker {
    exclusion: Mutex<()>,
}

#[must_use]
pub fn child_broker() -> &'static ChildBroker {
    CHILD_BROKER.get_or_init(|| ChildBroker {
        exclusion: Mutex::new(()),
    })
}

fn require_parent_edge_standard_io(
    parent_edge: ParentEdge,
    standard_io: ChildStandardIoPolicy,
) -> io::Result<()> {
    match (parent_edge, standard_io) {
        (ParentEdge::LifetimePipe, ChildStandardIoPolicy::PrivateControlWithLifetimePipe)
        | (
            ParentEdge::ParentDeathKill,
            ChildStandardIoPolicy::Inherit
            | ChildStandardIoPolicy::Silenced
            | ChildStandardIoPolicy::Terminal,
        ) => Ok(()),
        (ParentEdge::LifetimePipe, _) => Err(io::Error::other(
            "lifetime-pipe parent edge requires private control stdout and lifetime-pipe stderr",
        )),
        (ParentEdge::ParentDeathKill, ChildStandardIoPolicy::PrivateControlWithLifetimePipe) => {
            Err(io::Error::other(
                "parent-death edge cannot substitute a lifetime-pipe stdio policy",
            ))
        }
    }
}

fn require_captured_parent_edge(parent_edge: ParentEdge) -> io::Result<()> {
    if parent_edge == ParentEdge::ParentDeathKill {
        Ok(())
    } else {
        Err(io::Error::other(
            "captured output requires the parent-death edge",
        ))
    }
}

mod broker;
mod launch;
mod recovery;

pub use broker::{dispatch_private_process_role, verify_direct_child_pidfd};
pub use launch::{
    BarrierRelease, BarrierWait, append_parent_identity, exec_target, launch_barrier,
    parse_wrapper_arguments, suspend_current_process,
};
pub use recovery::{
    arm_parent_death_and_recheck, current_process_identity, match_recovery_process, parse_boot_id,
    read_boot_id, read_start_time, recover_process, recover_process_action, require_pidfs_magic,
};

#[cfg(test)]
mod tests;