lazaret/crates/lazaret-client/src/terminal.rs

use std::fs::{File, OpenOptions};
use std::io::{self, Read, Write};
use std::num::NonZeroU16;
use std::time::Duration;

use lazaret_core::records::{
    DeveloperDecisionKind, InteractiveSessionBatch, InteractiveSessionEvent,
    InteractiveSessionMessage, InteractiveSessionRequest, ProtocolVersion, SessionBinding,
    SessionClose, SessionExitOutcome, SessionTerminalEvent, SessionTerminalMessage, TerminalInput,
    TerminalPrompt, TerminalResize, TerminalSignal, TerminalSignalKind,
};
use nix::sys::signal::{SigSet, SigmaskHow, Signal};
use nix::sys::signalfd::{SfdFlags, SignalFd};
use rustix::fs::{OFlags, fcntl_getfl, fcntl_setfl};
use rustix::termios::{LocalModes, OptionalActions, Termios, tcgetattr, tcgetwinsize, tcsetattr};

use crate::{ClientError, ClientSessionIdentity};

pub trait InteractiveSessionTransport {
    type Error;

    /// # Errors
    /// Returns an authenticated session or framing failure.
    fn start(
        &mut self,
        request: InteractiveSessionRequest,
    ) -> Result<InteractiveSessionBatch, ClientError<Self::Error>>;

    /// # Errors
    /// Returns a lost connection, stale session, or framing failure.
    fn exchange(
        &mut self,
        event: InteractiveSessionEvent,
    ) -> Result<InteractiveSessionBatch, ClientError<Self::Error>>;
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum HostTerminalEvent {
    Input(Vec<u8>),
    Signal(TerminalSignalKind),
    Resize,
    LocalSuspend,
    Terminate,
    Idle,
}

pub trait HostTerminal {
    /// # Errors
    /// Returns an error when raw mode cannot be entered.
    fn enter_raw(&mut self) -> io::Result<()>;
    /// # Errors
    /// Returns an error when the original terminal mode cannot be restored.
    fn restore(&mut self) -> io::Result<()>;
    /// # Errors
    /// Returns an error when the current terminal size cannot be observed.
    fn size(&self, session_id: lazaret_core::identity::SessionId) -> io::Result<TerminalResize>;
    /// # Errors
    /// Returns an error while observing terminal input or signals.
    fn next_event(&mut self) -> io::Result<HostTerminalEvent>;
    /// # Errors
    /// Returns an error while writing merged terminal output.
    fn write_output(&mut self, bytes: &[u8]) -> io::Result<()>;
    /// # Errors
    /// Returns an error while collecting one explicit developer decision.
    fn prompt(&mut self, prompt: &TerminalPrompt) -> io::Result<PromptOutcome>;
    /// # Errors
    /// Returns an error restoring, stopping, or re-entering raw mode.
    fn suspend_and_resume(&mut self) -> io::Result<()>;
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PromptOutcome {
    pub decision: DeveloperDecisionKind,
    pub resumed_size: Option<TerminalResize>,
}

pub struct SystemTerminal {
    terminal: File,
    original: Termios,
    original_flags: OFlags,
    raw: bool,
    signals: SignalQueue,
}

impl SystemTerminal {
    /// Opens the controlling terminal and captures its exact initial mode.
    ///
    /// # Errors
    /// Returns an error when the process has no usable controlling terminal.
    pub fn open() -> io::Result<Self> {
        let terminal = OpenOptions::new().read(true).write(true).open("/dev/tty")?;
        Ok(Self {
            original: tcgetattr(&terminal)?,
            original_flags: fcntl_getfl(&terminal)?,
            terminal,
            raw: false,
            signals: SignalQueue::register()?,
        })
    }
}

impl HostTerminal for SystemTerminal {
    fn enter_raw(&mut self) -> io::Result<()> {
        if self.raw {
            return Ok(());
        }
        let mut raw = self.original.clone();
        raw.make_raw();
        raw.local_modes.insert(LocalModes::ISIG);
        tcsetattr(&self.terminal, OptionalActions::Flush, &raw)?;
        if let Err(error) = fcntl_setfl(&self.terminal, self.original_flags | OFlags::NONBLOCK) {
            let _ = tcsetattr(&self.terminal, OptionalActions::Flush, &self.original);
            return Err(error.into());
        }
        self.raw = true;
        Ok(())
    }

