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 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 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 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 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}