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}
72
73impl Command {
74    /// Create a new command
75    pub fn new<S: AsRef<OsStr>>(program: S) -> Self {
76        let (in_tx, in_rx) = async_channel::bounded(BUF_SIZE);
77        let (out_tx, out_rx) = async_channel::bounded(BUF_SIZE);
78        Self {
79            pid: None,
80            program: program.as_ref().to_os_string(),
81            args: <_>::default(),
82            env: <_>::default(),
83            current_dir: None,
84            in_tx,
85            in_rx,
86            out_tx,
87            out_rx,
88            terminal_id: "screen-256color".to_string(),
89            terminal_size: (80, 24),
90        }
91    }
92    /// Get the sender for sending input to the command
93    pub fn in_tx(&self) -> async_channel::Sender<Input> {
94        self.in_tx.clone()
95    }
96    /// Get the output sender, useful e.g. to manually send empty output frames to the terminal handler
97    pub fn out_tx(&self) -> async_channel::Sender<Output> {
98        self.out_tx.clone()
99    }
100    /// Get the receiver for receiving output from the command
101    pub fn out_rx(&self) -> async_channel::Receiver<Output> {
102        self.out_rx.clone()
103    }
104    /// Set the terminal id
105    pub fn terminal_id<S: Into<String>>(mut self, terminal_id: S) -> Self {
106        self.terminal_id = terminal_id.into();
107        self
108    }
109    /// Set the terminal size
110    pub fn terminal_size(mut self, terminal_size: (usize, usize)) -> Self {
111        self.terminal_size = terminal_size;
112        self
113    }
114    /// Set the program arguments
115    pub fn args<I, S>(mut self, args: I) -> Self
116    where
117        I: IntoIterator<Item = S>,
118        S: AsRef<OsStr>,
119    {
120        self.args = args
121            .into_iter()
122            .map(|s| s.as_ref().to_os_string())
123            .collect();
124        self
125    }
126    /// Add a program argument
127    pub fn arg<S: AsRef<OsStr>>(mut self, arg: S) -> Self {
128        self.args.push(arg.as_ref().to_os_string());
129        self
130    }
131    /// Set the environment variables
132    pub fn envs<I, K, V>(mut self, env: I) -> Self
133    where
134        I: IntoIterator<Item = (K, V)>,
135        K: AsRef<OsStr>,
136        V: AsRef<OsStr>,
137    {
138        self.env = env
139            .into_iter()
140            .map(|(k, v)| (k.as_ref().to_os_string(), v.as_ref().to_os_string()))
141            .collect();
142        self
143    }
144    /// Add an environment variable
145    pub fn env<K: AsRef<OsStr>, V: AsRef<OsStr>>(mut self, key: K, value: V) -> Self {
146        self.env
147            .insert(key.as_ref().to_os_string(), value.as_ref().to_os_string());
148        self
149    }
150    /// Set the working directory
151    pub fn current_dir<P: AsRef<Path>>(mut self, current_dir: P) -> Self {
152        self.current_dir = Some(current_dir.as_ref().to_path_buf());
153        self
154    }
155    /// Run the command
156    pub async fn run(self) {
157        let out_tx = self.out_tx.clone();
158        match self.run_subprocess().await {
159            Ok(v) => {
160                out_tx.send(Output::Terminated(v)).await.ok();
161            }
162            Err(e) => {
163                out_tx.send(Output::Error(e.to_string())).await.ok();
164            }
165        }
166    }
167
168    #[allow(clippy::too_many_lines)]
169    async fn run_subprocess(
170        mut self,
171    ) -> Result<Option<i32>, Box<dyn std::error::Error + Send + Sync + 'static>> {
172        let win_size = rustix::termios::Winsize {
173            ws_col: self.terminal_size.0.try_into()?,
174            ws_row: self.terminal_size.1.try_into()?,
175            ws_xpixel: 0,
176            ws_ypixel: 0,
177        };
178        let pty = rustix_openpty::openpty(None, Some(&win_size))?;
179        let (master, slave) = (pty.controller, pty.user);
180
181        let master_fd = master.as_raw_fd();
182        let slave_fd = slave.as_raw_fd();
183
184        if let Ok(mut termios) = rustix::termios::tcgetattr(&master) {
185            // Set character encoding to UTF-8.
186            termios
187                .input_modes
188                .set(rustix::termios::InputModes::IUTF8, true);
189            let _ = rustix::termios::tcsetattr(
190                &master,
191                rustix::termios::OptionalActions::Now,
192                &termios,
193            );
194        }
195
196        let mut builder = tokio::process::Command::new(&self.program);
197
198        if let Some(ref current_dir) = self.current_dir {
199            builder.current_dir(current_dir);
200        }
201
202        builder
203            .args(&self.args)
204            .envs(&self.env)
205            .env("COLUMNS", self.terminal_size.0.to_string())
206            .env("LINES", self.terminal_size.1.to_string())
207            .env("TERM", &self.terminal_id);
208
209        builder.stdin(slave.try_clone()?);
210        builder.stderr(slave.try_clone()?);
211        builder.stdout(slave);
212
213        unsafe {
214            builder.pre_exec(move || {
215                let err = libc::setsid();
216                if err == -1 {
217                    return Err(std::io::Error::new(
218                        std::io::ErrorKind::Other,
219                        "Failed to set session id",
220                    ));
221                }
222
223                let res = libc::ioctl(slave_fd, TIOCSCTTY as _, 0);
224                if res == -1 {
225                    return Err(std::io::Error::new(
226                        std::io::ErrorKind::Other,
227                        format!("Failed to set controlling terminal: {}", res),
228                    ));
229                }
230
231                libc::close(slave_fd);
232                libc::close(master_fd);
233
234                libc::signal(libc::SIGCHLD, libc::SIG_DFL);
235                libc::signal(libc::SIGHUP, libc::SIG_DFL);
236                libc::signal(libc::SIGINT, libc::SIG_DFL);
237                libc::signal(libc::SIGQUIT, libc::SIG_DFL);
238                libc::signal(libc::SIGTERM, libc::SIG_DFL);
239                libc::signal(libc::SIGALRM, libc::SIG_DFL);
240
241                Ok(())
242            });
243        }
244
245        let mut child = builder.spawn()?;
246
247        let pid = child.id().ok_or("unable to get child pid")?;
248
249        self.out_tx.send(Output::Pid(pid)).await?;
250        self.pid = Some(pid);
251
252        let mut stdout = File::from_std(std::fs::File::from(master));
253
254        let mut stdin = stdout.try_clone().await?;
255
256        let tx_stdout = self.out_tx.clone();
257
258        let fut_out = tokio::spawn(async move {
259            let mut buf = [0u8; BUF_SIZE];
260            while let Ok(b) = stdout.read(&mut buf).await {
261                if b == 0 {
262                    break;
263                }
264                if tx_stdout
265                    .send(Output::Stdout(buf[..b].to_vec()))
266                    .await
267                    .is_err()
268                {
269                    break;
270                }
271            }
272        });
273
274        let fut_in = tokio::spawn(async move {
275            while let Ok(input) = self.in_rx.recv().await {
276                let mut data = match input {
277                    Input::Data(d) => d,
278                    Input::Resize(size) => {
279                        set_term_size(stdin.as_raw_fd(), size).ok();
280                        bmart::process::kill_pstree_with_signal(
281                            pid,
282                            bmart::process::Signal::SIGWINCH,
283                            true,
284                        );
285                        continue;
286                    }
287                    Input::Terminate => {
288                        break;
289                    }
290                };
291                // TODO: remove this input hack
292                if data == [0x0a] {
293                    data[0] = 0x0d;
294                }
295                if stdin.write_all(&data).await.is_err() {
296                    break;
297                }
298            }
299        });
300
301        let result = child.wait().await?;
302
303        fut_out.abort();
304        fut_in.abort();
305
306        let exit_code = result.code();
307
308        Ok(exit_code)
309    }
310}
311
312impl Drop for Command {
313    fn drop(&mut self) {
314        if let Some(pid) = self.pid {
315            tokio::spawn(bmart::process::kill_pstree(
316                pid,
317                Some(Duration::from_secs(1)),
318                true,
319            ));
320        }
321    }
322}
virtual_terminal/lib.rs

virtual_terminal/
lib.rs