    fn restore(&mut self) -> io::Result<()> {
        if !self.raw {
            return Ok(());
        }
        fcntl_setfl(&self.terminal, self.original_flags)?;
        tcsetattr(&self.terminal, OptionalActions::Flush, &self.original)?;
        self.raw = false;
        Ok(())
    }

    fn size(&self, session_id: lazaret_core::identity::SessionId) -> io::Result<TerminalResize> {
        let size = tcgetwinsize(&self.terminal)?;
        Ok(TerminalResize {
            session_id,
            rows: NonZeroU16::new(size.ws_row)
                .ok_or_else(|| io::Error::other("terminal reported zero rows"))?,
            columns: NonZeroU16::new(size.ws_col)
                .ok_or_else(|| io::Error::other("terminal reported zero columns"))?,
            pixel_width: size.ws_xpixel,
            pixel_height: size.ws_ypixel,
        })
    }

    fn next_event(&mut self) -> io::Result<HostTerminalEvent> {
        if let Some(event) = self.signals.next_event()? {
            return Ok(event);
        }
        let mut bytes = vec![0_u8; 16 * 1024];
        match self.terminal.read(&mut bytes) {
            Ok(0) => Ok(HostTerminalEvent::Terminate),
            Ok(length) => {
                bytes.truncate(length);
                Ok(HostTerminalEvent::Input(bytes))
            }
            Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
                std::thread::sleep(Duration::from_millis(5));
                Ok(HostTerminalEvent::Idle)
            }
            Err(error) => Err(error),
        }
    }

    fn write_output(&mut self, bytes: &[u8]) -> io::Result<()> {
        self.terminal.write_all(bytes)?;
        self.terminal.flush()
    }

    fn prompt(&mut self, prompt: &TerminalPrompt) -> io::Result<PromptOutcome> {
        self.write_output(
            format!(
                "\r\nLazaret host request\r\nprofile: {}\r\nsession: {}\r\ncommand: {}\r\njustification: {}\r\nApprove? [y/N] ",
                prompt.agent_profile,
                hex(prompt.session_id.as_bytes()),
                display_argv(&prompt.command_argv),
                prompt.justification
            )
            .as_bytes(),
        )?;
        let mut resumed_size = None;
        loop {
            match self.next_event()? {
                HostTerminalEvent::Input(bytes) => {
                    let approve = bytes
                        .first()
                        .is_some_and(|byte| matches!(byte, b'y' | b'Y'));
                    self.write_output(if approve { b"y\r\n" } else { b"n\r\n" })?;
                    return Ok(PromptOutcome {
                        decision: if approve {
                            DeveloperDecisionKind::Approve
                        } else {
                            DeveloperDecisionKind::Deny
                        },
                        resumed_size,
                    });
                }
                HostTerminalEvent::LocalSuspend => {
                    self.suspend_and_resume()?;
                    resumed_size = Some(self.size(prompt.session_id)?);
                }
                HostTerminalEvent::Terminate => return Err(io::ErrorKind::Interrupted.into()),
                HostTerminalEvent::Idle
                | HostTerminalEvent::Resize
                | HostTerminalEvent::Signal(_) => {}
            }
        }
    }

    fn suspend_and_resume(&mut self) -> io::Result<()> {
        self.restore()?;
        lazaret_linux::process::suspend_current_process()?;
        self.enter_raw()
    }
}

impl Drop for SystemTerminal {
    fn drop(&mut self) {
        let _ = self.restore();
        self.signals.restore();
    }
}

struct SignalQueue {
    fd: SignalFd,
    original_mask: SigSet,
    restored: bool,
}

impl SignalQueue {
    fn register() -> io::Result<Self> {
        let mut mask = SigSet::empty();
        for signal in [
            Signal::SIGINT,
            Signal::SIGQUIT,
            Signal::SIGTSTP,
            Signal::SIGWINCH,
            Signal::SIGTERM,
            Signal::SIGHUP,
        ] {
            mask.add(signal);
        }
        let original_mask = mask
            .thread_swap_mask(SigmaskHow::SIG_BLOCK)
            .map_err(errno_to_io)?;
        let fd = match SignalFd::with_flags(&mask, SfdFlags::SFD_NONBLOCK | SfdFlags::SFD_CLOEXEC) {
            Ok(fd) => fd,
            Err(error) => {
                let _ = original_mask.thread_swap_mask(SigmaskHow::SIG_SETMASK);
                return Err(errno_to_io(error));
            }
        };
        Ok(Self {
            fd,
            original_mask,
            restored: false,
        })
    }

