Add interactive client terminal transport

crates/lazaret-client/src/terminal.rs

@@ -0,0 +1,389 @@
+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
+    })
+}