feat(runtime): add interactive local PTY sessions

rabby-runtime/Cargo.toml

@@ -16,3 +16,4 @@ base64 = "0.22"
 rand = "0.8"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
+portable-pty = "0.8"

rabby-runtime/src/lib.rs

@@ -2,4 +2,5 @@ pub mod config_store;
 pub mod local_terminal;
 pub mod pty;
 pub mod ssh_runtime;
+pub mod terminal_session;
 pub mod wallet_vault;

rabby-runtime/src/terminal_session/mod.rs

@@ -0,0 +1,243 @@
+use portable_pty::{native_pty_system, Child, CommandBuilder, MasterPty, PtySize};
+use serde::Serialize;
+use std::collections::HashMap;
+use std::fmt;
+use std::io::{Read, Write};
+use std::time::{Duration, Instant};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct TerminalSessionId(pub u64);
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
+pub struct TerminalSessionSize {
+    pub cols: u16,
+    pub rows: u16,
+}
+
+#[derive(Debug, PartialEq, Eq)]
+pub enum TerminalSessionError {
+    InvalidSize,
+    MissingSession(u64),
+    Io(String),
+    Timeout(String),
+}
+
+impl fmt::Display for TerminalSessionError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::InvalidSize => write!(f, "invalid terminal session size"),
+            Self::MissingSession(id) => write!(f, "terminal session {id} does not exist"),
+            Self::Io(message) => write!(f, "terminal session I/O error: {message}"),
+            Self::Timeout(needle) => write!(f, "timed out waiting for terminal output: {needle}"),
+        }
+    }
+}
+
+impl std::error::Error for TerminalSessionError {}
+
+pub struct TerminalSessionRegistry {
+    next_id: u64,
+    sessions: HashMap<TerminalSessionId, TerminalSession>,
+}
+
+struct TerminalSession {
+    master: Box<dyn MasterPty + Send>,
+    child: Box<dyn Child + Send + Sync>,
+    reader: Box<dyn Read + Send>,
+    writer: Box<dyn Write + Send>,
+    size: TerminalSessionSize,
+}
+
+impl Default for TerminalSessionRegistry {
+    fn default() -> Self {
+        Self {
+            next_id: 1,
+            sessions: HashMap::new(),
+        }
+    }
+}
+
+impl TerminalSessionRegistry {
+    pub fn spawn_local(
+        &mut self,
+        shell: impl AsRef<str>,
+        cols: u16,
+        rows: u16,
+    ) -> Result<TerminalSessionId, TerminalSessionError> {
+        let size = session_size(cols, rows)?;
+        let pty_system = native_pty_system();
+        let pair = pty_system
+            .openpty(PtySize {
+                rows,
+                cols,
+                pixel_width: 0,
+                pixel_height: 0,
+            })
+            .map_err(|err| TerminalSessionError::Io(err.to_string()))?;
+        let command = CommandBuilder::new(shell.as_ref());
+        let child = pair
+            .slave
+            .spawn_command(command)
+            .map_err(|err| TerminalSessionError::Io(err.to_string()))?;
+        let reader = pair
+            .master
+            .try_clone_reader()
+            .map_err(|err| TerminalSessionError::Io(err.to_string()))?;
+        let writer = pair
+            .master
+            .take_writer()
+            .map_err(|err| TerminalSessionError::Io(err.to_string()))?;
+        let id = TerminalSessionId(self.next_id);
+        self.next_id += 1;
+        self.sessions.insert(
+            id,
+            TerminalSession {
+                master: pair.master,
+                child,
+                reader,
+                writer,
+                size,
+            },
+        );
+        Ok(id)
+    }
+
+    pub fn write(
+        &mut self,
+        id: TerminalSessionId,
+        bytes: &[u8],
+    ) -> Result<(), TerminalSessionError> {
+        let session = self.session_mut(id)?;
+        session
+            .writer
+            .write_all(bytes)
+            .map_err(|err| TerminalSessionError::Io(err.to_string()))?;
+        session
+            .writer
+            .flush()
+            .map_err(|err| TerminalSessionError::Io(err.to_string()))
+    }
+
+    pub fn read_until(
+        &mut self,
+        id: TerminalSessionId,
+        needle: &str,
+        timeout: Duration,
+    ) -> Result<String, TerminalSessionError> {
+        let deadline = Instant::now() + timeout;
+        let mut output = String::new();
+        let mut buf = [0u8; 4096];
+        while Instant::now() < deadline {
+            let session = self.session_mut(id)?;
+            match session.reader.read(&mut buf) {
+                Ok(0) => std::thread::sleep(Duration::from_millis(10)),
+                Ok(n) => {
+                    output.push_str(&String::from_utf8_lossy(&buf[..n]));
+                    if output.contains(needle) {
+                        return Ok(output);
+                    }
+                }
+                Err(err) if err.kind() == std::io::ErrorKind::Interrupted => continue,
+                Err(err) => return Err(TerminalSessionError::Io(err.to_string())),
+            }
+        }
+        Err(TerminalSessionError::Timeout(needle.to_string()))
+    }
+
+    pub fn resize(
+        &mut self,
+        id: TerminalSessionId,
+        cols: u16,
+        rows: u16,
+    ) -> Result<(), TerminalSessionError> {
+        let size = session_size(cols, rows)?;
+        let session = self.session_mut(id)?;
+        session
+            .master
+            .resize(PtySize {
+                rows,
+                cols,
+                pixel_width: 0,
+                pixel_height: 0,
+            })
+            .map_err(|err| TerminalSessionError::Io(err.to_string()))?;
+        session.size = size;
+        Ok(())
+    }
+
+    pub fn size(&self, id: TerminalSessionId) -> Result<TerminalSessionSize, TerminalSessionError> {
+        self.sessions
+            .get(&id)
+            .map(|session| session.size)
+            .ok_or(TerminalSessionError::MissingSession(id.0))
+    }
+
+    pub fn kill(&mut self, id: TerminalSessionId) -> Result<(), TerminalSessionError> {
+        let mut session = self
+            .sessions
+            .remove(&id)
+            .ok_or(TerminalSessionError::MissingSession(id.0))?;
+        session
+            .child
+            .kill()
+            .map_err(|err| TerminalSessionError::Io(err.to_string()))
+    }
+
+    fn session_mut(
+        &mut self,
+        id: TerminalSessionId,
+    ) -> Result<&mut TerminalSession, TerminalSessionError> {
+        self.sessions
+            .get_mut(&id)
+            .ok_or(TerminalSessionError::MissingSession(id.0))
+    }
+}
+
+fn session_size(cols: u16, rows: u16) -> Result<TerminalSessionSize, TerminalSessionError> {
+    if cols == 0 || rows == 0 {
+        return Err(TerminalSessionError::InvalidSize);
+    }
+    Ok(TerminalSessionSize { cols, rows })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::time::Duration;
+
+    #[test]
+    fn session_registry_spawns_shell_and_reads_output() {
+        let mut registry = TerminalSessionRegistry::default();
+        let id = registry.spawn_local("/bin/sh", 80, 24).unwrap();
+
+        registry
+            .write(id, b"printf rabby-pty-ready\nexit\n")
+            .unwrap();
+        let output = registry
+            .read_until(id, "rabby-pty-ready", Duration::from_secs(3))
+            .unwrap();
+
+        assert!(output.contains("rabby-pty-ready"));
+    }
+
+    #[test]
+    fn session_registry_resizes_running_session() {
+        let mut registry = TerminalSessionRegistry::default();
+        let id = registry.spawn_local("/bin/sh", 80, 24).unwrap();
+
+        registry.resize(id, 120, 40).unwrap();
+        let size = registry.size(id).unwrap();
+
+        assert_eq!(size.cols, 120);
+        assert_eq!(size.rows, 40);
+        registry.kill(id).unwrap();
+    }
+
+    #[test]
+    fn session_registry_rejects_missing_session() {
+        let mut registry = TerminalSessionRegistry::default();
+        let err = registry.write(TerminalSessionId(999), b"x").unwrap_err();
+
+        assert_eq!(err, TerminalSessionError::MissingSession(999));
+    }
+}