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

use super::launch::{BarrierWait, append_parent_identity, exec_target, parse_wrapper_arguments};
use super::recovery::{
    arm_parent_death_and_recheck, authority_from_pidfd, current_process_identity, read_boot_id,
    read_start_time, require_default_sigchld, require_host_pid_namespace, require_reapable_sigchld,
};
use super::{
    AcknowledgedChildAcquisition, AsFd, BarrierRelease, BootstrapChild, CapturedProcess,
    ChildBroker, ChildStandardIo, ChildStandardIoPolicy, File, HostPidLocator, OFlags, OsString,
    OwnedFd, PRIVATE_PROCESS_ROLE, ParentEdge, Pid, PidfdFlags, PidfsProcessIdentityV1, PollFd,
    PollFlags, ProcessExit, ProcessOwnershipCommitAck, ProcessOwnershipContext,
    ProcessStartTimeTicks, RetainedProcess, SIGCHLD_PROBE_ROLE, Signal, StagedProcessOwnership,
    Stdio, TargetCommand, Timespec, WRAPPER_ROLE, WaitId, WaitIdOptions, fcntl_getfl, fcntl_setfl,
    io, launch_barrier, poll, recover_process, require_captured_parent_edge,
    require_parent_edge_standard_io,
};

impl ChildBroker {
    /// Spawns the private Host Coordinator bootstrap with exact pipe ownership and pidfd reaping.
    ///
    /// # Errors
    /// Refuses a command without the lifetime-pipe edge or any spawn, pidfd, or status failure.
    #[allow(clippy::disallowed_methods, clippy::disallowed_types)]
    pub fn spawn_bootstrap(
        &self,
        target: &TargetCommand,
        lifetime_reader: File,
    ) -> io::Result<BootstrapChild> {
        let _guard = self
            .exclusion
            .lock()
            .map_err(|_| io::Error::other("child broker lock poisoned"))?;
        if target.parent_edge() != ParentEdge::LifetimePipe {
            return Err(io::Error::other(
                "bootstrap child requires the lifetime-pipe edge",
            ));
        }
        require_default_sigchld()?;
        require_reapable_sigchld()?;
        let mut command = std::process::Command::new(&target.program);
        command
            .args(&target.arguments)
            .stdin(Stdio::from(lifetime_reader))
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit());
        let mut child = command.spawn()?;
        let status = child
            .stdout
            .take()
            .ok_or_else(|| io::Error::other("bootstrap status pipe is absent"))?;
        let process = match acquire_spawned_child(child.id()) {
            Ok(process) => process,
            Err(error) => {
                let _ = child.kill();
                let _ = child.wait();
                return Err(error);
            }
        };
        drop(child);
        Ok(BootstrapChild {
            process,
            status: File::from(OwnedFd::from(status)),
        })
    }

    /// Runs a non-durable direct child under the same spawn/reap exclusion as owned children.
    ///
    /// # Errors
    /// Returns an error when causal launch, exact waiting, or the command exit status fails.
    #[allow(clippy::disallowed_methods, clippy::disallowed_types)]
    pub fn run_transient(&self, target: &TargetCommand) -> io::Result<ProcessExit> {
        self.run_transient_with_stdout(target, None)
    }

    /// Runs a non-durable direct child with one private inherited stdout descriptor.
    ///
    /// # Errors
    /// Returns an error when causal launch, exact waiting, or the command exit status fails.
    pub fn run_transient_with_private_stdout(
        &self,
        target: &TargetCommand,
        stdout: OwnedFd,
    ) -> io::Result<ProcessExit> {
        self.run_transient_with_stdout(target, Some(stdout))
    }

    #[allow(clippy::disallowed_methods, clippy::disallowed_types)]
    fn run_transient_with_stdout(
        &self,
        target: &TargetCommand,
        stdout: Option<OwnedFd>,
    ) -> io::Result<ProcessExit> {
        let _guard = self
            .exclusion
            .lock()
            .map_err(|_| io::Error::other("child broker lock poisoned"))?;
        require_parent_edge_standard_io(target.parent_edge(), ChildStandardIoPolicy::Silenced)?;
        require_default_sigchld()?;
        require_reapable_sigchld()?;
        let parent = current_process_identity()?;
        let (release, wait) = launch_barrier()?;
        let mut command = private_role_command(WRAPPER_ROLE);
        append_parent_identity(&mut command, parent);
        target.append_to(&mut command);
        command
            .stdin(wait.into_stdio())
            .stdout(stdout.map_or_else(Stdio::null, Stdio::from))
            .stderr(Stdio::null());
        let mut child = command.spawn()?;
        let process = match acquire_spawned_child(child.id()) {
            Ok(process) => process,
            Err(error) => return fail_before_pidfd(release, &mut child, error),
        };
        if let Err(error) = release.release() {
            return fail_after_pidfd_closed(process, error);
        }
        drop(child);
        let status =
            rustix::process::waitid(WaitId::PidFd(process.pidfd.as_fd()), WaitIdOptions::EXITED)?
                .ok_or_else(|| io::Error::other("transient child exited without wait status"))?;
        let exit = ProcessExit {
            exit_code: status.exit_status(),
            terminating_signal: status.terminating_signal(),
        };
        drop(process);
        if exit.exit_code == Some(0) && exit.terminating_signal.is_none() {
            Ok(exit)
        } else {
            Err(io::Error::other(format!(
                "transient child failed with exit {:?} and signal {:?}",
                exit.exit_code, exit.terminating_signal
            )))
        }
    }

    /// Starts a non-durable direct child with split captured output and exact retained authority.
    ///
    /// # Errors
    /// Refuses invalid parent-edge or SIGCHLD state and cleans up on spawn, pidfd, barrier, or
    /// pipe-acquisition failure.
    #[allow(clippy::disallowed_methods, clippy::disallowed_types)]
    pub fn run_transient_captured(&self, target: &TargetCommand) -> io::Result<CapturedProcess> {
        let _guard = self
            .exclusion
            .lock()
            .map_err(|_| io::Error::other("child broker lock poisoned"))?;
        require_captured_parent_edge(target.parent_edge())?;
        require_default_sigchld()?;
        require_reapable_sigchld()?;
        let parent = current_process_identity()?;
        let (release, wait) = launch_barrier()?;
        let mut command = private_role_command(WRAPPER_ROLE);
        append_parent_identity(&mut command, parent);
        target.append_to(&mut command);
        command
            .stdin(wait.into_stdio())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());
        let mut child = command.spawn()?;
        let process = match acquire_spawned_child(child.id()) {
            Ok(process) => process,
            Err(error) => return fail_before_pidfd(release, &mut child, error),
        };
        if let Err(error) = release.release() {
            return fail_after_pidfd_closed(process, error);
        }
        captured_process(process, child)
    }

    /// Spawns the repository wrapper, stages exact ownership, accepts its acknowledgment,
    /// and only then releases target execution.
    ///
    /// # Errors
    /// Refuses and reaps the blocked wrapper on every failed acquisition or callback.
    #[allow(clippy::disallowed_methods, clippy::disallowed_types)]
    pub fn acquire<Stage, Acknowledge>(
        &self,
        expected_parent: PidfsProcessIdentityV1,
        ownership: ProcessOwnershipContext,
        target: &TargetCommand,
        stage: Stage,
        acknowledge: Acknowledge,
    ) -> io::Result<RetainedProcess>
    where
        Stage: FnOnce(StagedProcessOwnership) -> io::Result<()>,
        Acknowledge: FnOnce(StagedProcessOwnership) -> io::Result<ProcessOwnershipCommitAck>,
    {
        self.acquire_with_stdio(
            expected_parent,
            ownership,
            target,
            ChildStandardIo::inherit(),
            stage,
            acknowledge,
        )
    }

    /// # Errors
    /// Refuses and reaps the blocked wrapper on every failed acquisition or callback.
    #[allow(clippy::disallowed_methods, clippy::disallowed_types)]
    pub fn acquire_with_stdio<Stage, Acknowledge>(
        &self,
        expected_parent: PidfsProcessIdentityV1,
        ownership: ProcessOwnershipContext,
        target: &TargetCommand,
        standard_io: ChildStandardIo,
        stage: Stage,
        acknowledge: Acknowledge,
    ) -> io::Result<RetainedProcess>
    where
        Stage: FnOnce(StagedProcessOwnership) -> io::Result<()>,
        Acknowledge: FnOnce(StagedProcessOwnership) -> io::Result<ProcessOwnershipCommitAck>,
    {
        self.acquire_with_stdio_after_ack(
            AcknowledgedChildAcquisition {
                expected_parent,
                ownership,
                target,
                standard_io,
            },
            stage,
            acknowledge,
            || Ok(()),
        )
    }

    /// Acquires an exact child while retaining its blocked wrapper across one acknowledged
    /// parent transition.
    ///
    /// # Errors
    /// Refuses and reaps the blocked wrapper when staging, acknowledgment, transition, or release
    /// fails.
    #[allow(clippy::disallowed_methods, clippy::disallowed_types)]
    pub fn acquire_with_stdio_after_ack<Stage, Acknowledge, Transition>(
        &self,
        acquisition: AcknowledgedChildAcquisition<'_>,
        stage: Stage,
        acknowledge: Acknowledge,
        transition: Transition,
    ) -> io::Result<RetainedProcess>
    where
        Stage: FnOnce(StagedProcessOwnership) -> io::Result<()>,
        Acknowledge: FnOnce(StagedProcessOwnership) -> io::Result<ProcessOwnershipCommitAck>,
        Transition: FnOnce() -> io::Result<()>,
    {
        let _guard = self
            .exclusion
            .lock()
            .map_err(|_| io::Error::other("child broker lock poisoned"))?;
        require_parent_edge_standard_io(
            acquisition.target.parent_edge(),
            acquisition.standard_io.policy(),
        )?;
        require_default_sigchld()?;
        require_reapable_sigchld()?;
        let (release, wait) = launch_barrier()?;
        let mut command = private_role_command(WRAPPER_ROLE);
        append_parent_identity(&mut command, acquisition.expected_parent);
        acquisition.target.append_to(&mut command);
        command.stdin(wait.into_stdio());
        let mut standard_io = acquisition.standard_io;
        if let Some(stdout) = standard_io.stdout.take() {
            command.stdout(stdout);
        }
        if let Some(stderr) = standard_io.stderr.take() {
            command.stderr(stderr);
        }
        let mut child = command.spawn()?;
        let host_pid = child.id();
        let acquired = acquire_spawned_child(host_pid);
        let process = match acquired {
            Ok(process) => process,
            Err(error) => return fail_before_pidfd(release, &mut child, error),
        };
        let staged = StagedProcessOwnership::v1(
            acquisition.ownership.domain_id,
            acquisition.ownership.generation_id,
            acquisition.ownership.coordinator_id,
            acquisition.ownership.slot,
            process.identity,
        );
        if let Err(error) = stage(staged) {
            return fail_after_pidfd(release, process, error);
        }
        let acknowledgment = match acknowledge(staged) {
            Ok(acknowledgment) => acknowledgment,
            Err(error) => return fail_after_pidfd(release, process, error),
        };
        if !acknowledgment.acknowledges(&staged) {
            return fail_after_pidfd(
                release,
                process,
                io::Error::other(
                    "process ownership acknowledgment does not match staged ownership",
                ),
            );
        }
        if let Err(error) = transition() {
            return fail_after_pidfd(release, process, error);
        }
        if let Err(error) = release.release() {
            return fail_after_pidfd_closed(process, error);
        }
        drop(child);
        Ok(process)
    }

    /// Acquires an exact child with bounded-output pipe ownership retained by the caller.
    ///
    /// # Errors
    /// Refuses and reaps the blocked wrapper on every failed acquisition or callback.
    #[allow(clippy::disallowed_methods, clippy::disallowed_types)]
    pub fn acquire_captured<Stage, Acknowledge>(
        &self,
        expected_parent: PidfsProcessIdentityV1,
        ownership: ProcessOwnershipContext,
        target: &TargetCommand,
        stage: Stage,
        acknowledge: Acknowledge,
    ) -> io::Result<CapturedProcess>
    where
        Stage: FnOnce(StagedProcessOwnership) -> io::Result<()>,
        Acknowledge: FnOnce(StagedProcessOwnership) -> io::Result<ProcessOwnershipCommitAck>,
    {
        let _guard = self
            .exclusion
            .lock()
            .map_err(|_| io::Error::other("child broker lock poisoned"))?;
        require_captured_parent_edge(target.parent_edge())?;
        require_default_sigchld()?;
        require_reapable_sigchld()?;
        let (release, wait) = launch_barrier()?;
        let mut command = private_role_command(WRAPPER_ROLE);
        append_parent_identity(&mut command, expected_parent);
        target.append_to(&mut command);
        command
            .stdin(wait.into_stdio())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());
        let mut child = command.spawn()?;
        let host_pid = child.id();
        let process = match acquire_spawned_child(host_pid) {
            Ok(process) => process,
            Err(error) => return fail_before_pidfd(release, &mut child, error),
        };
        let staged = StagedProcessOwnership::v1(
            ownership.domain_id,
            ownership.generation_id,
            ownership.coordinator_id,
            ownership.slot,
            process.identity,
        );
        if let Err(error) = stage(staged) {
            return fail_after_pidfd(release, process, error);
        }
        let acknowledgment = match acknowledge(staged) {
            Ok(acknowledgment) => acknowledgment,
            Err(error) => return fail_after_pidfd(release, process, error),
        };
        if !acknowledgment.acknowledges(&staged) {
            return fail_after_pidfd(
                release,
                process,
                io::Error::other(
                    "process ownership acknowledgment does not match staged ownership",
                ),
            );
        }
        if let Err(error) = release.release() {
            return fail_after_pidfd_closed(process, error);
        }
        captured_process(process, child)
    }

    /// # Errors
    /// Returns an error when serialization or exact pidfd reaping fails.
    #[allow(clippy::disallowed_methods)]
    pub fn reap(&self, process: RetainedProcess) -> io::Result<()> {
        self.reap_status(process).map(|_| ())
    }

    /// # Errors
    /// Returns an error when serialization or exact pidfd reaping fails.
    #[allow(clippy::disallowed_methods)]
    pub fn reap_status(&self, process: RetainedProcess) -> io::Result<ProcessExit> {
        let _guard = self
            .exclusion
            .lock()
            .map_err(|_| io::Error::other("child broker lock poisoned"))?;
        let status =
            rustix::process::waitid(WaitId::PidFd(process.pidfd.as_fd()), WaitIdOptions::EXITED)?
                .ok_or_else(|| io::Error::other("exact child exited without wait status"))?;
        let exit = ProcessExit {
            exit_code: status.exit_status(),
            terminating_signal: status.terminating_signal(),
        };
        drop(process);
        Ok(exit)
    }
}

