1use std::io;
7use std::pin::Pin;
8use std::task::{Context, Poll};
9
10use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
11
12use crate::backend::ChildExit;
13use crate::config::SessionConfig;
14use crate::error::{ExpectError, Result, SpawnError};
15use crate::types::ProcessExitStatus;
16
17pub struct PtyTransport {
19 reader: Box<dyn AsyncRead + Unpin + Send>,
21 writer: Box<dyn AsyncWrite + Unpin + Send>,
23 pid: Option<u32>,
25}
26
27impl PtyTransport {
28 pub fn new<R, W>(reader: R, writer: W) -> Self
30 where
31 R: AsyncRead + Unpin + Send + 'static,
32 W: AsyncWrite + Unpin + Send + 'static,
33 {
34 Self {
35 reader: Box::new(reader),
36 writer: Box::new(writer),
37 pid: None,
38 }
39 }
40
41 pub const fn set_pid(&mut self, pid: u32) {
43 self.pid = Some(pid);
44 }
45
46 #[must_use]
48 pub const fn pid(&self) -> Option<u32> {
49 self.pid
50 }
51}
52
53impl AsyncRead for PtyTransport {
54 fn poll_read(
55 mut self: Pin<&mut Self>,
56 cx: &mut Context<'_>,
57 buf: &mut ReadBuf<'_>,
58 ) -> Poll<io::Result<()>> {
59 Pin::new(&mut self.reader).poll_read(cx, buf)
60 }
61}
62
63impl AsyncWrite for PtyTransport {
64 fn poll_write(
65 mut self: Pin<&mut Self>,
66 cx: &mut Context<'_>,
67 buf: &[u8],
68 ) -> Poll<io::Result<usize>> {
69 Pin::new(&mut self.writer).poll_write(cx, buf)
70 }
71
72 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
73 Pin::new(&mut self.writer).poll_flush(cx)
74 }
75
76 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
77 Pin::new(&mut self.writer).poll_shutdown(cx)
78 }
79}
80
81#[derive(Debug, Clone)]
83#[non_exhaustive]
84pub struct PtyConfig {
85 pub dimensions: (u16, u16),
87 pub login_shell: bool,
89 pub env_mode: EnvMode,
91 pub env: std::collections::HashMap<String, String>,
94 pub working_directory: Option<std::path::PathBuf>,
97}
98
99impl Default for PtyConfig {
100 fn default() -> Self {
101 Self {
102 dimensions: (80, 24),
103 login_shell: false,
104 env_mode: EnvMode::Inherit,
105 env: std::collections::HashMap::new(),
106 working_directory: None,
107 }
108 }
109}
110
111impl From<&SessionConfig> for PtyConfig {
112 fn from(config: &SessionConfig) -> Self {
113 Self {
114 dimensions: config.dimensions,
115 login_shell: false,
116 env_mode: match (config.inherit_env, config.env.is_empty()) {
117 (false, _) => EnvMode::Clear,
118 (true, true) => EnvMode::Inherit,
119 (true, false) => EnvMode::Extend,
120 },
121 env: config.env.clone(),
122 working_directory: config.working_dir.clone(),
123 }
124 }
125}
126
127#[derive(Debug, Clone, Copy, PartialEq, Eq)]
129pub enum EnvMode {
130 Inherit,
132 Clear,
134 Extend,
136}
137
138pub struct PtySpawner {
140 config: PtyConfig,
141}
142
143impl PtySpawner {
144 #[must_use]
146 pub fn new() -> Self {
147 Self {
148 config: PtyConfig::default(),
149 }
150 }
151
152 #[must_use]
154 pub const fn with_config(config: PtyConfig) -> Self {
155 Self { config }
156 }
157
158 pub const fn set_dimensions(&mut self, cols: u16, rows: u16) {
160 self.config.dimensions = (cols, rows);
161 }
162
163 #[cfg(unix)]
176 pub async fn spawn(&self, command: &str, args: &[String]) -> Result<PtyHandle> {
177 use rust_pty::{PtySystem, UnixPtySystem};
178
179 if let Some(dir) = &self.config.working_directory
183 && !dir.is_dir()
184 {
185 return Err(ExpectError::Spawn(SpawnError::InvalidWorkingDir {
186 path: dir.display().to_string(),
187 }));
188 }
189
190 let built_env: Option<std::collections::HashMap<std::ffi::OsString, std::ffi::OsString>> =
195 match self.config.env_mode {
196 EnvMode::Inherit if self.config.env.is_empty() => None,
197 EnvMode::Inherit | EnvMode::Extend => {
198 let mut m: std::collections::HashMap<_, _> = std::env::vars_os().collect();
199 for (k, v) in &self.config.env {
200 m.insert(std::ffi::OsString::from(k), std::ffi::OsString::from(v));
201 }
202 Some(m)
203 }
204 EnvMode::Clear => Some(
205 self.config
206 .env
207 .iter()
208 .map(|(k, v)| (std::ffi::OsString::from(k), std::ffi::OsString::from(v)))
209 .collect(),
210 ),
211 };
212
213 let pty_config = rust_pty::PtyConfig {
214 window_size: self.config.dimensions,
215 env: match self.config.env_mode {
216 EnvMode::Clear if self.config.env.is_empty() => {
217 Some(std::collections::HashMap::new())
218 }
219 _ => built_env,
220 },
221 working_directory: self.config.working_directory.clone(),
222 ..Default::default()
223 };
224
225 let (master, child) =
226 UnixPtySystem::spawn(command, args.iter().map(String::as_str), &pty_config)
227 .await
228 .map_err(|e| {
229 ExpectError::Spawn(SpawnError::PtyAllocation {
230 reason: format!("Unix PTY spawn failed: {e}"),
231 })
232 })?;
233
234 Ok(PtyHandle {
235 master,
236 child,
237 dimensions: self.config.dimensions,
238 })
239 }
240
241 #[cfg(windows)]
250 pub async fn spawn(&self, command: &str, args: &[String]) -> Result<WindowsPtyHandle> {
251 use rust_pty::{PtySystem, WindowsPtySystem};
252
253 let built_env: Option<std::collections::HashMap<std::ffi::OsString, std::ffi::OsString>> =
259 match self.config.env_mode {
260 EnvMode::Inherit if self.config.env.is_empty() => None,
261 EnvMode::Inherit | EnvMode::Extend => {
262 let mut m: std::collections::HashMap<_, _> = std::env::vars_os().collect();
263 for (k, v) in &self.config.env {
264 m.insert(std::ffi::OsString::from(k), std::ffi::OsString::from(v));
265 }
266 Some(m)
267 }
268 EnvMode::Clear => Some(
269 self.config
270 .env
271 .iter()
272 .map(|(k, v)| (std::ffi::OsString::from(k), std::ffi::OsString::from(v)))
273 .collect(),
274 ),
275 };
276
277 let pty_config = rust_pty::PtyConfig {
279 window_size: self.config.dimensions,
280 env: match self.config.env_mode {
281 EnvMode::Clear if self.config.env.is_empty() => {
282 Some(std::collections::HashMap::new())
283 }
284 _ => built_env,
285 },
286 working_directory: self.config.working_directory.clone(),
287 ..Default::default()
288 };
289
290 let (master, child) =
292 WindowsPtySystem::spawn(command, args.iter().map(|s| s.as_str()), &pty_config)
293 .await
294 .map_err(|e| {
295 ExpectError::Spawn(SpawnError::PtyAllocation {
296 reason: format!("Windows ConPTY spawn failed: {e}"),
297 })
298 })?;
299
300 Ok(WindowsPtyHandle {
301 master,
302 child,
303 dimensions: self.config.dimensions,
304 })
305 }
306}
307
308impl Default for PtySpawner {
309 fn default() -> Self {
310 Self::new()
311 }
312}
313
314#[cfg(unix)]
316#[derive(Debug)]
317pub struct PtyHandle {
318 pub(crate) master: rust_pty::UnixPtyMaster,
320 pub(crate) child: rust_pty::UnixPtyChild,
322 dimensions: (u16, u16),
324}
325
326#[cfg(windows)]
328pub struct WindowsPtyHandle {
329 pub(crate) master: rust_pty::WindowsPtyMaster,
331 pub(crate) child: rust_pty::WindowsPtyChild,
333 dimensions: (u16, u16),
335}
336
337#[cfg(windows)]
338impl std::fmt::Debug for WindowsPtyHandle {
339 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
340 f.debug_struct("WindowsPtyHandle")
341 .field("dimensions", &self.dimensions)
342 .finish_non_exhaustive()
343 }
344}
345
346#[cfg(unix)]
347impl PtyHandle {
348 #[must_use]
350 pub const fn pid(&self) -> u32 {
351 self.child.pid()
352 }
353
354 #[must_use]
356 pub const fn dimensions(&self) -> (u16, u16) {
357 self.dimensions
358 }
359
360 pub fn resize(&mut self, cols: u16, rows: u16) -> Result<()> {
362 use rust_pty::{PtyMaster, WindowSize};
363 self.master
364 .resize(WindowSize::new(cols, rows))
365 .map_err(|e| ExpectError::Io(io::Error::other(format!("resize failed: {e}"))))?;
366 self.dimensions = (cols, rows);
367 Ok(())
368 }
369
370 }
375
376#[cfg(windows)]
377impl WindowsPtyHandle {
378 #[must_use]
380 pub fn pid(&self) -> u32 {
381 self.child.pid()
382 }
383
384 #[must_use]
386 pub const fn dimensions(&self) -> (u16, u16) {
387 self.dimensions
388 }
389
390 pub fn resize(&mut self, cols: u16, rows: u16) -> Result<()> {
392 use rust_pty::{PtyMaster, WindowSize};
393 let size = WindowSize::new(cols, rows);
394 self.master
395 .resize(size)
396 .map_err(|e| ExpectError::Io(io::Error::other(format!("resize failed: {e}"))))?;
397 self.dimensions = (cols, rows);
398 Ok(())
399 }
400
401 #[must_use]
403 pub fn is_running(&self) -> bool {
404 self.child.is_running()
405 }
406
407 pub fn kill(&mut self) -> Result<()> {
409 self.child
410 .kill()
411 .map_err(|e| ExpectError::Io(io::Error::other(format!("kill failed: {e}"))))
412 }
413}
414
415#[cfg(unix)]
420pub struct AsyncPty {
421 master: rust_pty::UnixPtyMaster,
423 child: rust_pty::UnixPtyChild,
425 pid: u32,
427 dimensions: (u16, u16),
429}
430
431#[cfg(unix)]
432impl AsyncPty {
433 pub fn from_handle(handle: PtyHandle) -> io::Result<Self> {
441 let pid = handle.child.pid();
442 let dimensions = handle.dimensions;
443 Ok(Self {
444 master: handle.master,
445 child: handle.child,
446 pid,
447 dimensions,
448 })
449 }
450
451 pub fn try_wait(&mut self) -> Option<ProcessExitStatus> {
457 match self.child.try_wait() {
458 Ok(Some(rust_pty::ExitStatus::Exited(code))) => Some(ProcessExitStatus::Exited(code)),
459 Ok(Some(rust_pty::ExitStatus::Signaled(sig))) => Some(ProcessExitStatus::Signaled(sig)),
460 Ok(None) | Err(_) => None,
461 }
462 }
463
464 pub fn is_running(&mut self) -> bool {
469 self.try_wait().is_none()
470 }
471
472 #[must_use]
474 pub const fn pid(&self) -> u32 {
475 self.pid
476 }
477
478 #[must_use]
480 pub const fn dimensions(&self) -> (u16, u16) {
481 self.dimensions
482 }
483
484 pub fn resize(&mut self, cols: u16, rows: u16) -> Result<()> {
486 use rust_pty::{PtyMaster, WindowSize};
487 self.master
488 .resize(WindowSize::new(cols, rows))
489 .map_err(|e| ExpectError::Io(io::Error::other(format!("resize failed: {e}"))))?;
490 self.dimensions = (cols, rows);
491 Ok(())
492 }
493
494 #[allow(unsafe_code)]
505 pub fn signal(&mut self, signal: i32) -> Result<()> {
506 if self.try_wait().is_some() {
508 return Err(ExpectError::SessionClosed);
509 }
510 let result = unsafe { libc::kill(self.pid as i32, signal) };
512 if result == 0 {
513 Ok(())
514 } else {
515 let err = io::Error::last_os_error();
516 if err.raw_os_error() == Some(libc::ESRCH) {
519 Err(ExpectError::SessionClosed)
520 } else {
521 Err(ExpectError::Io(err))
522 }
523 }
524 }
525
526 pub fn kill(&mut self) -> Result<()> {
528 self.signal(libc::SIGKILL)
529 }
530}
531
532#[cfg(unix)]
533impl AsyncRead for AsyncPty {
534 fn poll_read(
535 mut self: Pin<&mut Self>,
536 cx: &mut Context<'_>,
537 buf: &mut ReadBuf<'_>,
538 ) -> Poll<io::Result<()>> {
539 Pin::new(&mut self.master).poll_read(cx, buf)
540 }
541}
542
543#[cfg(unix)]
544impl AsyncWrite for AsyncPty {
545 fn poll_write(
546 mut self: Pin<&mut Self>,
547 cx: &mut Context<'_>,
548 buf: &[u8],
549 ) -> Poll<io::Result<usize>> {
550 if matches!(self.child.try_wait(), Ok(Some(_))) {
552 return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()));
553 }
554 Pin::new(&mut self.master).poll_write(cx, buf)
555 }
556
557 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
558 Pin::new(&mut self.master).poll_flush(cx)
559 }
560
561 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
562 Pin::new(&mut self.master).poll_shutdown(cx)
563 }
564}
565
566#[cfg(unix)]
567impl std::fmt::Debug for AsyncPty {
568 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
569 f.debug_struct("AsyncPty")
570 .field("pid", &self.pid)
571 .field("dimensions", &self.dimensions)
572 .finish_non_exhaustive()
573 }
574}
575
576#[cfg(unix)]
577impl ChildExit for AsyncPty {
578 fn try_exit_status(&mut self) -> Option<ProcessExitStatus> {
579 self.try_wait()
580 }
581}
582
583#[cfg(windows)]
588pub struct WindowsAsyncPty {
589 master: rust_pty::WindowsPtyMaster,
591 child: rust_pty::WindowsPtyChild,
593 pid: u32,
595 dimensions: (u16, u16),
597}
598
599#[cfg(windows)]
600impl WindowsAsyncPty {
601 pub fn from_handle(handle: WindowsPtyHandle) -> Self {
605 let pid = handle.child.pid();
606 let dimensions = handle.dimensions;
607 Self {
608 master: handle.master,
609 child: handle.child,
610 pid,
611 dimensions,
612 }
613 }
614
615 #[must_use]
617 pub const fn pid(&self) -> u32 {
618 self.pid
619 }
620
621 #[must_use]
623 pub const fn dimensions(&self) -> (u16, u16) {
624 self.dimensions
625 }
626
627 pub fn resize(&mut self, cols: u16, rows: u16) -> Result<()> {
629 use rust_pty::{PtyMaster, WindowSize};
630 let size = WindowSize::new(cols, rows);
631 self.master
632 .resize(size)
633 .map_err(|e| ExpectError::Io(io::Error::other(format!("resize failed: {e}"))))?;
634 self.dimensions = (cols, rows);
635 Ok(())
636 }
637
638 #[must_use]
640 pub fn is_running(&self) -> bool {
641 self.child.is_running()
642 }
643
644 pub fn kill(&mut self) -> Result<()> {
646 self.child
647 .kill()
648 .map_err(|e| ExpectError::Io(io::Error::other(format!("kill failed: {e}"))))
649 }
650}
651
652#[cfg(windows)]
653impl ChildExit for WindowsAsyncPty {
654 fn try_exit_status(&mut self) -> Option<ProcessExitStatus> {
655 match self.child.try_wait() {
660 Ok(Some(rust_pty::ExitStatus::Exited(code))) => Some(ProcessExitStatus::Exited(code)),
661 Ok(Some(rust_pty::ExitStatus::Terminated(code))) => {
663 Some(ProcessExitStatus::Exited(code as i32))
664 }
665 Ok(None) | Err(_) => None,
667 }
668 }
669}
670
671#[cfg(windows)]
672impl AsyncRead for WindowsAsyncPty {
673 fn poll_read(
674 mut self: Pin<&mut Self>,
675 cx: &mut Context<'_>,
676 buf: &mut ReadBuf<'_>,
677 ) -> Poll<io::Result<()>> {
678 Pin::new(&mut self.master).poll_read(cx, buf)
680 }
681}
682
683#[cfg(windows)]
684impl AsyncWrite for WindowsAsyncPty {
685 fn poll_write(
686 mut self: Pin<&mut Self>,
687 cx: &mut Context<'_>,
688 buf: &[u8],
689 ) -> Poll<io::Result<usize>> {
690 if matches!(self.child.try_wait(), Ok(Some(_))) {
692 return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()));
693 }
694 Pin::new(&mut self.master).poll_write(cx, buf)
695 }
696
697 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
698 Pin::new(&mut self.master).poll_flush(cx)
699 }
700
701 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
702 Pin::new(&mut self.master).poll_shutdown(cx)
703 }
704}
705
706#[cfg(windows)]
707impl std::fmt::Debug for WindowsAsyncPty {
708 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
709 f.debug_struct("WindowsAsyncPty")
710 .field("pid", &self.pid)
711 .field("dimensions", &self.dimensions)
712 .finish_non_exhaustive()
713 }
714}
715
716#[cfg(test)]
717mod tests {
718 use super::*;
719
720 #[test]
721 fn pty_config_default() {
722 let config = PtyConfig::default();
723 assert_eq!(config.dimensions.0, 80);
724 assert_eq!(config.dimensions.1, 24);
725 assert_eq!(config.env_mode, EnvMode::Inherit);
726 }
727
728 #[test]
729 fn pty_config_from_session() {
730 let session_config = SessionConfig {
731 dimensions: (120, 40),
732 ..Default::default()
733 };
734
735 let pty_config = PtyConfig::from(&session_config);
736 assert_eq!(pty_config.dimensions.0, 120);
737 assert_eq!(pty_config.dimensions.1, 40);
738 }
739
740 #[cfg(unix)]
741 #[tokio::test]
742 async fn spawn_rejects_null_byte_in_command() {
743 let spawner = PtySpawner::new();
744 let result = spawner.spawn("test\0command", &[]).await;
745
746 assert!(result.is_err());
747 let err = result.unwrap_err();
748 let err_str = err.to_string();
749 assert!(
750 err_str.contains("nul byte"),
751 "Expected error about a nul byte, got: {err_str}"
752 );
753 }
754
755 #[cfg(unix)]
756 #[tokio::test]
757 async fn spawn_rejects_null_byte_in_args() {
758 let spawner = PtySpawner::new();
759 let result = spawner
760 .spawn("/bin/echo", &["hello\0world".to_string()])
761 .await;
762
763 assert!(result.is_err());
764 let err = result.unwrap_err();
765 let err_str = err.to_string();
766 assert!(
767 err_str.contains("nul byte"),
768 "Expected error about a nul byte, got: {err_str}"
769 );
770 }
771}