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    pub fn new(program: impl Into<String>) -> Self {
33        Self {
34            program: program.into(),
35            args: Vec::new(),
36            cwd: None,
37            env: BTreeMap::new(),
38            cols: 120,
39            rows: 40,
40        }
41    }
42}
43
44pub struct PtySession {
45    child: Box<dyn portable_pty::Child + Send + Sync>,
46    master: Box<dyn MasterPty + Send>,
47    writer: Box<dyn Write + Send>,
48}
49
50pub struct PtyReader {
51    reader: Box<dyn Read + Send>,
52}
53
54pub struct SpawnedPty {
55    pub session: PtySession,
56    pub reader: PtyReader,
57}
58
59impl PtySession {
60    pub fn spawn(spec: &CommandSpec) -> Result<SpawnedPty> {
61        let pty_system = native_pty_system();
62        let pair = pty_system.openpty(PtySize {
63            rows: spec.rows,
64            cols: spec.cols,
65            pixel_width: 0,
66            pixel_height: 0,
67        })?;
68
69        let mut command = CommandBuilder::new(&spec.program);
70        for arg in &spec.args {
71            command.arg(arg);
72        }
73        if let Some(cwd) = &spec.cwd {
74            command.cwd(cwd);
75        }
76        for (key, value) in &spec.env {
77            command.env(key, value);
78        }
79
80        let child = pair.slave.spawn_command(command)?;
81        drop(pair.slave);
82
83        let writer = pair.master.take_writer()?;
84        let reader = pair.master.try_clone_reader()?;
85
86        Ok(SpawnedPty {
87            session: Self {
88                child,
89                master: pair.master,
90                writer,
91            },
92            reader: PtyReader { reader },
93        })
94    }
95
96    pub fn resize(&self, cols: u16, rows: u16) -> Result<()> {
97        self.master.resize(PtySize {
98            rows,
99            cols,
100            pixel_width: 0,
101            pixel_height: 0,
102        })?;
103        Ok(())
104    }
105
106    pub fn process_id(&self) -> Option<u32> {
107        self.child.process_id()
108    }
109
110    pub fn write_all(&mut self, data: &[u8]) -> std::io::Result<()> {
111        self.writer.write_all(data)
112    }
113}
114
115impl PtyReader {
116    pub fn read_into(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
117        self.reader.read(buffer)
118    }
119}