monitor.rs

use std::{convert::Infallible, ffi::c_int, io, process::Command};

use crate::exec::{opt_fmt, signal_fmt};
use crate::system::signal::{
    consts::*, register_handlers, SignalHandler, SignalHandlerBehavior, SignalNumber, SignalSet,
    SignalStream,
};
use crate::{
    common::bin_serde::BinPipe,
    exec::{
        event::{EventRegistry, Process},
        exec_command,
        io_util::{retry_while_interrupted, was_interrupted},
        use_pty::backchannel::{MonitorBackchannel, MonitorMessage, ParentMessage},
    },
};
use crate::{
    exec::{
        event::{PollEvent, StopReason},
        use_pty::{SIGCONT_BG, SIGCONT_FG},
    },
    log::{dev_error, dev_info, dev_warn},
};
use crate::{
    exec::{handle_sigchld, terminate_process, HandleSigchld},
    system::{
        _exit, fork, getpgid, getpgrp,
        interface::ProcessId,
        kill, setpgid, setsid,
        term::{PtyFollower, Terminal},
        wait::{Wait, WaitError, WaitOptions},
        ForkResult,
    },
};

use super::CommandStatus;

pub(super) fn exec_monitor(
    pty_follower: PtyFollower,
    command: Command,
    foreground: bool,
    backchannel: &mut MonitorBackchannel,
    original_set: Option<SignalSet>,
) -> io::Result<Infallible> {
    // SIGTTIN and SIGTTOU are ignored here but the docs state that it shouldn't
    // be possible to receive them in the first place. Investigate
    match SignalHandler::register(SIGTTIN, SignalHandlerBehavior::Ignore) {
        Ok(handler) => handler.forget(),
        Err(err) => dev_warn!("cannot set handler for SIGTTIN: {err}"),
    }
    match SignalHandler::register(SIGTTOU, SignalHandlerBehavior::Ignore) {
        Ok(handler) => handler.forget(),
        Err(err) => dev_warn!("cannot set handler for SIGTTOU: {err}"),
    }

    // Start a new terminal session with the monitor as the leader.
    setsid().map_err(|err| {
        dev_warn!("cannot start a new session: {err}");
        err
    })?;

    // Set the follower side of the pty as the controlling terminal for the session.
    pty_follower.make_controlling_terminal().map_err(|err| {
        dev_warn!("cannot set the controlling terminal: {err}");
        err
    })?;

    // Use a pipe to get the IO error if `exec_command` fails.
    let (errpipe_tx, errpipe_rx) = BinPipe::pair()?;

    // Wait for the parent to give us green light before spawning the command. This avoids race
    // conditions when the command exits quickly.
    let event = retry_while_interrupted(|| backchannel.recv()).map_err(|err| {
        dev_warn!("cannot receive green light from parent: {err}");
        err
    })?;
    // Given that `UnixStream` delivers messages in order it shouldn't be possible to
    // receive an event different to `Edge` at the beginning.
    debug_assert_eq!(event, MonitorMessage::Edge);

    // FIXME (ogsudo): Some extra config happens here if selinux is available.

    // SAFETY: There should be no other threads at this point.
    let ForkResult::Parent(command_pid) = unsafe { fork() }.map_err(|err| {
        dev_warn!("unable to fork command process: {err}");
        err
    })?
    else {
        drop(errpipe_rx);

        // FIXME (ogsudo): Do any additional configuration that needs to be run after `fork` but before `exec`
        let command_pid = ProcessId::new(std::process::id() as i32);

        setpgid(ProcessId::new(0), command_pid).ok();

        // Wait for the monitor to set us as the foreground group for the pty if we are in the
        // foreground.
        if foreground {
            while !pty_follower.tcgetpgrp().is_ok_and(|pid| pid == command_pid) {
                std::thread::yield_now();
            }
        }

        // Done with the pty follower.
        drop(pty_follower);

        exec_command(command, original_set, errpipe_tx)
    };

    // Send the command's PID to the parent.
    if let Err(err) = backchannel.send(&ParentMessage::CommandPid(command_pid)) {
        dev_warn!("cannot send command PID to parent: {err}");
    }

    let mut registry = EventRegistry::new();

    let mut closure = MonitorClosure::new(
        command_pid,
        pty_follower,
        errpipe_rx,
        backchannel,
        &mut registry,
    )?;

    // Restore the signal mask now that the handlers have been setup.
    if let Some(set) = original_set {
        if let Err(err) = set.set_mask() {
            dev_warn!("cannot restore signal mask: {err}");
        }
    }

    // Set the foreground group for the pty follower.
    if foreground {
        if let Err(err) = closure.pty_follower.tcsetpgrp(closure.command_pgrp) {
            dev_error!(
                "cannot set foreground progess group to {} (command): {err}",
                closure.command_pgrp
            );
        }
    }

    // FIXME (ogsudo): Here's where the signal mask is removed because the handlers for the signals
    // have been setup after initializing the closure.

    // Start the event loop.
    let reason = registry.event_loop(&mut closure);

    // Terminate the command if it's not terminated.
    if let Some(command_pid) = closure.command_pid {
        terminate_process(command_pid, true);

        loop {
            match command_pid.wait(WaitOptions::new()) {
                Err(WaitError::Io(err)) if was_interrupted(&err) => {}
                _ => break,
            }
        }
    }

    // Take the controlling tty so the command's children don't receive SIGHUP when we exit.
    if let Err(err) = closure.pty_follower.tcsetpgrp(closure.monitor_pgrp) {
        dev_error!(
            "cannot set foreground process group to {} (monitor): {err}",
            closure.monitor_pgrp
        );
    }

    // Disable nonblocking assetions as we will not poll the backchannel anymore.
    closure.backchannel.set_nonblocking_assertions(false);

    match reason {
        StopReason::Break(err) => match err.try_into() {
            Ok(msg) => {
                if let Err(err) = closure.backchannel.send(&msg) {
                    dev_warn!("cannot send message over backchannel: {err}")
                }
            }
            Err(err) => {
                dev_warn!("socket error `{err:?}` cannot be converted to a message")
            }
        },
        StopReason::Exit(command_status) => {
            if let Err(err) = closure.backchannel.send(&command_status.into()) {
                dev_warn!("cannot send message over backchannel: {err}")
            }
        }
    }

    // Wait for the parent to give us red light before shutting down. This avoids missing
    // output when the monitor exits too quickly.
    let event = retry_while_interrupted(|| backchannel.recv()).map_err(|err| {
        dev_warn!("cannot receive red light from parent: {err}");
        err
    })?;
    debug_assert_eq!(event, MonitorMessage::Edge);

    // FIXME (ogsudo): The tty is restored here if selinux is available.

    // We call `_exit` instead of `exit` to avoid flushing the parent's IO streams by accident.
    _exit(1);
}