    fn next_event(&mut self) -> io::Result<Option<HostTerminalEvent>> {
        let Some(info) = self.fd.read_signal().map_err(errno_to_io)? else {
            return Ok(None);
        };
        let signal = Signal::try_from(info.ssi_signo.cast_signed()).map_err(errno_to_io)?;
        Ok(match signal {
            Signal::SIGINT => Some(HostTerminalEvent::Signal(TerminalSignalKind::Interrupt)),
            Signal::SIGQUIT => Some(HostTerminalEvent::Signal(TerminalSignalKind::Quit)),
            Signal::SIGTSTP => Some(HostTerminalEvent::LocalSuspend),
            Signal::SIGWINCH => Some(HostTerminalEvent::Resize),
            Signal::SIGTERM | Signal::SIGHUP => Some(HostTerminalEvent::Terminate),
            _ => None,
        })
    }

    fn restore(&mut self) {
        if !self.restored {
            let _ = self.original_mask.thread_swap_mask(SigmaskHow::SIG_SETMASK);
            self.restored = true;
        }
    }
}

impl Drop for SignalQueue {
    fn drop(&mut self) {
        self.restore();
    }
}

fn errno_to_io(error: nix::errno::Errno) -> io::Error {
    io::Error::from_raw_os_error(error as i32)
}

/// Runs one interactive terminal route until its exact Session exit.
///
/// # Errors
/// Returns terminal, prompt-signing, stale-session, or lost-connection failures.
pub fn run_interactive_session<T: InteractiveSessionTransport, H: HostTerminal>(
    transport: &mut T,
    terminal: &mut H,
    identity: &ClientSessionIdentity,
    binding: &SessionBinding,
    argv: &[Vec<u8>],
) -> Result<i32, ClientError<T::Error>> {
    terminal.enter_raw()?;
    let result = (|| {
        let session_id = identity.session_id();
        let initial_size = terminal.size(session_id)?;
        let mut batch = transport.start(InteractiveSessionRequest {
            protocol_version: ProtocolVersion::V1,
            session_id,
            argv: argv.to_vec(),
            initial_size,
        })?;
        loop {
            if let Some(exit) = handle_batch(transport, terminal, identity, binding, batch)? {
                return match exit.outcome() {
                    SessionExitOutcome::Exited => exit
                        .exit_code()
                        .ok_or(ClientError::Protocol("Exited omitted its exit code")),
                    SessionExitOutcome::Terminated | SessionExitOutcome::RuntimeLost => {
                        Err(ClientError::LostConnection)
                    }
                };
            }
            let event =
                match terminal.next_event()? {
                    HostTerminalEvent::Input(bytes) => InteractiveSessionEvent::Terminal(
                        SessionTerminalEvent::Input(TerminalInput { session_id, bytes }),
                    ),
                    HostTerminalEvent::Signal(signal) => InteractiveSessionEvent::Terminal(
                        SessionTerminalEvent::Signal(TerminalSignal { session_id, signal }),
                    ),
                    HostTerminalEvent::Resize => InteractiveSessionEvent::Terminal(
                        SessionTerminalEvent::Resize(terminal.size(session_id)?),
                    ),
                    HostTerminalEvent::LocalSuspend => {
                        terminal.suspend_and_resume()?;
                        InteractiveSessionEvent::Terminal(SessionTerminalEvent::Resize(
                            terminal.size(session_id)?,
                        ))
                    }
                    HostTerminalEvent::Terminate => InteractiveSessionEvent::Terminal(
                        SessionTerminalEvent::Close(SessionClose { session_id }),
                    ),
                    HostTerminalEvent::Idle => InteractiveSessionEvent::Poll { session_id },
                };
            batch = transport.exchange(event)?;
        }
    })();
    let restoration = terminal.restore().map_err(ClientError::Io);
    restoration?;
    result
}

fn handle_batch<T: InteractiveSessionTransport, H: HostTerminal>(
    transport: &mut T,
    terminal: &mut H,
    identity: &ClientSessionIdentity,
    binding: &SessionBinding,
    batch: InteractiveSessionBatch,
) -> Result<Option<lazaret_core::records::SessionExit>, ClientError<T::Error>> {
    if batch.session_id != identity.session_id() {
        return Err(ClientError::Protocol("response Session ID differs"));
    }
    for message in batch.messages {
        match message {
            InteractiveSessionMessage::Terminal(SessionTerminalMessage::Output(output)) => {
                require_session(output.session_id, identity.session_id())?;
                terminal.write_output(&output.bytes)?;
            }
            InteractiveSessionMessage::Terminal(SessionTerminalMessage::Exit(exit)) => {
                require_session(exit.session_id(), identity.session_id())?;
                return Ok(Some(exit));
            }
            InteractiveSessionMessage::Prompt(prompt) => {
                require_session(prompt.session_id, identity.session_id())?;
                let outcome = terminal.prompt(&prompt)?;
                if let Some(resize) = outcome.resumed_size {
                    let response = transport.exchange(InteractiveSessionEvent::Terminal(
                        SessionTerminalEvent::Resize(resize),
                    ))?;
                    if let Some(exit) =
                        handle_batch(transport, terminal, identity, binding, response)?
                    {
                        return Ok(Some(exit));
                    }
                }
                let signed_decision =
                    identity.sign_developer_decision(binding.clone(), &prompt, outcome.decision)?;
                let response = transport.exchange(InteractiveSessionEvent::PromptDecision {
                    session_id: identity.session_id(),
                    signed_decision,
                })?;
                if let Some(exit) = handle_batch(transport, terminal, identity, binding, response)?
                {
                    return Ok(Some(exit));
                }
            }
        }
    }
    Ok(None)
}

fn require_session<E>(
    actual: lazaret_core::identity::SessionId,
    expected: lazaret_core::identity::SessionId,
) -> Result<(), ClientError<E>> {
    if actual == expected {
        Ok(())
    } else {
        Err(ClientError::Protocol("response Session ID differs"))
    }
}

fn display_argv(argv: &[Vec<u8>]) -> String {
    argv.iter()
        .map(|argument| String::from_utf8_lossy(argument))
        .collect::<Vec<_>>()
        .join(" ")
}

fn hex(bytes: &[u8]) -> String {
    use std::fmt::Write as _;

    bytes.iter().fold(String::new(), |mut output, byte| {
        write!(output, "{byte:02x}").expect("writing to a String cannot fail");
        output
    })
}

#[cfg(test)]
#[allow(clippy::disallowed_methods, clippy::disallowed_types)]
mod tests {
    use std::os::unix::process::ExitStatusExt as _;
    use std::process::Command;

