| 1 | //! PTY (pseudo-terminal) management |
| 2 | //! |
| 3 | //! Handles spawning the shell process and I/O with it. |
| 4 | |
| 5 | use anyhow::Result; |
| 6 | use portable_pty::{native_pty_system, CommandBuilder, PtyPair, PtySize}; |
| 7 | use std::io::{Read, Write}; |
| 8 | use std::sync::atomic::{AtomicBool, Ordering}; |
| 9 | use std::sync::mpsc::{self, Receiver, Sender}; |
| 10 | use std::sync::Arc; |
| 11 | use std::thread; |
| 12 | |
| 13 | /// Manages a PTY connection to a shell process |
| 14 | pub struct Pty { |
| 15 | pair: PtyPair, |
| 16 | writer: Box<dyn Write + Send>, |
| 17 | output_rx: Receiver<Vec<u8>>, |
| 18 | _output_thread: thread::JoinHandle<()>, |
| 19 | /// Flag indicating the shell has exited |
| 20 | shell_exited: Arc<AtomicBool>, |
| 21 | } |
| 22 | |
| 23 | impl Pty { |
| 24 | /// Spawn a new PTY with the user's shell |
| 25 | pub fn spawn(cols: u16, rows: u16) -> Result<Self> { |
| 26 | let pty_system = native_pty_system(); |
| 27 | |
| 28 | let pair = pty_system.openpty(PtySize { |
| 29 | rows, |
| 30 | cols, |
| 31 | pixel_width: 0, |
| 32 | pixel_height: 0, |
| 33 | })?; |
| 34 | |
| 35 | // Get the user's shell from $SHELL, fallback to /bin/sh |
| 36 | let shell = std::env::var("SHELL").unwrap_or_else(|_| "/bin/sh".to_string()); |
| 37 | |
| 38 | let mut cmd = CommandBuilder::new(&shell); |
| 39 | // Start shell as login shell |
| 40 | cmd.arg("-l"); |
| 41 | |
| 42 | // Set working directory to current directory |
| 43 | if let Ok(cwd) = std::env::current_dir() { |
| 44 | cmd.cwd(cwd); |
| 45 | } |
| 46 | |
| 47 | // Spawn the shell |
| 48 | let _child = pair.slave.spawn_command(cmd)?; |
| 49 | |
| 50 | // Get writer for sending input to the PTY |
| 51 | let writer = pair.master.take_writer()?; |
| 52 | |
| 53 | // Set up a thread to read output from the PTY |
| 54 | let mut reader = pair.master.try_clone_reader()?; |
| 55 | let (output_tx, output_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = mpsc::channel(); |
| 56 | |
| 57 | // Flag to signal when shell exits |
| 58 | let shell_exited = Arc::new(AtomicBool::new(false)); |
| 59 | let shell_exited_clone = Arc::clone(&shell_exited); |
| 60 | |
| 61 | let output_thread = thread::spawn(move || { |
| 62 | let mut buf = [0u8; 4096]; |
| 63 | loop { |
| 64 | match reader.read(&mut buf) { |
| 65 | Ok(0) => { |
| 66 | // EOF - shell exited |
| 67 | shell_exited_clone.store(true, Ordering::SeqCst); |
| 68 | break; |
| 69 | } |
| 70 | Ok(n) => { |
| 71 | if output_tx.send(buf[..n].to_vec()).is_err() { |
| 72 | break; // Receiver dropped |
| 73 | } |
| 74 | } |
| 75 | Err(_) => { |
| 76 | // Error - likely shell exited |
| 77 | shell_exited_clone.store(true, Ordering::SeqCst); |
| 78 | break; |
| 79 | } |
| 80 | } |
| 81 | } |
| 82 | }); |
| 83 | |
| 84 | Ok(Self { |
| 85 | pair, |
| 86 | writer, |
| 87 | output_rx, |
| 88 | _output_thread: output_thread, |
| 89 | shell_exited, |
| 90 | }) |
| 91 | } |
| 92 | |
| 93 | /// Send input bytes to the PTY |
| 94 | pub fn write(&mut self, data: &[u8]) -> Result<()> { |
| 95 | self.writer.write_all(data)?; |
| 96 | self.writer.flush()?; |
| 97 | Ok(()) |
| 98 | } |
| 99 | |
| 100 | /// Read any available output from the PTY (non-blocking) |
| 101 | pub fn read(&mut self) -> Option<Vec<u8>> { |
| 102 | // Collect all available output |
| 103 | let mut output = Vec::new(); |
| 104 | while let Ok(data) = self.output_rx.try_recv() { |
| 105 | output.extend(data); |
| 106 | } |
| 107 | if output.is_empty() { |
| 108 | None |
| 109 | } else { |
| 110 | Some(output) |
| 111 | } |
| 112 | } |
| 113 | |
| 114 | /// Resize the PTY |
| 115 | pub fn resize(&self, cols: u16, rows: u16) -> Result<()> { |
| 116 | self.pair.master.resize(PtySize { |
| 117 | rows, |
| 118 | cols, |
| 119 | pixel_width: 0, |
| 120 | pixel_height: 0, |
| 121 | })?; |
| 122 | Ok(()) |
| 123 | } |
| 124 | |
| 125 | /// Check if the shell is still alive |
| 126 | pub fn is_alive(&self) -> bool { |
| 127 | !self.shell_exited.load(Ordering::SeqCst) |
| 128 | } |
| 129 | } |
| 130 |