caco3_pty/
nixpty.rs

1use nix::fcntl::OFlag;
2use nix::pty::PtyMaster;
3use std::fs::File;
4use std::io::{self, Read, Write};
5use std::os::fd::AsFd;
6use std::task::{ready, Poll};
7use tokio::io::unix::AsyncFd;
8use tokio::process::Child;
9
10use crate::sys::get_child_terminal_path;
11use crate::{sys, AllocateError, ResizeError, SpawnError};
12
13pub struct PtyPair {
14    pty_master: PtyMaster,
15    child_pty: File,
16}
17
18impl PtyPair {
19    /// Allocate a new pseudo terminal with file descriptors for the parent and child end of the terminal.
20    fn new() -> Result<Self, AllocateError> {
21        let pty_master = nix::pty::posix_openpt(
22            OFlag::O_RDWR | OFlag::O_NOCTTY | OFlag::O_NONBLOCK | OFlag::O_CLOEXEC,
23        )
24        .map_err(io::Error::from)
25        .map_err(AllocateError::Open)?;
26        nix::pty::grantpt(&pty_master)
27            .map_err(io::Error::from)
28            .map_err(AllocateError::Grant)?;
29        nix::pty::unlockpt(&pty_master)
30            .map_err(io::Error::from)
31            .map_err(AllocateError::Unlock)?;
32        let child_pty_path =
33            get_child_terminal_path(&pty_master).map_err(AllocateError::GetChildName)?;
34        let child_pty = std::fs::OpenOptions::new()
35            .read(true)
36            .write(true)
37            .open(child_pty_path)
38            .map_err(AllocateError::OpenChild)?;
39        Ok(Self {
40            pty_master,
41            child_pty,
42        })
43    }
44
45    /// Spawn a child process as the session leader of a new process group with the pseudo terminal as controlling terminal.
46    ///
47    /// Also returns the parent side of the pseudo terminal as [`PseudoTerminal`] object.
48    pub async fn spawn(
49        self,
50        mut command: tokio::process::Command,
51    ) -> Result<(PseudoTerminal, Child), SpawnError> {
52        let Self {
53            pty_master,
54            child_pty: child_tty_file,
55        } = self;
56        let stdin = child_tty_file;
57        let stdout = stdin.try_clone().map_err(SpawnError::DuplicateStdio)?;
58        let stderr = stdin.try_clone().map_err(SpawnError::DuplicateStdio)?;
59        command.stdin(stdin);
60        command.stdout(stdout);
61        command.stderr(stderr);
62
63        unsafe {
64            command.pre_exec(move || {
65                sys::create_process_group()
66                    .map_err(SpawnError::CreateSession)
67                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
68                sys::set_controlling_terminal_to_stdin()
69                    .map_err(SpawnError::SetControllingTerminal)
70                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
71                Ok(())
72            });
73        };
74        let child = command.spawn().map_err(SpawnError::Spawn)?;
75        let pty = PseudoTerminal::new(pty_master)?;
76        Ok((pty, child))
77    }
78}
79
80pub struct PseudoTerminal {
81    inner: AsyncFd<PtyMaster>,
82}
83
84impl PseudoTerminal {
85    /// Allocate a new pseudo terminal.
86    pub fn allocate() -> Result<PtyPair, AllocateError> {
87        PtyPair::new()
88    }
89
90    fn new(pty_master: PtyMaster) -> Result<Self, SpawnError> {
91        Ok(Self {
92            inner: AsyncFd::new(pty_master).map_err(SpawnError::WrapAsyncFd)?,
93        })
94    }
95
96    /// Resize the pseudo-terminal.
97    ///
98    /// Should be called when the terminal emulator changes size.
99    pub fn resize(&self, width: u32, height: u32) -> Result<(), ResizeError> {
100        sys::resize_pty(self.inner.as_fd(), width, height).map_err(ResizeError)
101    }
102}
103
104impl tokio::io::AsyncRead for PseudoTerminal {
105    fn poll_read(
106        self: std::pin::Pin<&mut Self>,
107        cx: &mut std::task::Context<'_>,
108        buf: &mut tokio::io::ReadBuf<'_>,
109    ) -> Poll<io::Result<()>> {
110        poll_read_impl(&self.inner, cx, buf)
111    }
112}
113
114impl tokio::io::AsyncRead for &PseudoTerminal {
115    fn poll_read(
116        self: std::pin::Pin<&mut Self>,
117        cx: &mut std::task::Context<'_>,
118        buf: &mut tokio::io::ReadBuf<'_>,
119    ) -> Poll<io::Result<()>> {
120        poll_read_impl(&self.inner, cx, buf)
121    }
122}
123
124fn poll_read_impl(
125    fd: &AsyncFd<PtyMaster>,
126    cx: &mut std::task::Context<'_>,
127    buf: &mut tokio::io::ReadBuf<'_>,
128) -> Poll<io::Result<()>> {
129    loop {
130        let mut guard = ready!(fd.poll_read_ready(cx))?;
131
132        let unfilled = buf.initialize_unfilled();
133        match guard.try_io(|inner| inner.get_ref().read(unfilled)) {
134            Ok(Ok(len)) => {
135                buf.advance(len);
136                return Poll::Ready(Ok(()));
137            }
138            Ok(Err(err)) => return Poll::Ready(Err(err)),
139            Err(_would_block) => continue,
140        }
141    }
142}
143
144impl tokio::io::AsyncWrite for PseudoTerminal {
145    fn poll_write(
146        self: std::pin::Pin<&mut Self>,
147        cx: &mut std::task::Context<'_>,
148        buf: &[u8],
149    ) -> Poll<Result<usize, io::Error>> {
150        poll_write_impl(&self.inner, cx, buf)
151    }
152
153    fn poll_flush(
154        self: std::pin::Pin<&mut Self>,
155        _cx: &mut std::task::Context<'_>,
156    ) -> Poll<Result<(), io::Error>> {
157        Poll::Ready(Ok(()))
158    }
159
160    fn poll_shutdown(
161        self: std::pin::Pin<&mut Self>,
162        _cx: &mut std::task::Context<'_>,
163    ) -> Poll<Result<(), io::Error>> {
164        Poll::Ready(Ok(()))
165    }
166}
167
168impl tokio::io::AsyncWrite for &PseudoTerminal {
169    fn poll_write(
170        self: std::pin::Pin<&mut Self>,
171        cx: &mut std::task::Context<'_>,
172        buf: &[u8],
173    ) -> Poll<Result<usize, io::Error>> {
174        poll_write_impl(&self.inner, cx, buf)
175    }
176
177    fn poll_flush(
178        self: std::pin::Pin<&mut Self>,
179        _cx: &mut std::task::Context<'_>,
180    ) -> Poll<Result<(), io::Error>> {
181        Poll::Ready(Ok(()))
182    }
183
184    fn poll_shutdown(
185        self: std::pin::Pin<&mut Self>,
186        _cx: &mut std::task::Context<'_>,
187    ) -> Poll<Result<(), io::Error>> {
188        Poll::Ready(Ok(()))
189    }
190}
191
192fn poll_write_impl(
193    fd: &AsyncFd<PtyMaster>,
194    cx: &mut std::task::Context<'_>,
195    buf: &[u8],
196) -> Poll<Result<usize, io::Error>> {
197    loop {
198        let mut guard = ready!(fd.poll_write_ready(cx))?;
199        match guard.try_io(|inner| inner.get_ref().write(buf)) {
200            Ok(result) => return Poll::Ready(result),
201            Err(_would_block) => continue,
202        }
203    }
204}