    use nix::sys::signal::{SigSet, SigmaskHow, Signal, kill};
    use nix::unistd::Pid;

    use super::SignalQueue;

    #[test]
    fn signal_scope_restores_mask_and_default_termination() {
        if std::env::var_os("LAZARET_SIGNAL_SCOPE_CHILD").is_some() {
            let signals = SignalQueue::register().unwrap();
            drop(signals);
            kill(Pid::this(), Signal::SIGTERM).unwrap();
            panic!("restored SIGTERM did not terminate the process");
        }

        let observed = SigSet::empty()
            .thread_swap_mask(SigmaskHow::SIG_BLOCK)
            .unwrap();
        let signals = SignalQueue::register().unwrap();
        drop(signals);
        let restored = SigSet::empty()
            .thread_swap_mask(SigmaskHow::SIG_BLOCK)
            .unwrap();
        assert_eq!(restored, observed);

        let status = Command::new(std::env::current_exe().unwrap())
            .arg("--exact")
            .arg("terminal::tests::signal_scope_restores_mask_and_default_termination")
            .arg("--nocapture")
            .env("LAZARET_SIGNAL_SCOPE_CHILD", "1")
            .status()
            .unwrap();
        assert_eq!(status.signal(), Some(nix::libc::SIGTERM));
    }
}