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

virtual_terminal/
lib.rs
