use std::{ io::Write, sync::{Arc, Mutex}, }; use bytes::Bytes; use eyre::{eyre, Result}; use portable_pty::{CommandBuilder, MasterPty, PtySize}; pub struct Pty { tx: tokio::sync::mpsc::Sender, pub master: Arc>>, pub reader: Arc>>, } impl Pty { pub async fn open<'a>(rows: u16, cols: u16) -> Result { let sys = portable_pty::native_pty_system(); let pair = sys .openpty(PtySize { rows, cols, pixel_width: 0, pixel_height: 0, }) .map_err(|e| eyre!("Failed to open pty: {}", e))?; let cmd = CommandBuilder::new_default_prog(); tokio::task::spawn_blocking(move || { let mut child = pair.slave.spawn_command(cmd).unwrap(); // Wait for the child to exit let _ = child.wait().unwrap(); // Ensure slave is dropped // This closes file handles, we can deadlock if this is not done correctly. drop(pair.slave); }); // Handle input -> write to master writer let (master_tx, mut master_rx) = tokio::sync::mpsc::channel::(32); let mut writer = pair.master.take_writer().unwrap(); let reader = pair .master .try_clone_reader() .map_err(|e| e.to_string()) .expect("Failed to clone reader"); tokio::spawn(async move { while let Some(bytes) = master_rx.recv().await { writer.write_all(&bytes).unwrap(); writer.flush().unwrap(); } // When the channel has been closed, we won't be getting any more input. Close the // writer and the master. // This will also close the writer, which sends EOF to the underlying shell. Ensuring // that is also closed. drop(writer); }); Ok(Pty { tx: master_tx, master: Arc::new(Mutex::new(pair.master)), reader: Arc::new(Mutex::new(reader)), }) } pub async fn resize(&self, rows: u16, cols: u16) -> Result<()> { let master = self .master .lock() .map_err(|e| eyre!("Failed to lock pty master: {e}"))?; master .resize(PtySize { rows, cols, pixel_width: 0, pixel_height: 0, }) .map_err(|e| eyre!("Failed to resize terminal: {e}"))?; Ok(()) } pub async fn send_bytes(&self, bytes: Bytes) -> Result<()> { self.tx .send(bytes) .await .map_err(|e| eyre!("Failed to write to master tx: {}", e)) } pub async fn send_string(&self, cmd: &str) -> Result<()> { let bytes: Vec = cmd.bytes().collect(); let bytes = Bytes::from(bytes); self.send_bytes(bytes).await } pub async fn send_single_string(&self, cmd: &str) -> Result<()> { let mut bytes: Vec = cmd.bytes().collect(); bytes.push(0x04); let bytes = Bytes::from(bytes); self.send_bytes(bytes).await } }