zwire_host/
pty.rs

1//! Multiplexed PTY sessions.
2//!
3//! `pty_spawn` opens a login shell (or any program) in a pseudo-terminal and
4//! streams its output as `{"ev":"output","b64":…}` frames. Sessions are keyed by
5//! the request's `id`, so one connection can drive many terminals at once —
6//! `pty_write` / `pty_resize` / `pty_kill` route by the same `id`. An absent
7//! `id` uses the default (empty) key and omits the `pty` field from frames, so
8//! the original single-terminal zwire protocol is unchanged.
9use crate::proto::{b64_encode, send_msg, Out};
10use crate::store::expand;
11use portable_pty::{native_pty_system, Child, CommandBuilder, MasterPty, PtySize};
12use serde_json::{json, Value};
13use std::io::{Read, Write};
14use std::thread::JoinHandle;
15
16/// One live pseudo-terminal: the master write side, a resize handle, the child
17/// process, and the reader thread pumping output frames.
18pub struct PtySession {
19    writer: Box<dyn Write + Send>,
20    master: Box<dyn MasterPty + Send>,
21    child: Box<dyn Child + Send + Sync>,
22    reader: Option<JoinHandle<()>>,
23}
24
25fn default_shell() -> String {
26    if cfg!(windows) {
27        std::env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".into())
28    } else {
29        std::env::var("SHELL").unwrap_or_else(|_| "/bin/sh".into())
30    }
31}
32
33fn output_frame(key: &str, bytes: &[u8]) -> Value {
34    let b64 = b64_encode(bytes);
35    if key.is_empty() {
36        json!({ "ev": "output", "b64": b64 })
37    } else {
38        json!({ "ev": "output", "pty": key, "b64": b64 })
39    }
40}
41
42fn exit_frame(key: &str) -> Value {
43    if key.is_empty() {
44        json!({ "ev": "exit" })
45    } else {
46        json!({ "ev": "exit", "pty": key })
47    }
48}
49
50impl PtySession {
51    /// Spawn a terminal from a `pty_spawn` request. `key` is the session id used
52    /// to tag frames (empty = legacy default). Returns `None` if the OS refuses
53    /// to open the PTY or spawn the program.
54    ///
55    /// Recognised request fields: `rows`, `cols`, `shell` (program to run,
56    /// defaults to `$SHELL`), `args` (defaults to `["-l"]`), `cwd`, `env`.
57    pub fn spawn(out: &Out, req: &Value, key: String) -> Option<PtySession> {
58        let rows = req["rows"].as_u64().unwrap_or(24) as u16;
59        let cols = req["cols"].as_u64().unwrap_or(80) as u16;
60        let pair = native_pty_system()
61            .openpty(PtySize {
62                rows,
63                cols,
64                pixel_width: 0,
65                pixel_height: 0,
66            })
67            .ok()?;
68
69        let program = req["shell"]
70            .as_str()
71            .map_or_else(default_shell, String::from);
72        let mut cmd = CommandBuilder::new(&program);
73        match req["args"].as_array() {
74            Some(args) => {
75                for a in args {
76                    if let Some(s) = a.as_str() {
77                        cmd.arg(s);
78                    }
79                }
80            }
81            None => cmd.arg("-l"),
82        }
83        if let Some(cwd) = req["cwd"].as_str() {
84            cmd.cwd(expand(cwd));
85        }
86        cmd.env("TERM", "xterm-256color");
87        if let Some(env) = req["env"].as_object() {
88            for (k, v) in env {
89                if let Some(s) = v.as_str() {
90                    cmd.env(k, s);
91                }
92            }
93        }
94
95        let child = pair.slave.spawn_command(cmd).ok()?;
96        drop(pair.slave);
97        let mut reader = pair.master.try_clone_reader().ok()?;
98        let writer = pair.master.take_writer().ok()?;
99
100        let out2 = out.clone();
101        let handle = std::thread::spawn(move || {
102            let mut buf = [0u8; 65536];
103            loop {
104                match reader.read(&mut buf) {
105                    Ok(0) | Err(_) => break,
106                    Ok(n) => {
107                        if send_msg(&out2, &output_frame(&key, &buf[..n])).is_err() {
108                            break;
109                        }
110                    }
111                }
112            }
113            let _ = send_msg(&out2, &exit_frame(&key));
114        });
115
116        Some(PtySession {
117            writer,
118            master: pair.master,
119            child,
120            reader: Some(handle),
121        })
122    }
123
124    /// Feed bytes to the terminal (`pty_write`). Accepts `data` (UTF-8 string)
125    /// or `b64` (binary) on the request.
126    pub fn write(&mut self, req: &Value) {
127        let bytes = if let Some(s) = req["data"].as_str() {
128            s.as_bytes().to_vec()
129        } else if let Some(b) = req["b64"].as_str() {
130            crate::proto::b64_decode(b).unwrap_or_default()
131        } else {
132            return;
133        };
134        let _ = self.writer.write_all(&bytes);
135        let _ = self.writer.flush();
136    }
137
138    /// Resize the terminal (`pty_resize`).
139    pub fn resize(&self, req: &Value) {
140        let rows = req["rows"].as_u64().unwrap_or(24) as u16;
141        let cols = req["cols"].as_u64().unwrap_or(80) as u16;
142        let _ = self.master.resize(PtySize {
143            rows,
144            cols,
145            pixel_width: 0,
146            pixel_height: 0,
147        });
148    }
149}
150
151impl Drop for PtySession {
152    fn drop(&mut self) {
153        let _ = self.child.kill();
154        let _ = self.child.wait();
155        if let Some(h) = self.reader.take() {
156            let _ = h.join();
157        }
158    }
159}
zwire_host/pty.rs

zwire_host/
pty.rs
