149 lines
4.5 KiB
Rust
149 lines
4.5 KiB
Rust
use std::io::{Read, Write};
|
|
use anyhow::Result;
|
|
use portable_pty::{CommandBuilder, MasterPty, PtySize, native_pty_system};
|
|
use tokio::sync::mpsc;
|
|
|
|
/// A spawned PTY with its child process. Output chunks arrive on `output`.
|
|
pub struct PtyHandle {
|
|
master: Box<dyn MasterPty + Send>,
|
|
writer: Box<dyn Write + Send>,
|
|
child: Box<dyn portable_pty::Child + Send + Sync>,
|
|
/// Raw output chunks read off the PTY master (already on the async side).
|
|
pub output: mpsc::Receiver<Vec<u8>>,
|
|
}
|
|
|
|
/// Parameters for spawning a surface's process.
|
|
pub struct SpawnSpec {
|
|
pub command: String,
|
|
pub args: Vec<String>,
|
|
pub cwd: std::path::PathBuf,
|
|
pub cols: u16,
|
|
pub rows: u16,
|
|
/// Extra environment variables (e.g. SPACESH_SURFACE_ID).
|
|
pub env: Vec<(String, String)>,
|
|
}
|
|
|
|
impl PtyHandle {
|
|
pub fn spawn(spec: SpawnSpec) -> Result<Self> {
|
|
let pty_system = native_pty_system();
|
|
let pair = pty_system.openpty(PtySize {
|
|
rows: spec.rows,
|
|
cols: spec.cols,
|
|
pixel_width: 0,
|
|
pixel_height: 0,
|
|
})?;
|
|
|
|
let mut cmd = CommandBuilder::new(&spec.command);
|
|
for a in &spec.args {
|
|
cmd.arg(a);
|
|
}
|
|
cmd.cwd(&spec.cwd);
|
|
for (k, v) in &spec.env {
|
|
cmd.env(k, v);
|
|
}
|
|
|
|
let child = pair.slave.spawn_command(cmd)?;
|
|
// The slave handle must be dropped so the child is the only holder; otherwise
|
|
// EOF is never observed on the master after the child exits.
|
|
drop(pair.slave);
|
|
|
|
let writer = pair.master.take_writer()?;
|
|
let mut reader = pair.master.try_clone_reader()?;
|
|
|
|
let (tx, rx) = mpsc::channel::<Vec<u8>>(256);
|
|
std::thread::spawn(move || {
|
|
let mut buf = [0u8; 8192];
|
|
loop {
|
|
match reader.read(&mut buf) {
|
|
Ok(0) => break, // EOF: child closed the pty
|
|
Ok(n) => {
|
|
if tx.blocking_send(buf[..n].to_vec()).is_err() {
|
|
break; // receiver gone
|
|
}
|
|
}
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
});
|
|
|
|
Ok(Self {
|
|
master: pair.master,
|
|
writer,
|
|
child,
|
|
output: rx,
|
|
})
|
|
}
|
|
|
|
pub fn write_input(&mut self, bytes: &[u8]) -> Result<()> {
|
|
self.writer.write_all(bytes)?;
|
|
self.writer.flush()?;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn resize(&self, cols: u16, rows: u16) -> Result<()> {
|
|
self.master.resize(PtySize { rows, cols, pixel_width: 0, pixel_height: 0 })?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Best-effort wait for the child's exit code (blocking).
|
|
pub fn wait(&mut self) -> i32 {
|
|
match self.child.wait() {
|
|
Ok(status) => status.exit_code() as i32,
|
|
Err(_) => -1,
|
|
}
|
|
}
|
|
|
|
pub fn kill(&mut self) {
|
|
let _ = self.child.kill();
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn shell_spec(script: &str) -> SpawnSpec {
|
|
SpawnSpec {
|
|
command: "/bin/sh".into(),
|
|
args: vec!["-c".into(), script.into()],
|
|
cwd: std::env::temp_dir(),
|
|
cols: 80,
|
|
rows: 24,
|
|
env: vec![("SPACESH_SURFACE_ID".into(), "s_test".into())],
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn spawn_echo_produces_output() {
|
|
let mut handle = PtyHandle::spawn(shell_spec("printf SPACESH_OK")).unwrap();
|
|
let mut collected = Vec::new();
|
|
// Drain until EOF (channel closes when the reader thread sees EOF).
|
|
while let Some(chunk) = handle.output.recv().await {
|
|
collected.extend_from_slice(&chunk);
|
|
}
|
|
let text = String::from_utf8_lossy(&collected);
|
|
assert!(text.contains("SPACESH_OK"), "got: {text:?}");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn resize_does_not_error() {
|
|
let handle = PtyHandle::spawn(shell_spec("sleep 0.2")).unwrap();
|
|
handle.resize(120, 40).unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn input_is_echoed_back() {
|
|
// `cat` echoes stdin back to stdout on a pty.
|
|
let mut handle = PtyHandle::spawn(shell_spec("cat")).unwrap();
|
|
handle.write_input(b"hello\n").unwrap();
|
|
let mut collected = Vec::new();
|
|
// Read a few chunks then kill cat to end the stream.
|
|
if let Some(chunk) = handle.output.recv().await {
|
|
collected.extend_from_slice(&chunk);
|
|
}
|
|
handle.kill();
|
|
let text = String::from_utf8_lossy(&collected);
|
|
assert!(text.contains("hello"), "got: {text:?}");
|
|
}
|
|
}
|