Add agent pty support

twizmwazin · Apr 12, 2024 · f6b5517 · f6b5517
1 parent 8690150
commit f6b5517
Show file tree

Hide file tree

Showing 8 changed files with 49 additions and 51 deletions.
diff --git a/.gitmodules b/.gitmodules
@@ -0,0 +1,3 @@
+[submodule "crates/subprocess"]
+	path = crates/subprocess
+	url = [email protected]:twizmwazin/rust-subprocess.git
diff --git a/crates/bh_agent_client/src/bindings.rs b/crates/bh_agent_client/src/bindings.rs
@@ -98,9 +98,10 @@ impl BhAgentClient {
         setuid: Option<u32>,
         setgid: Option<u32>,
         setpgid: Option<bool>,
+        use_pty: Option<bool>,
     ) -> PyResult<ProcessId> {
         debug!(
-            "Running process with argv {:?}, stdin {}, stdout {}, stderr {}, executable {:?}, env {:?}, cwd {:?}, setuid {:?}, setgid {:?}, setpgid {:?}",
+            "Running process with argv {:?}, stdin {}, stdout {}, stderr {}, executable {:?}, env {:?}, cwd {:?}, setuid {:?}, setgid {:?}, setpgid {:?}, use_pty {:?}",
             argv,
             stdin,
             stdout,
@@ -110,7 +111,9 @@ impl BhAgentClient {
             cwd,
             setuid,
             setgid,
-            setpgid,);
+            setpgid,
+            use_pty,
+        );
 
         let config = RemotePOpenConfig {
             argv,
@@ -132,6 +135,7 @@ impl BhAgentClient {
             setuid,
             setgid,
             setpgid: setpgid.unwrap_or(false),
+            use_pty: use_pty.unwrap_or(false),
         };
         run_in_runtime(
             self,

diff --git a/crates/bh_agent_common/src/types.rs b/crates/bh_agent_common/src/types.rs
@@ -32,6 +32,7 @@ pub struct RemotePOpenConfig {
     pub setuid: Option<u32>,
     pub setgid: Option<u32>,
     pub setpgid: bool,
+    pub use_pty: bool,
 }
 
 #[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)]

diff --git a/crates/bh_agent_server/Cargo.toml b/crates/bh_agent_server/Cargo.toml
@@ -8,7 +8,7 @@ edition = "2021"
 [dependencies]
 anyhow = "1.0.75"
 bh_agent_common = { path = "../bh_agent_common" }
-subprocess = "0.2.9"
+subprocess = { path = "../subprocess" }
 tarpc = { version = "0.34.0", features = ["full"] }
 tokio = { version = "1.32.0", features = ["rt-multi-thread"] }
 futures = "0.3.28"

diff --git a/crates/bh_agent_server/src/state.rs b/crates/bh_agent_server/src/state.rs
@@ -3,6 +3,7 @@ use log::{debug, trace};
 use std::collections::HashMap;
 use std::ffi::OsString;
 use std::fs::{File, OpenOptions};
+use std::rc::Rc;
 use std::sync::{Arc, RwLock};
 use std::thread::sleep;
 use std::time::Duration;
@@ -116,19 +117,15 @@ impl BhAgentState {
     }
 
     pub fn run_command(&self, config: RemotePOpenConfig) -> Result<ProcessId, AgentError> {
+        let select_io = |redirection: &Redirection| match (redirection, &config.use_pty) {
+            (Redirection::None, _) => subprocess::Redirection::None,
+            (Redirection::Save, false) => subprocess::Redirection::Pipe,
+            (Redirection::Save, true) => subprocess::Redirection::Pty,
+        };
         let mut popenconfig = PopenConfig {
-            stdin: match config.stdin {
-                Redirection::None => subprocess::Redirection::None,
-                Redirection::Save => subprocess::Redirection::Pipe,
-            },
-            stdout: match config.stdout {
-                Redirection::None => subprocess::Redirection::None,
-                Redirection::Save => subprocess::Redirection::Pipe,
-            },
-            stderr: match config.stderr {
-                Redirection::None => subprocess::Redirection::None,
-                Redirection::Save => subprocess::Redirection::Pipe,
-            },
+            stdin: select_io(&config.stdin),
+            stdout: select_io(&config.stdout),
+            stderr: select_io(&config.stderr),
             detached: false,
             executable: config
                 .executable
@@ -160,42 +157,32 @@ impl BhAgentState {
 
         let proc_id = self.take_proc_id()?;
 
+        let save_io = |io: &Option<File>,
+                       io_map: &RwLock<BiMap<ProcessId, FileId>>,
+                       name: &str,
+                       write: bool|
+         -> Result<(), AgentError> {
+            if io.is_some() {
+                trace!("Saving {} for process {}", name, proc_id);
+                let file_id = self.take_file_id()?;
+                io_map.write()?.insert(proc_id, file_id);
+                if write {
+                    self.file_modes
+                        .write()?
+                        .insert(file_id, FileOpenMode::Write);
+                }
+                self.file_types
+                    .write()?
+                    .insert(file_id, FileOpenType::Binary);
+            } else {
+                trace!("Process {} has no {}", proc_id, name);
+            }
+            Ok(())
+        };
         // Stick the process channels into the file map
-        if proc.stdin.is_some() {
-            trace!("Saving stdin for process {}", proc_id);
-            let file_id = self.take_file_id()?;
-            self.proc_stdin_ids.write()?.insert(proc_id, file_id);
-            self.file_modes
-                .write()?
-                .insert(file_id, FileOpenMode::Write);
-            self.file_types
-                .write()?
-                .insert(file_id, FileOpenType::Binary);
-        } else {
-            trace!("Process {} has no stdin", proc_id);
-        }
-        if proc.stdout.is_some() {
-            trace!("Saving stdout for process {}", proc_id);
-            let file_id = self.take_file_id()?;
-            self.proc_stdout_ids.write()?.insert(proc_id, file_id);
-            self.file_modes.write()?.insert(file_id, FileOpenMode::Read);
-            self.file_types
-                .write()?
-                .insert(file_id, FileOpenType::Binary);
-        } else {
-            trace!("Process {} has no stdout", proc_id);
-        }
-        if proc.stderr.is_some() {
-            trace!("Saving stderr for process {}", proc_id);
-            let file_id = self.take_file_id()?;
-            self.proc_stderr_ids.write()?.insert(proc_id, file_id);
-            self.file_modes.write()?.insert(file_id, FileOpenMode::Read);
-            self.file_types
-                .write()?
-                .insert(file_id, FileOpenType::Binary);
-        } else {
-            trace!("Process {} has no stderr", proc_id);
-        }
+        save_io(&proc.stdin, &self.proc_stdin_ids, "stdin", true)?;
+        save_io(&proc.stdout, &self.proc_stdout_ids, "stdout", false)?;
+        save_io(&proc.stderr, &self.proc_stderr_ids, "stderr", false)?;
 
         // Move the proc to the process map
         self.processes

diff --git a/crates/subprocess b/crates/subprocess
diff --git a/python/bh_agent_client.pyi b/python/bh_agent_client.pyi
@@ -25,6 +25,7 @@ class BhAgentClient:
         setuid: int | None,
         setgid: int | None,
         setpgid: int | None,
+        use_pty: bool | None,
     ) -> int: ...
     def get_process_ids(self, env_id: int) -> list[int]: ...
     def get_process_channel(self, env_id: int, proc_id: int, channel: int) -> int: ...

diff --git a/python/binharness/agentenvironment.py b/python/binharness/agentenvironment.py
@@ -183,7 +183,7 @@ def run_command(
         *args: Path | str | Sequence[Path | str],
         env: dict[str, str] | None = None,
         cwd: Path | None = None,
-        pty: bool = False,  # noqa: ARG002 TODO: Implement pty support
+        pty: bool = False,
     ) -> AgentProcess:
         """Run a command in the environment."""
         normalized_args = list(normalize_args(*args))
@@ -200,6 +200,7 @@ def run_command(
             setuid=None,
             setgid=None,
             setpgid=False,
+            use_pty=pty,
         )
         return AgentProcess(
             self._client,