struct MonitorClosure<'a> {
    /// The command PID.
    ///
    /// This is `Some` iff the process is still running.
    command_pid: Option<ProcessId>,
    command_pgrp: ProcessId,
    monitor_pgrp: ProcessId,
    pty_follower: PtyFollower,
    errpipe_rx: BinPipe<i32>,
    backchannel: &'a mut MonitorBackchannel,
    signal_stream: &'static SignalStream,
    _signal_handlers: [SignalHandler; MonitorClosure::SIGNALS.len()],
}

impl<'a> MonitorClosure<'a> {
    const SIGNALS: [SignalNumber; 8] = [
        SIGINT, SIGQUIT, SIGTSTP, SIGTERM, SIGHUP, SIGUSR1, SIGUSR2, SIGCHLD,
    ];

    fn new(
        command_pid: ProcessId,
        pty_follower: PtyFollower,
        errpipe_rx: BinPipe<i32>,
        backchannel: &'a mut MonitorBackchannel,
        registry: &mut EventRegistry<Self>,
    ) -> io::Result<Self> {
        // Store the pgid of the monitor.
        let monitor_pgrp = getpgrp();

        // Register the callback to receive the IO error if the command fails to execute.
        registry.register_event(&errpipe_rx, PollEvent::Readable, |_| {
            MonitorEvent::ReadableErrPipe
        });

        // Enable nonblocking assertions as we will poll this inside the event loop.
        backchannel.set_nonblocking_assertions(true);

        // Register the callback to receive events from the backchannel
        registry.register_event(backchannel, PollEvent::Readable, |_| {
            MonitorEvent::ReadableBackchannel
        });

        let signal_stream = SignalStream::init()?;

        registry.register_event(signal_stream, PollEvent::Readable, |_| MonitorEvent::Signal);

        let signal_handlers = register_handlers(Self::SIGNALS)?;

        // Put the command in its own process group.
        let command_pgrp = command_pid;
        if let Err(err) = setpgid(command_pid, command_pgrp) {
            dev_warn!("cannot set process group ID for process: {err}");
        };

        Ok(Self {
            command_pid: Some(command_pid),
            command_pgrp,
            monitor_pgrp,
            pty_follower,
            errpipe_rx,
            backchannel,
            signal_stream,
            _signal_handlers: signal_handlers,
        })
    }

    /// Based on `mon_backchannel_cb`
    fn read_backchannel(&mut self, registry: &mut EventRegistry<Self>) {
        match self.backchannel.recv() {
            Err(err) => {
                // We can try later if receive is interrupted.
                if err.kind() != io::ErrorKind::Interrupted {
                    // There's something wrong with the backchannel, break the event loop.
                    dev_warn!("cannot read from backchannel: {err}");
                    registry.set_break(err);
                }
            }
            Ok(event) => {
                match event {
                    // We shouldn't receive this event at this point in the protocol
                    MonitorMessage::Edge => unreachable!(),
                    // Forward signal to the command.
                    MonitorMessage::Signal(signal) => {
                        if let Some(command_pid) = self.command_pid {
                            self.send_signal(signal, command_pid, true)
                        }
                    }
                }
            }
        }
    }

