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

use std::fs::{File, OpenOptions};
use std::io::{self, Read, Write};
use std::num::NonZeroU16;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
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 rustix::fs::{OFlags, fcntl_getfl, fcntl_setfl};
use rustix::termios::{LocalModes, OptionalActions, Termios, tcgetattr, tcgetwinsize, tcsetattr};
use signal_hook::SigId;
use signal_hook::consts::signal::{SIGHUP, SIGINT, SIGQUIT, SIGTERM, SIGTSTP, SIGWINCH};

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<DeveloperDecisionKind>;
    /// # Errors
    /// Returns an error restoring, stopping, or re-entering raw mode.
    fn suspend_and_resume(&mut self) -> io::Result<()>;
}

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

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: SignalFlags::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 self.signals.terminate.swap(false, Ordering::SeqCst) {
            return Ok(HostTerminalEvent::Terminate);
        }
        if self.signals.suspend.swap(false, Ordering::SeqCst) {
            return Ok(HostTerminalEvent::LocalSuspend);
        }
        if self.signals.resize.swap(false, Ordering::SeqCst) {
            return Ok(HostTerminalEvent::Resize);
        }
        if self.signals.interrupt.swap(false, Ordering::SeqCst) {
            return Ok(HostTerminalEvent::Signal(TerminalSignalKind::Interrupt));
        }
        if self.signals.quit.swap(false, Ordering::SeqCst) {
            return Ok(HostTerminalEvent::Signal(TerminalSignalKind::Quit));
        }
        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<DeveloperDecisionKind> {
        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(),
        )?;
        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(if approve {
                        DeveloperDecisionKind::Approve
                    } else {
                        DeveloperDecisionKind::Deny
                    });
                }
                HostTerminalEvent::LocalSuspend => self.suspend_and_resume()?,
                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();
    }
}

struct SignalFlags {
    interrupt: Arc<AtomicBool>,
    quit: Arc<AtomicBool>,
    suspend: Arc<AtomicBool>,
    resize: Arc<AtomicBool>,
    terminate: Arc<AtomicBool>,
    registrations: Vec<SigId>,
}

impl SignalFlags {
    fn register() -> io::Result<Self> {
        let interrupt = Arc::new(AtomicBool::new(false));
        let quit = Arc::new(AtomicBool::new(false));
        let suspend = Arc::new(AtomicBool::new(false));
        let resize = Arc::new(AtomicBool::new(false));
        let terminate = Arc::new(AtomicBool::new(false));
        let registrations = vec![
            signal_hook::flag::register(SIGINT, Arc::clone(&interrupt))?,
            signal_hook::flag::register(SIGQUIT, Arc::clone(&quit))?,
            signal_hook::flag::register(SIGTSTP, Arc::clone(&suspend))?,
            signal_hook::flag::register(SIGWINCH, Arc::clone(&resize))?,
            signal_hook::flag::register(SIGTERM, Arc::clone(&terminate))?,
            signal_hook::flag::register(SIGHUP, Arc::clone(&terminate))?,
        ];
        Ok(Self {
            interrupt,
            quit,
            suspend,
            resize,
            terminate,
            registrations,
        })
    }
}

impl Drop for SignalFlags {
    fn drop(&mut self) {
        for registration in self.registrations.drain(..) {
            signal_hook::low_level::unregister(registration);
        }
    }
}

/// 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 decision = terminal.prompt(&prompt)?;
                let signed_decision =
                    identity.sign_developer_decision(binding.clone(), &prompt, 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
    })
}