virtual_terminal/
lib.rs

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

virtual_terminal/
lib.rs