    fn read_errpipe(&mut self, registry: &mut EventRegistry<Self>) {
        match self.errpipe_rx.read() {
            Err(err) if was_interrupted(&err) => { /* Retry later */ }
            Err(err) => registry.set_break(err),
            Ok(error_code) => {
                // Received error code from the command, forward it to the parent.
                self.backchannel
                    .send(&ParentMessage::IoError(error_code))
                    .ok();
            }
        }
    }

    /// Send a signal to the command.
    fn send_signal(&self, signal: c_int, command_pid: ProcessId, from_parent: bool) {
        dev_info!(
            "sending {}{} to command",
            signal_fmt(signal),
            opt_fmt(from_parent, " from parent"),
        );
        // FIXME: We should call `killpg` instead of `kill`.
        match signal {
            SIGALRM => {
                terminate_process(command_pid, false);
            }
            SIGCONT_FG => {
                // Continue with the command as the foreground process group
                if let Err(err) = self.pty_follower.tcsetpgrp(self.command_pgrp) {
                    dev_error!(
                        "cannot set the foreground process group to {} (command): {err}",
                        self.command_pgrp
                    );
                }
                kill(command_pid, SIGCONT).ok();
            }
            SIGCONT_BG => {
                // Continue with the monitor as the foreground process group
                if let Err(err) = self.pty_follower.tcsetpgrp(self.monitor_pgrp) {
                    dev_error!(
                        "cannot set the foreground process group to {} (monitor): {err}",
                        self.monitor_pgrp
                    );
                }
                kill(command_pid, SIGCONT).ok();
            }
            signal => {
                // Send the signal to the command.
                kill(command_pid, signal).ok();
            }
        }
    }

    fn on_signal(&mut self, registry: &mut EventRegistry<Self>) {
        let info = match self.signal_stream.recv() {
            Ok(info) => info,
            Err(err) => {
                dev_error!("could not receive signal: {err}");
                return;
            }
        };

        dev_info!("monitor received{}", info);

        // Don't do anything if the command has terminated already
        let Some(command_pid) = self.command_pid else {
            dev_info!("command was terminated, ignoring signal");
            return;
        };

        match info.signal() {
            SIGCHLD => handle_sigchld(self, registry, "command", command_pid),
            signal => {
                if let Some(pid) = info.signaler_pid() {
                    if is_self_terminating(pid, command_pid, self.command_pgrp) {
                        // Skip the signal if it was sent by the user and it is self-terminating.
                        return;
                    }
                }

                self.send_signal(signal, command_pid, false)
            }
        }
    }
}

/// Decides if the signal sent by the process with `signaler_pid` PID is self-terminating.
///
/// A signal is self-terminating if `signaler_pid`:
/// - is the same PID of the command, or
/// - is in the process group of the command and the command is the leader.
fn is_self_terminating(
    signaler_pid: ProcessId,
    command_pid: ProcessId,
    command_pgrp: ProcessId,
) -> bool {
    if signaler_pid.is_valid() {
        if signaler_pid == command_pid {
            return true;
        }

        if let Ok(grp_leader) = getpgid(signaler_pid) {
            if grp_leader == command_pgrp {
                return true;
            }
        }
    }

    false
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum MonitorEvent {
    Signal,
    ReadableErrPipe,
    ReadableBackchannel,
}

impl Process for MonitorClosure<'_> {
    type Event = MonitorEvent;
    type Break = io::Error;
    type Exit = CommandStatus;

    fn on_event(&mut self, event: Self::Event, registry: &mut EventRegistry<Self>) {
        match event {
            MonitorEvent::Signal => self.on_signal(registry),
            MonitorEvent::ReadableErrPipe => self.read_errpipe(registry),
            MonitorEvent::ReadableBackchannel => self.read_backchannel(registry),
        }
    }
}

impl HandleSigchld for MonitorClosure<'_> {
    const OPTIONS: WaitOptions = WaitOptions::new().untraced().no_hang();

    fn on_exit(&mut self, exit_code: c_int, registry: &mut EventRegistry<Self>) {
        registry.set_exit(CommandStatus::Exit(exit_code));
        self.command_pid = None;
    }

    fn on_term(&mut self, signal: c_int, registry: &mut EventRegistry<Self>) {
        registry.set_exit(CommandStatus::Term(signal));
        self.command_pid = None;
    }

    fn on_stop(&mut self, signal: c_int, _registry: &mut EventRegistry<Self>) {
        // Save the foreground process group ID so we can restore it later.
        if let Ok(pgrp) = self.pty_follower.tcgetpgrp() {
            if pgrp != self.monitor_pgrp {
                self.command_pgrp = pgrp;
            }
        }
        self.backchannel
            .send(&CommandStatus::Stop(signal).into())
            .ok();
    }
}