detached_shell/
pty_handler.rs

1use std::fs::File;
2use std::io::{Read, Write};
3use std::os::unix::io::{FromRawFd, RawFd};
4use std::os::unix::net::UnixStream;
5use std::sync::{Arc, Mutex};
6
7use crate::error::{NdsError, Result};
8use crate::pty_buffer::PtyBuffer;
9
10/// Enhanced PTY handler that maintains state across connections
11pub struct PtyHandler {
12    master_fd: RawFd,
13    output_buffer: PtyBuffer,
14    last_prompt_position: Arc<Mutex<Option<usize>>>,
15}
16
17impl PtyHandler {
18    pub fn new(master_fd: RawFd, buffer_size: usize) -> Self {
19        PtyHandler {
20            master_fd,
21            output_buffer: PtyBuffer::new(buffer_size),
22            last_prompt_position: Arc::new(Mutex::new(None)),
23        }
24    }
25
26    /// Read from PTY and handle the data
27    pub fn read_pty_data(&mut self, buffer: &mut [u8]) -> Result<Option<Vec<u8>>> {
28        let master_file = unsafe { File::from_raw_fd(self.master_fd) };
29        let mut master_file_clone = master_file.try_clone()?;
30        std::mem::forget(master_file); // Don't close the fd
31
32        match master_file_clone.read(buffer) {
33            Ok(0) => Ok(None), // EOF
34            Ok(n) => {
35                let data = buffer[..n].to_vec();
36
37                // Detect prompt patterns (simple heuristic)
38                if Self::looks_like_prompt(&data) {
39                    *self.last_prompt_position.lock().unwrap() = Some(n);
40                }
41
42                Ok(Some(data))
43            }
44            Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
45                Ok(Some(Vec::new())) // No data available
46            }
47            Err(e) => Err(NdsError::Io(e)),
48        }
49    }
50
51    /// Write data to PTY
52    pub fn write_to_pty(&mut self, data: &[u8]) -> Result<()> {
53        let mut master_file = unsafe { File::from_raw_fd(self.master_fd) };
54        master_file.write_all(data)?;
55        std::mem::forget(master_file); // Don't close the fd
56        Ok(())
57    }
58
59    /// Process PTY data and distribute to client or buffer
60    pub fn process_pty_output(
61        &mut self,
62        data: &[u8],
63        client: &mut Option<UnixStream>,
64    ) -> Result<()> {
65        if let Some(ref mut stream) = client {
66            // Try to send to connected client
67            if let Err(e) = stream.write_all(data) {
68                if e.kind() == std::io::ErrorKind::BrokenPipe {
69                    // Client disconnected, buffer the data
70                    self.output_buffer.push(data);
71                    *client = None;
72                }
73                return Err(NdsError::Io(e));
74            }
75        } else {
76            // No client connected, buffer the output
77            self.output_buffer.push(data);
78        }
79        Ok(())
80    }
81
82    /// Send buffered data to a newly connected client
83    pub fn send_buffered_data(&mut self, client: &mut UnixStream) -> Result<()> {
84        if !self.output_buffer.is_empty() {
85            let mut buffered_data = Vec::new();
86            self.output_buffer.drain_to(&mut buffered_data);
87
88            // Send in manageable chunks
89            const CHUNK_SIZE: usize = 4096;
90            for chunk in buffered_data.chunks(CHUNK_SIZE) {
91                client.write_all(chunk)?;
92                client.flush()?;
93                // Small delay to avoid overwhelming the client
94                std::thread::sleep(std::time::Duration::from_micros(100));
95            }
96        }
97        Ok(())
98    }
99
100    /// Check if data looks like a shell prompt
101    fn looks_like_prompt(data: &[u8]) -> bool {
102        // Simple heuristic: check for common prompt endings
103        // This could be made more sophisticated
104        if data.len() < 2 {
105            return false;
106        }
107
108        // Check for common prompt characters at the end
109        let last_chars = &data[data.len().saturating_sub(10)..];
110        last_chars.contains(&b'$')
111            || last_chars.contains(&b'#')
112            || last_chars.contains(&b'>')
113            || last_chars.contains(&b'%')
114    }
115
116    /// Request terminal to refresh/redraw
117    pub fn refresh_terminal(&mut self) -> Result<()> {
118        // Send Ctrl+L to refresh the display
119        self.write_to_pty(b"\x0c")?;
120        Ok(())
121    }
122}
detached_shell/pty_handler.rs

detached_shell/
pty_handler.rs
