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 a3s_box_core::ExecutionProcessSignal;
8use tokio::sync::Mutex;
9
10const PTY_CONTROL_SIGNAL: &[u8] = b"signal:";
11
12type PtyFrameReader = a3s_transport::FrameReader<tokio::io::ReadHalf<tokio::net::UnixStream>>;
13type PtyFrameWriter = a3s_transport::FrameWriter<tokio::io::WriteHalf<tokio::net::UnixStream>>;
14
15/// Client for interactive PTY sessions in the guest over Unix socket.
16///
17/// Connects to the PTY server (vsock port 4090) and provides async
18/// frame-based communication for bidirectional terminal I/O.
19/// Uses `a3s_transport::FrameReader`/`FrameWriter` for wire I/O.
20#[derive(Debug)]
21pub struct PtyClient {
22    reader: PtyFrameReader,
23    writer: PtyFrameWriter,
24}
25
26impl PtyClient {
27    /// Connect to the PTY server via Unix socket.
28    pub async fn connect(socket_path: &Path) -> Result<Self> {
29        let stream = tokio::net::UnixStream::connect(socket_path)
30            .await
31            .map_err(|e| {
32                BoxError::ExecError(format!(
33                    "Failed to connect to PTY server at {}: {}",
34                    socket_path.display(),
35                    e,
36                ))
37            })?;
38
39        let (r, w) = tokio::io::split(stream);
40        Ok(Self {
41            reader: a3s_transport::FrameReader::new(r),
42            writer: a3s_transport::FrameWriter::new(w),
43        })
44    }
45
46    /// Send a PtyRequest to start an interactive session.
47    pub async fn send_request(&mut self, req: &a3s_box_core::pty::PtyRequest) -> Result<()> {
48        let payload = serde_json::to_vec(req)
49            .map_err(|e| BoxError::ExecError(format!("Failed to serialize PtyRequest: {}", e)))?;
50        self.write_raw_frame(a3s_box_core::pty::FRAME_PTY_REQUEST, &payload)
51            .await
52    }
53
54    /// Start a streaming PTY session and return a handle for supervision.
55    pub async fn start_stream(
56        mut self,
57        req: &a3s_box_core::pty::PtyRequest,
58    ) -> Result<StreamingPty> {
59        self.send_request(req).await?;
60        Ok(StreamingPty {
61            reader: self.reader,
62            writer: Arc::new(Mutex::new(self.writer)),
63            started: std::time::Instant::now(),
64            stdout_bytes: 0,
65            done: false,
66        })
67    }
68
69    /// Send terminal data to the guest.
70    pub async fn send_data(&mut self, data: &[u8]) -> Result<()> {
71        self.write_raw_frame(a3s_box_core::pty::FRAME_PTY_DATA, data)
72            .await
73    }
74
75    /// Send a terminal resize notification.
76    pub async fn send_resize(&mut self, cols: u16, rows: u16) -> Result<()> {
77        let resize = a3s_box_core::pty::PtyResize { cols, rows };
78        let payload = serde_json::to_vec(&resize)
79            .map_err(|e| BoxError::ExecError(format!("Failed to serialize PtyResize: {}", e)))?;
80        self.write_raw_frame(a3s_box_core::pty::FRAME_PTY_RESIZE, &payload)
81            .await
82    }
83
84    /// Read the next frame from the guest.
85    ///
86    /// Returns `Ok(None)` on EOF (guest disconnected).
87    pub async fn read_frame(&mut self) -> Result<Option<(u8, Vec<u8>)>> {
88        match self.reader.read_frame().await {
89            Ok(Some(frame)) => Ok(Some((frame.frame_type as u8, frame.payload))),
90            Ok(None) => Ok(None),
91            Err(e) => Err(BoxError::ExecError(format!("PTY frame read failed: {}", e))),
92        }
93    }
94
95    /// Split the client into read and write halves for concurrent I/O.
96    pub fn into_split(
97        self,
98    ) -> (
99        a3s_transport::FrameReader<tokio::io::ReadHalf<tokio::net::UnixStream>>,
100        a3s_transport::FrameWriter<tokio::io::WriteHalf<tokio::net::UnixStream>>,
101    ) {
102        (self.reader, self.writer)
103    }
104
105    /// Write a raw PTY frame using the transport writer.
106    async fn write_raw_frame(&mut self, frame_type: u8, payload: &[u8]) -> Result<()> {
107        // PTY uses custom frame type bytes (0x01-0x05) that map to transport FrameType
108        let ft = a3s_transport::FrameType::try_from(frame_type)
109            .unwrap_or(a3s_transport::FrameType::Data);
110        let frame = a3s_transport::Frame {
111            frame_type: ft,
112            payload: payload.to_vec(),
113        };
114        self.writer
115            .write_frame(&frame)
116            .await
117            .map_err(|e| BoxError::ExecError(format!("PTY frame write failed: {}", e)))
118    }
119}
120
121/// Cloneable input side for a running PTY session.
122#[derive(Clone, Debug)]
123pub struct StreamingPtyInput {
124    writer: Arc<Mutex<PtyFrameWriter>>,
125}
126
127impl StreamingPtyInput {
128    /// Write terminal data to the PTY.
129    pub async fn write_stdin(&self, data: &[u8]) -> Result<()> {
130        let frame = a3s_transport::Frame {
131            frame_type: a3s_transport::FrameType::Control,
132            payload: data.to_vec(),
133        };
134        self.writer
135            .lock()
136            .await
137            .write_frame(&frame)
138            .await
139            .map_err(|e| BoxError::ExecError(format!("PTY stdin write failed: {}", e)))
140    }
141
142    /// Request terminal resize.
143    pub async fn resize(&self, cols: u16, rows: u16) -> Result<()> {
144        let resize = a3s_box_core::pty::PtyResize { cols, rows };
145        let payload = serde_json::to_vec(&resize)
146            .map_err(|e| BoxError::ExecError(format!("Failed to serialize PtyResize: {}", e)))?;
147        let frame = a3s_transport::Frame {
148            frame_type: a3s_transport::FrameType::Heartbeat,
149            payload,
150        };
151        self.writer
152            .lock()
153            .await
154            .write_frame(&frame)
155            .await
156            .map_err(|e| BoxError::ExecError(format!("PTY resize write failed: {}", e)))
157    }
158
159    /// Close the PTY control stream. The guest treats this as a session stop.
160    pub async fn close(&self) -> Result<()> {
161        self.writer
162            .lock()
163            .await
164            .write_frame(&a3s_transport::Frame::close())
165            .await
166            .map_err(|e| BoxError::ExecError(format!("PTY close write failed: {}", e)))
167    }
168
169    /// Deliver a typed Linux workload signal to the PTY process group.
170    pub async fn send_signal(&self, signal: ExecutionProcessSignal) -> Result<()> {
171        if signal == ExecutionProcessSignal::Kill {
172            return self.close().await;
173        }
174        let mut payload = PTY_CONTROL_SIGNAL.to_vec();
175        payload.extend_from_slice(signal.linux_number().to_string().as_bytes());
176        let frame = a3s_transport::Frame {
177            // PTY frames are directional. Error/Close (0x05) is the existing
178            // host-to-guest control lane and remains an error guest-to-host.
179            frame_type: a3s_transport::FrameType::Close,
180            payload,
181        };
182        self.writer
183            .lock()
184            .await
185            .write_frame(&frame)
186            .await
187            .map_err(|e| BoxError::ExecError(format!("PTY signal write failed: {}", e)))
188    }
189}
190
191/// Handle for reading PTY session output and exit status.
192pub struct StreamingPty {
193    reader: PtyFrameReader,
194    writer: Arc<Mutex<PtyFrameWriter>>,
195    started: std::time::Instant,
196    stdout_bytes: u64,
197    done: bool,
198}
199
200impl StreamingPty {
201    /// Return a cloneable input handle for this PTY session.
202    pub fn input(&self) -> StreamingPtyInput {
203        StreamingPtyInput {
204            writer: self.writer.clone(),
205        }
206    }
207
208    /// Read the next PTY output or exit event.
209    pub async fn next_event(&mut self) -> Result<Option<a3s_box_core::exec::ExecEvent>> {
210        use a3s_box_core::exec::{ExecChunk, ExecEvent, ExecExit, StreamType};
211
212        if self.done {
213            return Ok(None);
214        }
215
216        let Some(frame) = self
217            .reader
218            .read_frame()
219            .await
220            .map_err(|e| BoxError::ExecError(format!("PTY frame read failed: {}", e)))?
221        else {
222            self.done = true;
223            return Ok(None);
224        };
225
226        match frame.frame_type as u8 {
227            a3s_box_core::pty::FRAME_PTY_DATA => {
228                self.stdout_bytes += frame.payload.len() as u64;
229                Ok(Some(ExecEvent::Chunk(ExecChunk {
230                    stream: StreamType::Stdout,
231                    data: frame.payload,
232                })))
233            }
234            a3s_box_core::pty::FRAME_PTY_EXIT => {
235                let exit: a3s_box_core::pty::PtyExit = serde_json::from_slice(&frame.payload)
236                    .map_err(|e| BoxError::ExecError(format!("Failed to parse PTY exit: {}", e)))?;
237                self.done = true;
238                Ok(Some(ExecEvent::Exit(ExecExit {
239                    exit_code: exit.exit_code,
240                    oom_killed: false,
241                })))
242            }
243            a3s_box_core::pty::FRAME_PTY_ERROR => {
244                let msg = String::from_utf8_lossy(&frame.payload);
245                self.done = true;
246                Err(BoxError::ExecError(format!("PTY session error: {}", msg)))
247            }
248            other => Err(BoxError::ExecError(format!(
249                "Unexpected PTY frame type in stream: 0x{other:02x}"
250            ))),
251        }
252    }
253
254    /// Request cancellation of the running PTY session.
255    pub async fn cancel(&mut self) -> Result<()> {
256        self.input().close().await
257    }
258
259    /// Whether the PTY stream has finished.
260    pub fn is_done(&self) -> bool {
261        self.done
262    }
263
264    /// Get execution metrics so far.
265    pub fn metrics(&self) -> a3s_box_core::exec::ExecMetrics {
266        a3s_box_core::exec::ExecMetrics {
267            duration_ms: self.started.elapsed().as_millis() as u64,
268            peak_memory_bytes: None,
269            stdout_bytes: self.stdout_bytes,
270            stderr_bytes: 0,
271        }
272    }
273}
274
275impl std::fmt::Debug for StreamingPty {
276    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
277        f.debug_struct("StreamingPty")
278            .field("done", &self.done)
279            .field("stdout_bytes", &self.stdout_bytes)
280            .finish()
281    }
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287    use tokio::io::{AsyncReadExt, AsyncWriteExt};
288    use tokio::net::UnixListener;
289
290    fn bind_test_listener(path: &Path) -> Option<UnixListener> {
291        match UnixListener::bind(path) {
292            Ok(listener) => Some(listener),
293            Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied => {
294                eprintln!(
295                    "skipping Unix socket test; sandbox denied bind at {}: {}",
296                    path.display(),
297                    e
298                );
299                None
300            }
301            Err(e) => panic!("failed to bind test socket {}: {}", path.display(), e),
302        }
303    }
304
305    fn test_pty_request() -> a3s_box_core::pty::PtyRequest {
306        a3s_box_core::pty::PtyRequest {
307            cmd: vec!["/bin/sh".to_string()],
308            env: vec!["TERM=xterm".to_string()],
309            working_dir: Some("/workspace".to_string()),
310            rootfs: Some("/run/a3s/rootfs".to_string()),
311            user: Some("1000:1000".to_string()),
312            cols: 100,
313            rows: 30,
314        }
315    }
316
317    #[tokio::test]
318    async fn test_pty_client_connect_nonexistent() {
319        let result = PtyClient::connect(Path::new("/tmp/nonexistent-pty-test.sock")).await;
320        assert!(result.is_err());
321    }
322
323    #[tokio::test]
324    async fn test_pty_frame_roundtrip() {
325        let tmp = tempfile::TempDir::new().unwrap();
326        let sock_path = tmp.path().join("pty.sock");
327        let Some(listener) = bind_test_listener(&sock_path) else {
328            return;
329        };
330
331        let sock_path_clone = sock_path.clone();
332        let server = tokio::spawn(async move {
333            let (mut stream, _) = listener.accept().await.unwrap();
334            // Read a frame: [type:1][len:4][payload]
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            if len > 0 {
341                stream.read_exact(&mut payload).await.unwrap();
342            }
343            // Echo it back
344            stream.write_all(&header).await.unwrap();
345            stream.write_all(&payload).await.unwrap();
346            (frame_type, payload)
347        });
348
349        let mut client = PtyClient::connect(&sock_path_clone).await.unwrap();
350        client.send_data(b"hello world").await.unwrap();
351
352        let frame = client.read_frame().await.unwrap().unwrap();
353        assert_eq!(frame.0, a3s_box_core::pty::FRAME_PTY_DATA);
354        assert_eq!(&frame.1[..], b"hello world");
355
356        let (server_type, server_payload) = server.await.unwrap();
357        assert_eq!(server_type, a3s_box_core::pty::FRAME_PTY_DATA);
358        assert_eq!(&server_payload[..], b"hello world");
359    }
360
361    #[tokio::test]
362    async fn test_pty_send_request_serializes_payload() {
363        let tmp = tempfile::TempDir::new().unwrap();
364        let sock_path = tmp.path().join("pty_request.sock");
365        let Some(listener) = bind_test_listener(&sock_path) else {
366            return;
367        };
368
369        let server = tokio::spawn(async move {
370            let (stream, _) = listener.accept().await.unwrap();
371            let (r, _w) = tokio::io::split(stream);
372            let mut reader = a3s_transport::FrameReader::new(r);
373            reader.read_frame().await.unwrap().unwrap()
374        });
375
376        let mut client = PtyClient::connect(&sock_path).await.unwrap();
377        let req = test_pty_request();
378        client.send_request(&req).await.unwrap();
379
380        let frame = server.await.unwrap();
381        assert_eq!(frame.frame_type as u8, a3s_box_core::pty::FRAME_PTY_REQUEST);
382        let parsed: a3s_box_core::pty::PtyRequest = serde_json::from_slice(&frame.payload).unwrap();
383        assert_eq!(parsed.cmd, req.cmd);
384        assert_eq!(parsed.env, req.env);
385        assert_eq!(parsed.working_dir, req.working_dir);
386        assert_eq!(parsed.rootfs, req.rootfs);
387        assert_eq!(parsed.user, req.user);
388        assert_eq!(parsed.cols, 100);
389        assert_eq!(parsed.rows, 30);
390    }
391
392    #[tokio::test]
393    async fn test_pty_send_resize() {
394        let tmp = tempfile::TempDir::new().unwrap();
395        let sock_path = tmp.path().join("pty_resize.sock");
396        let Some(listener) = bind_test_listener(&sock_path) else {
397            return;
398        };
399
400        tokio::spawn(async move {
401            let (mut stream, _) = listener.accept().await.unwrap();
402            let mut header = [0u8; 5];
403            stream.read_exact(&mut header).await.unwrap();
404            let frame_type = header[0];
405            let len = u32::from_be_bytes([header[1], header[2], header[3], header[4]]) as usize;
406            let mut payload = vec![0u8; len];
407            stream.read_exact(&mut payload).await.unwrap();
408
409            assert_eq!(frame_type, a3s_box_core::pty::FRAME_PTY_RESIZE);
410            let resize: a3s_box_core::pty::PtyResize = serde_json::from_slice(&payload).unwrap();
411            assert_eq!(resize.cols, 120);
412            assert_eq!(resize.rows, 40);
413        });
414
415        let mut client = PtyClient::connect(&sock_path).await.unwrap();
416        client.send_resize(120, 40).await.unwrap();
417    }
418
419    #[tokio::test]
420    async fn test_pty_stream_input_writes_stdin_resize_and_signal_frames() {
421        let tmp = tempfile::TempDir::new().unwrap();
422        let sock_path = tmp.path().join("pty_input.sock");
423        let Some(listener) = bind_test_listener(&sock_path) else {
424            return;
425        };
426
427        let server = tokio::spawn(async move {
428            let (stream, _) = listener.accept().await.unwrap();
429            let (r, _w) = tokio::io::split(stream);
430            let mut reader = a3s_transport::FrameReader::new(r);
431
432            let request = reader.read_frame().await.unwrap().unwrap();
433            assert_eq!(
434                request.frame_type as u8,
435                a3s_box_core::pty::FRAME_PTY_REQUEST
436            );
437            let stdin = reader.read_frame().await.unwrap().unwrap();
438            let resize = reader.read_frame().await.unwrap().unwrap();
439            let terminate = reader.read_frame().await.unwrap().unwrap();
440            let kill = reader.read_frame().await.unwrap().unwrap();
441            (stdin, resize, terminate, kill)
442        });
443
444        let client = PtyClient::connect(&sock_path).await.unwrap();
445        let stream = client.start_stream(&test_pty_request()).await.unwrap();
446        let input = stream.input();
447        input.write_stdin(b"echo hi\n").await.unwrap();
448        input.resize(132, 43).await.unwrap();
449        input
450            .send_signal(ExecutionProcessSignal::Terminate)
451            .await
452            .unwrap();
453        input
454            .send_signal(ExecutionProcessSignal::Kill)
455            .await
456            .unwrap();
457
458        let (stdin, resize, terminate, kill) = server.await.unwrap();
459        assert_eq!(stdin.frame_type, a3s_transport::FrameType::Control);
460        assert_eq!(stdin.payload, b"echo hi\n");
461
462        assert_eq!(resize.frame_type, a3s_transport::FrameType::Heartbeat);
463        let resize: a3s_box_core::pty::PtyResize = serde_json::from_slice(&resize.payload).unwrap();
464        assert_eq!(resize.cols, 132);
465        assert_eq!(resize.rows, 43);
466
467        assert_eq!(terminate.frame_type, a3s_transport::FrameType::Close);
468        assert_eq!(terminate.payload, b"signal:15");
469        assert_eq!(kill.frame_type, a3s_transport::FrameType::Close);
470        assert!(kill.payload.is_empty());
471    }
472
473    #[tokio::test]
474    async fn test_pty_read_frame_eof() {
475        let tmp = tempfile::TempDir::new().unwrap();
476        let sock_path = tmp.path().join("pty_eof.sock");
477        let Some(listener) = bind_test_listener(&sock_path) else {
478            return;
479        };
480
481        tokio::spawn(async move {
482            let (stream, _) = listener.accept().await.unwrap();
483            drop(stream); // Close immediately → EOF
484        });
485
486        let mut client = PtyClient::connect(&sock_path).await.unwrap();
487        let frame = client.read_frame().await.unwrap();
488        assert!(frame.is_none()); // EOF
489    }
490
491    #[tokio::test]
492    async fn test_pty_stream_eof_marks_done_and_metrics_count_stdout() {
493        let tmp = tempfile::TempDir::new().unwrap();
494        let sock_path = tmp.path().join("pty_stream_eof.sock");
495        let Some(listener) = bind_test_listener(&sock_path) else {
496            return;
497        };
498
499        tokio::spawn(async move {
500            let (stream, _) = listener.accept().await.unwrap();
501            let (r, w) = tokio::io::split(stream);
502            let mut reader = a3s_transport::FrameReader::new(r);
503            let mut writer = a3s_transport::FrameWriter::new(w);
504
505            let request = reader.read_frame().await.unwrap().unwrap();
506            assert_eq!(
507                request.frame_type as u8,
508                a3s_box_core::pty::FRAME_PTY_REQUEST
509            );
510            writer
511                .write_frame(&a3s_transport::Frame {
512                    frame_type: a3s_transport::FrameType::Control,
513                    payload: b"abc".to_vec(),
514                })
515                .await
516                .unwrap();
517        });
518
519        let client = PtyClient::connect(&sock_path).await.unwrap();
520        let mut stream = client.start_stream(&test_pty_request()).await.unwrap();
521        assert!(!stream.is_done());
522
523        match stream.next_event().await.unwrap().unwrap() {
524            a3s_box_core::exec::ExecEvent::Chunk(chunk) => assert_eq!(chunk.data, b"abc"),
525            other => panic!("unexpected event: {other:?}"),
526        }
527        assert_eq!(stream.metrics().stdout_bytes, 3);
528
529        assert!(stream.next_event().await.unwrap().is_none());
530        assert!(stream.is_done());
531        assert!(stream.next_event().await.unwrap().is_none());
532    }
533
534    #[tokio::test]
535    async fn test_pty_client_start_stream_reads_data_and_exit() {
536        let tmp = tempfile::TempDir::new().unwrap();
537        let sock_path = tmp.path().join("pty_stream.sock");
538        let Some(listener) = bind_test_listener(&sock_path) else {
539            return;
540        };
541
542        tokio::spawn(async move {
543            let (stream, _) = listener.accept().await.unwrap();
544            let (r, w) = tokio::io::split(stream);
545            let mut reader = a3s_transport::FrameReader::new(r);
546            let mut writer = a3s_transport::FrameWriter::new(w);
547
548            let request = reader.read_frame().await.unwrap().unwrap();
549            assert_eq!(
550                request.frame_type as u8,
551                a3s_box_core::pty::FRAME_PTY_REQUEST
552            );
553
554            writer
555                .write_frame(&a3s_transport::Frame {
556                    frame_type: a3s_transport::FrameType::Control,
557                    payload: b"tty output".to_vec(),
558                })
559                .await
560                .unwrap();
561
562            let exit = a3s_box_core::pty::PtyExit { exit_code: 9 };
563            writer
564                .write_frame(&a3s_transport::Frame {
565                    frame_type: a3s_transport::FrameType::Error,
566                    payload: serde_json::to_vec(&exit).unwrap(),
567                })
568                .await
569                .unwrap();
570        });
571
572        let client = PtyClient::connect(&sock_path).await.unwrap();
573        let req = a3s_box_core::pty::PtyRequest {
574            cmd: vec!["/bin/sh".to_string()],
575            env: vec![],
576            working_dir: None,
577            rootfs: None,
578            user: None,
579            cols: 80,
580            rows: 24,
581        };
582        let mut stream = client.start_stream(&req).await.unwrap();
583
584        match stream.next_event().await.unwrap().unwrap() {
585            a3s_box_core::exec::ExecEvent::Chunk(chunk) => {
586                assert_eq!(chunk.stream, a3s_box_core::exec::StreamType::Stdout);
587                assert_eq!(chunk.data, b"tty output");
588            }
589            other => panic!("unexpected event: {other:?}"),
590        }
591        match stream.next_event().await.unwrap().unwrap() {
592            a3s_box_core::exec::ExecEvent::Exit(exit) => assert_eq!(exit.exit_code, 9),
593            other => panic!("unexpected event: {other:?}"),
594        }
595    }
596
597    #[tokio::test]
598    async fn test_pty_stream_error_frame_marks_done() {
599        let tmp = tempfile::TempDir::new().unwrap();
600        let sock_path = tmp.path().join("pty_stream_error.sock");
601        let Some(listener) = bind_test_listener(&sock_path) else {
602            return;
603        };
604
605        tokio::spawn(async move {
606            let (stream, _) = listener.accept().await.unwrap();
607            let (r, w) = tokio::io::split(stream);
608            let mut reader = a3s_transport::FrameReader::new(r);
609            let mut writer = a3s_transport::FrameWriter::new(w);
610
611            let request = reader.read_frame().await.unwrap().unwrap();
612            assert_eq!(
613                request.frame_type as u8,
614                a3s_box_core::pty::FRAME_PTY_REQUEST
615            );
616            writer
617                .write_frame(&a3s_transport::Frame {
618                    frame_type: a3s_transport::FrameType::Close,
619                    payload: b"terminal failed".to_vec(),
620                })
621                .await
622                .unwrap();
623        });
624
625        let client = PtyClient::connect(&sock_path).await.unwrap();
626        let mut stream = client.start_stream(&test_pty_request()).await.unwrap();
627
628        let err = stream.next_event().await.unwrap_err();
629        assert!(err.to_string().contains("terminal failed"));
630        assert!(stream.is_done());
631        assert!(stream.next_event().await.unwrap().is_none());
632    }
633
634    #[tokio::test]
635    async fn test_pty_stream_unexpected_frame_type_errors_without_marking_done() {
636        let tmp = tempfile::TempDir::new().unwrap();
637        let sock_path = tmp.path().join("pty_stream_unexpected.sock");
638        let Some(listener) = bind_test_listener(&sock_path) else {
639            return;
640        };
641
642        tokio::spawn(async move {
643            let (stream, _) = listener.accept().await.unwrap();
644            let (r, w) = tokio::io::split(stream);
645            let mut reader = a3s_transport::FrameReader::new(r);
646            let mut writer = a3s_transport::FrameWriter::new(w);
647
648            let request = reader.read_frame().await.unwrap().unwrap();
649            assert_eq!(
650                request.frame_type as u8,
651                a3s_box_core::pty::FRAME_PTY_REQUEST
652            );
653            writer
654                .write_frame(&a3s_transport::Frame {
655                    frame_type: a3s_transport::FrameType::Data,
656                    payload: b"not a PTY data frame".to_vec(),
657                })
658                .await
659                .unwrap();
660        });
661
662        let client = PtyClient::connect(&sock_path).await.unwrap();
663        let mut stream = client.start_stream(&test_pty_request()).await.unwrap();
664
665        let err = stream.next_event().await.unwrap_err();
666        assert!(err.to_string().contains("Unexpected PTY frame type"));
667        assert!(!stream.is_done());
668    }
669
670    #[tokio::test]
671    async fn test_pty_stream_cancel_writes_close_frame() {
672        let tmp = tempfile::TempDir::new().unwrap();
673        let sock_path = tmp.path().join("pty_cancel.sock");
674        let Some(listener) = bind_test_listener(&sock_path) else {
675            return;
676        };
677
678        tokio::spawn(async move {
679            let (stream, _) = listener.accept().await.unwrap();
680            let (r, _w) = tokio::io::split(stream);
681            let mut reader = a3s_transport::FrameReader::new(r);
682
683            let request = reader.read_frame().await.unwrap().unwrap();
684            assert_eq!(
685                request.frame_type as u8,
686                a3s_box_core::pty::FRAME_PTY_REQUEST
687            );
688            let close = reader.read_frame().await.unwrap().unwrap();
689            assert_eq!(close.frame_type, a3s_transport::FrameType::Close);
690            assert!(close.payload.is_empty());
691        });
692
693        let client = PtyClient::connect(&sock_path).await.unwrap();
694        let req = a3s_box_core::pty::PtyRequest {
695            cmd: vec!["sleep".to_string(), "60".to_string()],
696            env: vec![],
697            working_dir: None,
698            rootfs: None,
699            user: None,
700            cols: 80,
701            rows: 24,
702        };
703        let mut stream = client.start_stream(&req).await.unwrap();
704        stream.cancel().await.unwrap();
705    }
706}
a3s_box_runtime/grpc/pty.rs

a3s_box_runtime/grpc/
pty.rs
