Skip to main content

taskers_runtime/
pty.rs

1use std::{
2    collections::BTreeMap,
3    io::{Read, Write},
4    path::PathBuf,
5};
6
7use anyhow::Result;
8use portable_pty::{CommandBuilder, MasterPty, PtySize, native_pty_system};
9
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub struct CommandSpec {
12    pub program: String,
13    pub args: Vec<String>,
14    pub cwd: Option<PathBuf>,
15    pub env: BTreeMap<String, String>,
16    pub cols: u16,
17    pub rows: u16,
18}
19
20impl CommandSpec {
21    pub fn shell() -> Self {
22        Self {
23            program: std::env::var("SHELL").unwrap_or_else(|_| "/bin/sh".into()),
24            args: Vec::new(),
25            cwd: None,
26            env: BTreeMap::new(),
27            cols: 120,
28            rows: 40,
29        }
30    }
31}
32
33pub struct PtySession {
34    child: Box<dyn portable_pty::Child + Send + Sync>,
35    master: Box<dyn MasterPty + Send>,
36    writer: Box<dyn Write + Send>,
37}
38
39pub struct PtyReader {
40    reader: Box<dyn Read + Send>,
41}
42
43pub struct SpawnedPty {
44    pub session: PtySession,
45    pub reader: PtyReader,
46}
47
48impl PtySession {
49    pub fn spawn(spec: &CommandSpec) -> Result<SpawnedPty> {
50        let pty_system = native_pty_system();
51        let pair = pty_system.openpty(PtySize {
52            rows: spec.rows,
53            cols: spec.cols,
54            pixel_width: 0,
55            pixel_height: 0,
56        })?;
57
58        let mut command = CommandBuilder::new(&spec.program);
59        for arg in &spec.args {
60            command.arg(arg);
61        }
62        if let Some(cwd) = &spec.cwd {
63            command.cwd(cwd);
64        }
65        for (key, value) in &spec.env {
66            command.env(key, value);
67        }
68
69        let child = pair.slave.spawn_command(command)?;
70        drop(pair.slave);
71
72        let writer = pair.master.take_writer()?;
73        let reader = pair.master.try_clone_reader()?;
74
75        Ok(SpawnedPty {
76            session: Self {
77                child,
78                master: pair.master,
79                writer,
80            },
81            reader: PtyReader { reader },
82        })
83    }
84
85    pub fn resize(&self, cols: u16, rows: u16) -> Result<()> {
86        self.master.resize(PtySize {
87            rows,
88            cols,
89            pixel_width: 0,
90            pixel_height: 0,
91        })?;
92        Ok(())
93    }
94
95    pub fn process_id(&self) -> Option<u32> {
96        self.child.process_id()
97    }
98
99    pub fn write_all(&mut self, data: &[u8]) -> std::io::Result<()> {
100        self.writer.write_all(data)
101    }
102}
103
104impl PtyReader {
105    pub fn read_into(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
106        self.reader.read(buffer)
107    }
108}