Add procedure-scoped stateful leaves

src/leaf/remote_shell/session.rs

@@ -1,20 +1,30 @@
-use std::io::{self, Read};
-use std::process::{Child, ChildStdin, ExitStatus};
-use std::sync::mpsc::{self, Receiver, TryRecvError};
+use std::io::{self, Read, Write};
+use std::process::Command;
+use std::sync::mpsc::{self, Receiver, SyncSender, TryRecvError};
 use std::thread;
 
-use unshell::protocol::tree::OutgoingData;
+use portable_pty::{CommandBuilder, ExitStatus, PtySize, native_pty_system};
+use unshell::protocol::tree::{IncomingData, OutgoingData, ProcedureEffect};
+
+use unshell::Procedure;
 
 use super::errors::ShellLeafError;
 
-pub(super) struct ShellSession {
-    pub(super) child: Child,
-    pub(super) stdin: Option<ChildStdin>,
+/// Per-hook shell session created by the `open` procedure.
+///
+/// The procedure type is also the stored session type. This keeps the mapping
+/// between protocol procedure and hook state direct and easy to inspect.
+#[derive(Procedure)]
+#[procedure(leaf = RemoteShellLeaf, name = "open")]
+pub struct ProcedureOpen {
+    pub(super) child: Box<dyn portable_pty::Child + Send>,
+    process_group_leader: Option<u32>,
+    stdin_tx: Option<SyncSender<Vec<u8>>>,
     output_rx: Receiver<OutputEvent>,
     return_path: Vec<String>,
     hook_id: u64,
     procedure_id: String,
-    pub(super) readers_closed: usize,
+    output_closed: bool,
     pub(super) exit_status: Option<ExitStatus>,
     pub(super) local_end_sent: bool,
 }
@@ -24,53 +34,62 @@ enum OutputEvent {
     ReaderClosed,
 }
 
-impl ShellSession {
+use super::RemoteShellLeaf;
+
+impl ProcedureOpen {
     pub(super) fn spawn(
         return_path: Vec<String>,
         hook_id: u64,
         procedure_id: String,
     ) -> Result<Self, ShellLeafError> {
-        let mut command = if cfg!(windows) {
-            let mut command = std::process::Command::new("cmd.exe");
+        let pty_system = native_pty_system();
+        let pair = pty_system
+            .openpty(PtySize {
+                rows: 24,
+                cols: 80,
+                pixel_width: 0,
+                pixel_height: 0,
+            })
+            .map_err(|error| io::Error::other(error.to_string()))?;
+
+        let command = if cfg!(windows) {
+            let mut command = CommandBuilder::new("cmd.exe");
             command.arg("/Q");
             command
         } else {
-            let mut command = std::process::Command::new("/bin/sh");
+            let mut command = CommandBuilder::new("/bin/sh");
             command.arg("-i");
             command
         };
 
-        let mut child = command
-            .stdin(std::process::Stdio::piped())
-            .stdout(std::process::Stdio::piped())
-            .stderr(std::process::Stdio::piped())
-            .spawn()?;
+        let child = pair
+            .slave
+            .spawn_command(command)
+            .map_err(|error| io::Error::other(error.to_string()))?;
+        let process_group_leader = child.process_id();
+        let stdin = pair
+            .master
+            .take_writer()
+            .map_err(|error| io::Error::other(error.to_string()))?;
+        let stdout = pair
+            .master
+            .try_clone_reader()
+            .map_err(|error| io::Error::other(error.to_string()))?;
 
-        let stdin = child
-            .stdin
-            .take()
-            .ok_or_else(|| io::Error::other("failed to capture shell stdin"))?;
-        let stdout = child
-            .stdout
-            .take()
-            .ok_or_else(|| io::Error::other("failed to capture shell stdout"))?;
-        let stderr = child
-            .stderr
-            .take()
-            .ok_or_else(|| io::Error::other("failed to capture shell stderr"))?;
-
-        let (tx, rx) = mpsc::channel();
-        spawn_pipe_reader(stdout, tx.clone());
-        spawn_pipe_reader(stderr, tx);
+        let (stdin_tx, stdin_rx) = mpsc::sync_channel(64);
+        let (tx, rx) = mpsc::sync_channel(64);
+        spawn_pipe_writer(stdin, stdin_rx);
+        spawn_pipe_reader(stdout, tx);
 
         Ok(Self {
             child,
-            stdin: Some(stdin),
+            process_group_leader,
+            stdin_tx: Some(stdin_tx),
             output_rx: rx,
             return_path,
             hook_id,
             procedure_id,
-            readers_closed: 0,
+            output_closed: false,
             exit_status: None,
             local_end_sent: false,
         })
@@ -87,15 +106,22 @@ impl ShellSession {
     }
 
     pub(super) fn terminate(&mut self) -> Result<(), ShellLeafError> {
-        self.stdin.take();
+        self.stdin_tx.take();
         match self.child.try_wait()? {
             Some(status) => {
                 self.exit_status = Some(status);
                 Ok(())
             }
             None => {
-                self.child.kill()?;
-                self.exit_status = Some(self.child.wait()?);
+                self.kill_process_group();
+                self.child
+                    .kill()
+                    .map_err(|error| io::Error::other(error.to_string()))?;
+                self.exit_status = Some(
+                    self.child
+                        .wait()
+                        .map_err(|error| io::Error::other(error.to_string()))?,
+                );
                 Ok(())
             }
         }
@@ -105,30 +131,113 @@ impl ShellSession {
         loop {
             match self.output_rx.try_recv() {
                 Ok(OutputEvent::Chunk(bytes)) => outgoing.push(self.packet(bytes, false)),
-                Ok(OutputEvent::ReaderClosed) => self.readers_closed += 1,
+                Ok(OutputEvent::ReaderClosed) => self.output_closed = true,
                 Err(TryRecvError::Empty) => break,
                 Err(TryRecvError::Disconnected) => {
-                    self.readers_closed = 2;
+                    self.output_closed = true;
                     break;
                 }
             }
         }
     }
-}
 
-pub(super) fn close_session(
-    mut session: ShellSession,
-) -> Result<Vec<OutgoingData>, ShellLeafError> {
-    session.terminate()?;
-    if session.local_end_sent {
-        return Ok(Vec::new());
+    /// Applies one inbound hook payload to the shell process.
+    pub(super) fn on_data(
+        &mut self,
+        data: IncomingData,
+    ) -> Result<ProcedureEffect, ShellLeafError> {
+        if !data.message.data.is_empty() {
+            let Some(stdin_tx) = self.stdin_tx.as_ref() else {
+                return Ok(ProcedureEffect::default());
+            };
+            stdin_tx.try_send(data.message.data).map_err(|_| {
+                io::Error::new(io::ErrorKind::WouldBlock, "shell stdin channel full")
+            })?;
+        }
+
+        if !data.message.end_hook {
+            return Ok(ProcedureEffect::default());
+        }
+
+        // Peer end means no more stdin from the caller. Keep the process alive so
+        // any buffered PTY output can drain through the normal poll path. On Unix
+        // we also send SIGHUP so an interactive shell treats this like terminal
+        // hangup instead of waiting forever on the still-open PTY master.
+        self.stdin_tx.take();
+        self.signal_peer_end();
+        Ok(ProcedureEffect::default())
     }
 
-    session.local_end_sent = true;
-    Ok(vec![session.packet(Vec::new(), true)])
+    /// Polls the shell for locally-generated output.
+    pub(super) fn poll(&mut self) -> Result<ProcedureEffect, ShellLeafError> {
+        let mut outgoing = Vec::new();
+        self.drain_output(&mut outgoing);
+
+        if self.local_end_sent {
+            return Ok(ProcedureEffect::outgoing(outgoing));
+        }
+
+        if self.exit_status.is_none() {
+            self.exit_status = self
+                .child
+                .try_wait()
+                .map_err(|error| io::Error::other(error.to_string()))?;
+        }
+
+        if self.exit_status.is_some() && !self.output_closed {
+            self.kill_process_group();
+        }
+
+        if self.exit_status.is_some() && self.output_closed {
+            outgoing.push(self.packet(Vec::new(), true));
+            self.local_end_sent = true;
+            return Ok(ProcedureEffect::close(outgoing));
+        }
+
+        Ok(ProcedureEffect::outgoing(outgoing))
+    }
+
+    fn kill_process_group(&self) {
+        #[cfg(unix)]
+        if let Some(process_group_leader) = self.process_group_leader {
+            let _ = Command::new("kill")
+                .arg("-KILL")
+                .arg(format!("-{}", process_group_leader))
+                .status();
+        }
+    }
+
+    fn signal_peer_end(&self) {
+        #[cfg(unix)]
+        if let Some(process_group_leader) = self.process_group_leader {
+            let _ = Command::new("kill")
+                .arg("-HUP")
+                .arg(format!("-{}", process_group_leader))
+                .status();
+        }
+    }
 }
 
-fn spawn_pipe_reader<R>(mut reader: R, tx: mpsc::Sender<OutputEvent>)
+impl Drop for ProcedureOpen {
+    fn drop(&mut self) {
+        let _ = self.terminate();
+    }
+}
+
+fn spawn_pipe_writer(mut stdin: Box<dyn Write + Send>, rx: Receiver<Vec<u8>>) {
+    thread::spawn(move || {
+        for bytes in rx {
+            if stdin.write_all(&bytes).is_err() {
+                break;
+            }
+            if stdin.flush().is_err() {
+                break;
+            }
+        }
+    });
+}
+
+fn spawn_pipe_reader<R>(mut reader: R, tx: mpsc::SyncSender<OutputEvent>)
 where
     R: Read + Send + 'static,
 {