virtual_terminal/
lib.rs

1use libc::TIOCSCTTY;
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeMap;
4use std::ffi::{OsStr, OsString};
5use std::os::fd::AsRawFd as _;
6use std::path::{Path, PathBuf};
7use std::time::Duration;
8use tokio::fs::File;
9use tokio::io::{AsyncReadExt, AsyncWriteExt};
10
11#[derive(Debug, Serialize)]
12#[serde(rename_all = "lowercase")]
13pub enum Output {
14    Pid(u32),
15    Stdout(Vec<u8>),
16    Error(String),
17    Terminated(Option<i32>),
18}
19
20#[derive(Debug, Deserialize, Eq, PartialEq)]
21pub enum Input {
22    Data(Vec<u8>),
23    Resize((usize, usize)),
24    Terminate,
25}
26
27const BUF_SIZE: usize = 8192;
28
29fn set_term_size(
30    fd: i32,
31    term_size: (usize, usize),
32) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
33    let ws = libc::winsize {
34        ws_row: u16::try_from(term_size.1)?,
35        ws_col: u16::try_from(term_size.0)?,
36        ws_xpixel: 0,
37        ws_ypixel: 0,
38    };
39
40    if unsafe { libc::ioctl(fd, libc::TIOCSWINSZ, &ws) } != 0 {
41        return Err("ioctl".into());
42    }
43
44    Ok(())
45}
46
47pub struct Command {
48    pid: Option<u32>,
49    program: OsString,
50    args: Vec<OsString>,
51    env: BTreeMap<OsString, OsString>,
52    current_dir: Option<PathBuf>,
53    in_tx: async_channel::Sender<Input>,
54    in_rx: async_channel::Receiver<Input>,
55    out_tx: async_channel::Sender<Output>,
56    out_rx: async_channel::Receiver<Output>,
57    terminal_id: String,
58    terminal_size: (usize, usize),
59}
60
61impl Command {
62    pub fn new<S: AsRef<OsStr>>(program: S) -> Self {
63        let (in_tx, in_rx) = async_channel::bounded(BUF_SIZE);
64        let (out_tx, out_rx) = async_channel::bounded(BUF_SIZE);
65        Self {
66            pid: None,
67            program: program.as_ref().to_os_string(),
68            args: <_>::default(),
69            env: <_>::default(),
70            current_dir: None,
71            in_tx,
72            in_rx,
73            out_tx,
74            out_rx,
75            terminal_id: "screen-256color".to_string(),
76            terminal_size: (80, 24),
77        }
78    }
79    pub fn in_tx(&self) -> async_channel::Sender<Input> {
80        self.in_tx.clone()
81    }
82    pub fn out_rx(&self) -> async_channel::Receiver<Output> {
83        self.out_rx.clone()
84    }
85    pub fn terminal_id<S: Into<String>>(mut self, terminal_id: S) -> Self {
86        self.terminal_id = terminal_id.into();
87        self
88    }
89    pub fn terminal_size(mut self, terminal_size: (usize, usize)) -> Self {
90        self.terminal_size = terminal_size;
91        self
92    }
93    pub fn args<I, S>(mut self, args: I) -> Self
94    where
95        I: IntoIterator<Item = S>,
96        S: AsRef<OsStr>,
97    {
98        self.args = args
99            .into_iter()
100            .map(|s| s.as_ref().to_os_string())
101            .collect();
102        self
103    }
104    pub fn arg<S: AsRef<OsStr>>(mut self, arg: S) -> Self {
105        self.args.push(arg.as_ref().to_os_string());
106        self
107    }
108    pub fn envs<I, K, V>(mut self, env: I) -> Self
109    where
110        I: IntoIterator<Item = (K, V)>,
111        K: AsRef<OsStr>,
112        V: AsRef<OsStr>,
113    {
114        self.env = env
115            .into_iter()
116            .map(|(k, v)| (k.as_ref().to_os_string(), v.as_ref().to_os_string()))
117            .collect();
118        self
119    }
120    pub fn env<K: AsRef<OsStr>, V: AsRef<OsStr>>(mut self, key: K, value: V) -> Self {
121        self.env
122            .insert(key.as_ref().to_os_string(), value.as_ref().to_os_string());
123        self
124    }
125    pub fn current_dir<P: AsRef<Path>>(mut self, current_dir: P) -> Self {
126        self.current_dir = Some(current_dir.as_ref().to_path_buf());
127        self
128    }
129    #[allow(clippy::too_many_arguments)]
130    pub async fn run(self) {
131        let out_tx = self.out_tx.clone();
132        match self.run_subprocess().await {
133            Ok(v) => {
134                out_tx.send(Output::Terminated(v)).await.ok();
135            }
136            Err(e) => {
137                out_tx.send(Output::Error(e.to_string())).await.ok();
138            }
139        }
140    }
141
142    #[allow(clippy::too_many_arguments, clippy::too_many_lines)]
143    async fn run_subprocess(
144        mut self,
145    ) -> Result<Option<i32>, Box<dyn std::error::Error + Send + Sync + 'static>> {
146        let win_size = rustix::termios::Winsize {
147            ws_col: self.terminal_size.0.try_into()?,
148            ws_row: self.terminal_size.1.try_into()?,
149            ws_xpixel: 0,
150            ws_ypixel: 0,
151        };
152        let pty = rustix_openpty::openpty(None, Some(&win_size))?;
153        let (master, slave) = (pty.controller, pty.user);
154
155        let master_fd = master.as_raw_fd();
156        let slave_fd = slave.as_raw_fd();
157
158        if let Ok(mut termios) = rustix::termios::tcgetattr(&master) {
159            // Set character encoding to UTF-8.
160            termios
161                .input_modes
162                .set(rustix::termios::InputModes::IUTF8, true);
163            let _ = rustix::termios::tcsetattr(
164                &master,
165                rustix::termios::OptionalActions::Now,
166                &termios,
167            );
168        }
169
170        let mut builder = tokio::process::Command::new(&self.program);
171
172        if let Some(ref current_dir) = self.current_dir {
173            builder.current_dir(current_dir);
174        }
175
176        builder
177            .args(&self.args)
178            .envs(&self.env)
179            .env("COLUMNS", self.terminal_size.0.to_string())
180            .env("LINES", self.terminal_size.1.to_string())
181            .env("TERM", &self.terminal_id);
182
183        builder.stdin(slave.try_clone()?);
184        builder.stderr(slave.try_clone()?);
185        builder.stdout(slave);
186
187        unsafe {
188            builder.pre_exec(move || {
189                let err = libc::setsid();
190                if err == -1 {
191                    return Err(std::io::Error::new(
192                        std::io::ErrorKind::Other,
193                        "Failed to set session id",
194                    ));
195                }
196
197                let res = libc::ioctl(slave_fd, TIOCSCTTY as _, 0);
198                if res == -1 {
199                    return Err(std::io::Error::new(
200                        std::io::ErrorKind::Other,
201                        format!("Failed to set controlling terminal: {}", res),
202                    ));
203                }
204
205                libc::close(slave_fd);
206                libc::close(master_fd);
207
208                libc::signal(libc::SIGCHLD, libc::SIG_DFL);
209                libc::signal(libc::SIGHUP, libc::SIG_DFL);
210                libc::signal(libc::SIGINT, libc::SIG_DFL);
211                libc::signal(libc::SIGQUIT, libc::SIG_DFL);
212                libc::signal(libc::SIGTERM, libc::SIG_DFL);
213                libc::signal(libc::SIGALRM, libc::SIG_DFL);
214
215                Ok(())
216            });
217        }
218
219        let mut child = builder.spawn()?;
220
221        let pid = child.id().ok_or("unable to get child pid")?;
222
223        self.out_tx.send(Output::Pid(pid)).await?;
224        self.pid = Some(pid);
225
226        let mut stdout = File::from_std(std::fs::File::from(master));
227
228        let mut stdin = stdout.try_clone().await?;
229
230        let tx_stdout = self.out_tx.clone();
231
232        let fut_out = tokio::spawn(async move {
233            let mut buf = [0u8; BUF_SIZE];
234            while let Ok(b) = stdout.read(&mut buf).await {
235                if b == 0 {
236                    break;
237                }
238                if tx_stdout
239                    .send(Output::Stdout(buf[..b].to_vec()))
240                    .await
241                    .is_err()
242                {
243                    break;
244                }
245            }
246        });
247
248        let fut_in = tokio::spawn(async move {
249            while let Ok(input) = self.in_rx.recv().await {
250                let mut data = match input {
251                    Input::Data(d) => d,
252                    Input::Resize(size) => {
253                        set_term_size(stdin.as_raw_fd(), size).ok();
254                        bmart::process::kill_pstree_with_signal(
255                            pid,
256                            bmart::process::Signal::SIGWINCH,
257                            true,
258                        );
259                        continue;
260                    }
261                    Input::Terminate => {
262                        break;
263                    }
264                };
265                // TODO: remove this input hack
266                if data == [0x0a] {
267                    data[0] = 0x0d;
268                }
269                if stdin.write_all(&data).await.is_err() {
270                    break;
271                }
272            }
273        });
274
275        let result = child.wait().await?;
276
277        fut_out.abort();
278        fut_in.abort();
279
280        let exit_code = result.code();
281
282        Ok(exit_code)
283    }
284}
285
286impl Drop for Command {
287    fn drop(&mut self) {
288        if let Some(pid) = self.pid {
289            tokio::spawn(bmart::process::kill_pstree(
290                pid,
291                Some(Duration::from_secs(1)),
292                true,
293            ));
294        }
295    }
296}
virtual_terminal/lib.rs

virtual_terminal/
lib.rs
