a3s_box_runtime/grpc/
pty.rs

1//! PTY terminal client for interactive sessions.
2
3use std::path::Path;
4use std::sync::Arc;
5
6use a3s_box_core::error::{BoxError, Result};
7use tokio::sync::Mutex;
8
9type PtyFrameReader = a3s_transport::FrameReader<tokio::io::ReadHalf<tokio::net::UnixStream>>;
10type PtyFrameWriter = a3s_transport::FrameWriter<tokio::io::WriteHalf<tokio::net::UnixStream>>;
11
12/// Client for interactive PTY sessions in the guest over Unix socket.
13///
14/// Connects to the PTY server (vsock port 4090) and provides async
15/// frame-based communication for bidirectional terminal I/O.
16/// Uses `a3s_transport::FrameReader`/`FrameWriter` for wire I/O.
17#[derive(Debug)]
18pub struct PtyClient {
19    reader: PtyFrameReader,
20    writer: PtyFrameWriter,
21}
22
23impl PtyClient {
24    /// Connect to the PTY server via Unix socket.
25    pub async fn connect(socket_path: &Path) -> Result<Self> {
26        let stream = tokio::net::UnixStream::connect(socket_path)
27            .await
28            .map_err(|e| {
29                BoxError::ExecError(format!(
30                    "Failed to connect to PTY server at {}: {}",
31                    socket_path.display(),
32                    e,
33                ))
34            })?;
35
36        let (r, w) = tokio::io::split(stream);
37        Ok(Self {
38            reader: a3s_transport::FrameReader::new(r),
39            writer: a3s_transport::FrameWriter::new(w),
40        })
41    }
42
43    /// Send a PtyRequest to start an interactive session.
44    pub async fn send_request(&mut self, req: &a3s_box_core::pty::PtyRequest) -> Result<()> {
45        let payload = serde_json::to_vec(req)
46            .map_err(|e| BoxError::ExecError(format!("Failed to serialize PtyRequest: {}", e)))?;
47        self.write_raw_frame(a3s_box_core::pty::FRAME_PTY_REQUEST, &payload)
48            .await
49    }
50
51    /// Start a streaming PTY session and return a handle for supervision.
52    pub async fn start_stream(
53        mut self,
54        req: &a3s_box_core::pty::PtyRequest,
55    ) -> Result<StreamingPty> {
56        self.send_request(req).await?;
57        Ok(StreamingPty {
58            reader: self.reader,
59            writer: Arc::new(Mutex::new(self.writer)),
60            started: std::time::Instant::now(),
61            stdout_bytes: 0,
62            done: false,
63        })
64    }
65
66    /// Send terminal data to the guest.
67    pub async fn send_data(&mut self, data: &[u8]) -> Result<()> {
68        self.write_raw_frame(a3s_box_core::pty::FRAME_PTY_DATA, data)
69            .await
70    }
71
72    /// Send a terminal resize notification.
73    pub async fn send_resize(&mut self, cols: u16, rows: u16) -> Result<()> {
74        let resize = a3s_box_core::pty::PtyResize { cols, rows };
75        let payload = serde_json::to_vec(&resize)
76            .map_err(|e| BoxError::ExecError(format!("Failed to serialize PtyResize: {}", e)))?;
77        self.write_raw_frame(a3s_box_core::pty::FRAME_PTY_RESIZE, &payload)
78            .await
79    }
80
81    /// Read the next frame from the guest.
82    ///
83    /// Returns `Ok(None)` on EOF (guest disconnected).
84    pub async fn read_frame(&mut self) -> Result<Option<(u8, Vec<u8>)>> {
85        match self.reader.read_frame().await {
86            Ok(Some(frame)) => Ok(Some((frame.frame_type as u8, frame.payload))),
87            Ok(None) => Ok(None),
88            Err(e) => Err(BoxError::ExecError(format!("PTY frame read failed: {}", e))),
89        }
90    }
91
92    /// Split the client into read and write halves for concurrent I/O.
93    pub fn into_split(
94        self,
95    ) -> (
96        a3s_transport::FrameReader<tokio::io::ReadHalf<tokio::net::UnixStream>>,
97        a3s_transport::FrameWriter<tokio::io::WriteHalf<tokio::net::UnixStream>>,
98    ) {
99        (self.reader, self.writer)
100    }
101
102    /// Write a raw PTY frame using the transport writer.
103    async fn write_raw_frame(&mut self, frame_type: u8, payload: &[u8]) -> Result<()> {
104        // PTY uses custom frame type bytes (0x01-0x05) that map to transport FrameType
105        let ft = a3s_transport::FrameType::try_from(frame_type)
106            .unwrap_or(a3s_transport::FrameType::Data);
107        let frame = a3s_transport::Frame {
108            frame_type: ft,
109            payload: payload.to_vec(),
110        };
111        self.writer
112            .write_frame(&frame)
113            .await
114            .map_err(|e| BoxError::ExecError(format!("PTY frame write failed: {}", e)))
115    }
116}
117
118/// Cloneable input side for a running PTY session.
119#[derive(Clone, Debug)]
120pub struct StreamingPtyInput {
121    writer: Arc<Mutex<PtyFrameWriter>>,
122}
123
124impl StreamingPtyInput {
125    /// Write terminal data to the PTY.
126    pub async fn write_stdin(&self, data: &[u8]) -> Result<()> {
127        let frame = a3s_transport::Frame {
128            frame_type: a3s_transport::FrameType::Control,
129            payload: data.to_vec(),
130        };
131        self.writer
132            .lock()
133            .await
134            .write_frame(&frame)
135            .await
136            .map_err(|e| BoxError::ExecError(format!("PTY stdin write failed: {}", e)))
137    }
138
139    /// Request terminal resize.
140    pub async fn resize(&self, cols: u16, rows: u16) -> Result<()> {
141        let resize = a3s_box_core::pty::PtyResize { cols, rows };
142        let payload = serde_json::to_vec(&resize)
143            .map_err(|e| BoxError::ExecError(format!("Failed to serialize PtyResize: {}", e)))?;
144        let frame = a3s_transport::Frame {
145            frame_type: a3s_transport::FrameType::Heartbeat,
146            payload,
147        };
148        self.writer
149            .lock()
150            .await
151            .write_frame(&frame)
152            .await
153            .map_err(|e| BoxError::ExecError(format!("PTY resize write failed: {}", e)))
154    }
155
156    /// Close the PTY control stream. The guest treats this as a session stop.
157    pub async fn close(&self) -> Result<()> {
158        self.writer
159            .lock()
160            .await
161            .write_frame(&a3s_transport::Frame::close())
162            .await
163            .map_err(|e| BoxError::ExecError(format!("PTY close write failed: {}", e)))
164    }
165}
166
167/// Handle for reading PTY session output and exit status.
168pub struct StreamingPty {
169    reader: PtyFrameReader,
170    writer: Arc<Mutex<PtyFrameWriter>>,
171    started: std::time::Instant,
172    stdout_bytes: u64,
173    done: bool,
174}
175
176impl StreamingPty {
177    /// Return a cloneable input handle for this PTY session.
178    pub fn input(&self) -> StreamingPtyInput {
179        StreamingPtyInput {
180            writer: self.writer.clone(),
181        }
182    }
183
184    /// Read the next PTY output or exit event.
185    pub async fn next_event(&mut self) -> Result<Option<a3s_box_core::exec::ExecEvent>> {
186        use a3s_box_core::exec::{ExecChunk, ExecEvent, ExecExit, StreamType};
187
188        if self.done {
189            return Ok(None);
190        }
191
192        let Some(frame) = self
193            .reader
194            .read_frame()
195            .await
196            .map_err(|e| BoxError::ExecError(format!("PTY frame read failed: {}", e)))?
197        else {
198            self.done = true;
199            return Ok(None);
200        };
201
202        match frame.frame_type as u8 {
203            a3s_box_core::pty::FRAME_PTY_DATA => {
204                self.stdout_bytes += frame.payload.len() as u64;
205                Ok(Some(ExecEvent::Chunk(ExecChunk {
206                    stream: StreamType::Stdout,
207                    data: frame.payload,
208                })))
209            }
210            a3s_box_core::pty::FRAME_PTY_EXIT => {
211                let exit: a3s_box_core::pty::PtyExit = serde_json::from_slice(&frame.payload)
212                    .map_err(|e| BoxError::ExecError(format!("Failed to parse PTY exit: {}", e)))?;
213                self.done = true;
214                Ok(Some(ExecEvent::Exit(ExecExit {
215                    exit_code: exit.exit_code,
216                    oom_killed: false,
217                })))
218            }
219            a3s_box_core::pty::FRAME_PTY_ERROR => {
220                let msg = String::from_utf8_lossy(&frame.payload);
221                self.done = true;
222                Err(BoxError::ExecError(format!("PTY session error: {}", msg)))
223            }
224            other => Err(BoxError::ExecError(format!(
225                "Unexpected PTY frame type in stream: 0x{other:02x}"
226            ))),
227        }
228    }
229
230    /// Request cancellation of the running PTY session.
231    pub async fn cancel(&mut self) -> Result<()> {
232        self.input().close().await
233    }
234
235    /// Whether the PTY stream has finished.
236    pub fn is_done(&self) -> bool {
237        self.done
238    }
239
240    /// Get execution metrics so far.
241    pub fn metrics(&self) -> a3s_box_core::exec::ExecMetrics {
242        a3s_box_core::exec::ExecMetrics {
243            duration_ms: self.started.elapsed().as_millis() as u64,
244            peak_memory_bytes: None,
245            stdout_bytes: self.stdout_bytes,
246            stderr_bytes: 0,
247        }
248    }
249}
250
251impl std::fmt::Debug for StreamingPty {
252    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
253        f.debug_struct("StreamingPty")
254            .field("done", &self.done)
255            .field("stdout_bytes", &self.stdout_bytes)
256            .finish()
257    }
258}
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263    use tokio::io::{AsyncReadExt, AsyncWriteExt};
264    use tokio::net::UnixListener;
265
266    fn bind_test_listener(path: &Path) -> Option<UnixListener> {
267        match UnixListener::bind(path) {
268            Ok(listener) => Some(listener),
269            Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied => {
270                eprintln!(
271                    "skipping Unix socket test; sandbox denied bind at {}: {}",
272                    path.display(),
273                    e
274                );
275                None
276            }
277            Err(e) => panic!("failed to bind test socket {}: {}", path.display(), e),
278        }
279    }
280
281    #[tokio::test]
282    async fn test_pty_client_connect_nonexistent() {
283        let result = PtyClient::connect(Path::new("/tmp/nonexistent-pty-test.sock")).await;
284        assert!(result.is_err());
285    }
286
287    #[tokio::test]
288    async fn test_pty_frame_roundtrip() {
289        let tmp = tempfile::TempDir::new().unwrap();
290        let sock_path = tmp.path().join("pty.sock");
291        let Some(listener) = bind_test_listener(&sock_path) else {
292            return;
293        };
294
295        let sock_path_clone = sock_path.clone();
296        let server = tokio::spawn(async move {
297            let (mut stream, _) = listener.accept().await.unwrap();
298            // Read a frame: [type:1][len:4][payload]
299            let mut header = [0u8; 5];
300            stream.read_exact(&mut header).await.unwrap();
301            let frame_type = header[0];
302            let len = u32::from_be_bytes([header[1], header[2], header[3], header[4]]) as usize;
303            let mut payload = vec![0u8; len];
304            if len > 0 {
305                stream.read_exact(&mut payload).await.unwrap();
306            }
307            // Echo it back
308            stream.write_all(&header).await.unwrap();
309            stream.write_all(&payload).await.unwrap();
310            (frame_type, payload)
311        });
312
313        let mut client = PtyClient::connect(&sock_path_clone).await.unwrap();
314        client.send_data(b"hello world").await.unwrap();
315
316        let frame = client.read_frame().await.unwrap().unwrap();
317        assert_eq!(frame.0, a3s_box_core::pty::FRAME_PTY_DATA);
318        assert_eq!(&frame.1[..], b"hello world");
319
320        let (server_type, server_payload) = server.await.unwrap();
321        assert_eq!(server_type, a3s_box_core::pty::FRAME_PTY_DATA);
322        assert_eq!(&server_payload[..], b"hello world");
323    }
324
325    #[tokio::test]
326    async fn test_pty_send_resize() {
327        let tmp = tempfile::TempDir::new().unwrap();
328        let sock_path = tmp.path().join("pty_resize.sock");
329        let Some(listener) = bind_test_listener(&sock_path) else {
330            return;
331        };
332
333        tokio::spawn(async move {
334            let (mut stream, _) = listener.accept().await.unwrap();
335            let mut header = [0u8; 5];
336            stream.read_exact(&mut header).await.unwrap();
337            let frame_type = header[0];
338            let len = u32::from_be_bytes([header[1], header[2], header[3], header[4]]) as usize;
339            let mut payload = vec![0u8; len];
340            stream.read_exact(&mut payload).await.unwrap();
341
342            assert_eq!(frame_type, a3s_box_core::pty::FRAME_PTY_RESIZE);
343            let resize: a3s_box_core::pty::PtyResize = serde_json::from_slice(&payload).unwrap();
344            assert_eq!(resize.cols, 120);
345            assert_eq!(resize.rows, 40);
346        });
347
348        let mut client = PtyClient::connect(&sock_path).await.unwrap();
349        client.send_resize(120, 40).await.unwrap();
350    }
351
352    #[tokio::test]
353    async fn test_pty_read_frame_eof() {
354        let tmp = tempfile::TempDir::new().unwrap();
355        let sock_path = tmp.path().join("pty_eof.sock");
356        let Some(listener) = bind_test_listener(&sock_path) else {
357            return;
358        };
359
360        tokio::spawn(async move {
361            let (stream, _) = listener.accept().await.unwrap();
362            drop(stream); // Close immediately → EOF
363        });
364
365        let mut client = PtyClient::connect(&sock_path).await.unwrap();
366        let frame = client.read_frame().await.unwrap();
367        assert!(frame.is_none()); // EOF
368    }
369
370    #[tokio::test]
371    async fn test_pty_client_start_stream_reads_data_and_exit() {
372        let tmp = tempfile::TempDir::new().unwrap();
373        let sock_path = tmp.path().join("pty_stream.sock");
374        let Some(listener) = bind_test_listener(&sock_path) else {
375            return;
376        };
377
378        tokio::spawn(async move {
379            let (stream, _) = listener.accept().await.unwrap();
380            let (r, w) = tokio::io::split(stream);
381            let mut reader = a3s_transport::FrameReader::new(r);
382            let mut writer = a3s_transport::FrameWriter::new(w);
383
384            let request = reader.read_frame().await.unwrap().unwrap();
385            assert_eq!(
386                request.frame_type as u8,
387                a3s_box_core::pty::FRAME_PTY_REQUEST
388            );
389
390            writer
391                .write_frame(&a3s_transport::Frame {
392                    frame_type: a3s_transport::FrameType::Control,
393                    payload: b"tty output".to_vec(),
394                })
395                .await
396                .unwrap();
397
398            let exit = a3s_box_core::pty::PtyExit { exit_code: 9 };
399            writer
400                .write_frame(&a3s_transport::Frame {
401                    frame_type: a3s_transport::FrameType::Error,
402                    payload: serde_json::to_vec(&exit).unwrap(),
403                })
404                .await
405                .unwrap();
406        });
407
408        let client = PtyClient::connect(&sock_path).await.unwrap();
409        let req = a3s_box_core::pty::PtyRequest {
410            cmd: vec!["/bin/sh".to_string()],
411            env: vec![],
412            working_dir: None,
413            rootfs: None,
414            user: None,
415            cols: 80,
416            rows: 24,
417        };
418        let mut stream = client.start_stream(&req).await.unwrap();
419
420        match stream.next_event().await.unwrap().unwrap() {
421            a3s_box_core::exec::ExecEvent::Chunk(chunk) => {
422                assert_eq!(chunk.stream, a3s_box_core::exec::StreamType::Stdout);
423                assert_eq!(chunk.data, b"tty output");
424            }
425            other => panic!("unexpected event: {other:?}"),
426        }
427        match stream.next_event().await.unwrap().unwrap() {
428            a3s_box_core::exec::ExecEvent::Exit(exit) => assert_eq!(exit.exit_code, 9),
429            other => panic!("unexpected event: {other:?}"),
430        }
431    }
432
433    #[tokio::test]
434    async fn test_pty_stream_cancel_writes_close_frame() {
435        let tmp = tempfile::TempDir::new().unwrap();
436        let sock_path = tmp.path().join("pty_cancel.sock");
437        let Some(listener) = bind_test_listener(&sock_path) else {
438            return;
439        };
440
441        tokio::spawn(async move {
442            let (stream, _) = listener.accept().await.unwrap();
443            let (r, _w) = tokio::io::split(stream);
444            let mut reader = a3s_transport::FrameReader::new(r);
445
446            let request = reader.read_frame().await.unwrap().unwrap();
447            assert_eq!(
448                request.frame_type as u8,
449                a3s_box_core::pty::FRAME_PTY_REQUEST
450            );
451            let close = reader.read_frame().await.unwrap().unwrap();
452            assert_eq!(close.frame_type, a3s_transport::FrameType::Close);
453            assert!(close.payload.is_empty());
454        });
455
456        let client = PtyClient::connect(&sock_path).await.unwrap();
457        let req = a3s_box_core::pty::PtyRequest {
458            cmd: vec!["sleep".to_string(), "60".to_string()],
459            env: vec![],
460            working_dir: None,
461            rootfs: None,
462            user: None,
463            cols: 80,
464            rows: 24,
465        };
466        let mut stream = client.start_stream(&req).await.unwrap();
467        stream.cancel().await.unwrap();
468    }
469}
a3s_box_runtime/grpc/pty.rs

a3s_box_runtime/grpc/
pty.rs