#[allow(clippy::disallowed_types)]
fn captured_process(
    process: RetainedProcess,
    mut child: std::process::Child,
) -> io::Result<CapturedProcess> {
    let Some(stdout) = child.stdout.take() else {
        return fail_running_captured(
            process,
            child,
            io::Error::other("captured child stdout is unavailable"),
        );
    };
    let Some(stderr) = child.stderr.take() else {
        return fail_running_captured(
            process,
            child,
            io::Error::other("captured child stderr is unavailable"),
        );
    };
    let stdout = match child_pipe_file(stdout.into()) {
        Ok(stdout) => stdout,
        Err(error) => return fail_running_captured(process, child, error),
    };
    let stderr = match child_pipe_file(stderr.into()) {
        Ok(stderr) => stderr,
        Err(error) => return fail_running_captured(process, child, error),
    };
    drop(child);
    Ok(CapturedProcess {
        process,
        stdout,
        stderr,
    })
}

fn child_pipe_file(fd: OwnedFd) -> io::Result<File> {
    let flags = fcntl_getfl(fd.as_fd())?;
    fcntl_setfl(fd.as_fd(), flags | OFlags::NONBLOCK)?;
    Ok(File::from(fd))
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn fail_running_captured<T>(
    process: RetainedProcess,
    child: std::process::Child,
    error: io::Error,
) -> io::Result<T> {
    let _ = process.send_signal(Signal::KILL);
    rustix::process::waitid(WaitId::PidFd(process.pidfd.as_fd()), WaitIdOptions::EXITED)?;
    drop(child);
    drop(process);
    Err(error)
}

#[allow(clippy::disallowed_types)]
pub(super) fn private_role_command(role: &str) -> std::process::Command {
    let mut command = std::process::Command::new("/proc/self/exe");
    command.arg(PRIVATE_PROCESS_ROLE).arg(role);
    command
}

/// Runs a private self-exec role before normal argument handling.
///
/// # Errors
/// Returns the role's fail-closed error when a private invocation is malformed.
pub fn dispatch_private_process_role(arguments: &[OsString]) -> Option<io::Result<()>> {
    if arguments
        .first()
        .is_none_or(|argument| argument != PRIVATE_PROCESS_ROLE)
    {
        return None;
    }
    let role = arguments.get(1).and_then(|argument| argument.to_str());
    match role {
        Some(SIGCHLD_PROBE_ROLE) if arguments.len() == 2 => {
            Some(BarrierWait::from_stdin().and_then(BarrierWait::wait))
        }
        Some(WRAPPER_ROLE) => Some(run_wrapper_role(arguments[2..].to_vec())),
        _ => Some(Err(io::Error::other(
            "unknown or malformed private process role",
        ))),
    }
}

fn run_wrapper_role(arguments: Vec<OsString>) -> io::Result<()> {
    let (expected_parent, target) = parse_wrapper_arguments(arguments)?;
    BarrierWait::from_stdin()?.wait()?;
    rustix::io::fcntl_setfd(std::io::stdin(), rustix::io::FdFlags::CLOEXEC)?;
    match target.parent_edge {
        ParentEdge::ParentDeathKill => {
            arm_parent_death_and_recheck(Signal::KILL, expected_parent)?;
        }
        ParentEdge::LifetimePipe => {
            require_lifetime_pipe_and_recheck(expected_parent)?;
        }
    }
    if let Some(stdin) = target.stdin.as_ref() {
        rustix::stdio::dup2_stdin(stdin)?;
    }
    Err(exec_target(target))
}

fn require_lifetime_pipe_and_recheck(
    expected_parent: PidfsProcessIdentityV1,
) -> io::Result<RetainedProcess> {
    let stderr = std::io::stderr();
    let mut descriptors = [PollFd::new(&stderr, PollFlags::IN)];
    poll(&mut descriptors, Some(&Timespec::default()))?;
    if descriptors[0]
        .revents()
        .intersects(PollFlags::IN | PollFlags::HUP | PollFlags::ERR)
    {
        return Err(io::Error::other("parent lifetime pipe is not held"));
    }
    let parent = recover_process(expected_parent)?;
    let observed_parent = rustix::process::getppid()
        .map(Pid::as_raw_nonzero)
        .map(i32::from)
        .and_then(|pid| u32::try_from(pid).ok());
    if observed_parent != Some(expected_parent.host_pid().get()) {
        return Err(io::Error::other(
            "parent changed while arming lifetime pipe",
        ));
    }
    Ok(parent)
}

#[allow(clippy::disallowed_methods)]
fn acquire_spawned_child(host_pid: u32) -> io::Result<RetainedProcess> {
    let raw_pid = i32::try_from(host_pid)
        .map_err(|_| io::Error::other("spawn returned out-of-range child PID"))?;
    let pid = Pid::from_raw(raw_pid)
        .ok_or_else(|| io::Error::other("spawn returned invalid child PID"))?;
    require_host_pid_namespace(host_pid)?;
    let pidfd = rustix::process::pidfd_open(pid, PidfdFlags::empty())?;
    verify_direct_child_pidfd(pidfd.as_fd())?;
    let authority = authority_from_pidfd(pidfd.as_fd(), read_boot_id()?)?;
    let start_time = read_start_time(host_pid)?;
    Ok(RetainedProcess {
        pidfd,
        identity: PidfsProcessIdentityV1::new(
            authority,
            HostPidLocator::new(host_pid),
            ProcessStartTimeTicks::new(start_time),
        ),
    })
}

/// # Errors
/// Refuses descriptors that do not name a waitable direct child of this broker process.
#[allow(clippy::disallowed_methods)]
pub fn verify_direct_child_pidfd(pidfd: impl AsFd) -> io::Result<()> {
    rustix::process::waitid(
        WaitId::PidFd(pidfd.as_fd()),
        WaitIdOptions::EXITED | WaitIdOptions::NOHANG | WaitIdOptions::NOWAIT,
    )?;
    Ok(())
}

#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
fn fail_before_pidfd<T>(
    release: BarrierRelease,
    child: &mut std::process::Child,
    error: io::Error,
) -> io::Result<T> {
    drop(release);
    child.wait()?;
    Err(error)
}

fn fail_after_pidfd<T>(
    release: BarrierRelease,
    process: RetainedProcess,
    error: io::Error,
) -> io::Result<T> {
    drop(release);
    fail_after_pidfd_closed(process, error)
}

#[allow(clippy::disallowed_methods)]
fn fail_after_pidfd_closed<T>(process: RetainedProcess, error: io::Error) -> io::Result<T> {
    rustix::process::waitid(WaitId::PidFd(process.pidfd.as_fd()), WaitIdOptions::EXITED)?;
    drop(process);
    Err(error)
}