1use std::collections::HashMap;
10use std::fs;
11use std::io::{Read, Write};
12use std::path::PathBuf;
13use std::sync::atomic::{AtomicU32, Ordering};
14use std::sync::{Arc, Mutex};
15use std::thread;
16
17use base64::Engine;
18use portable_pty::{CommandBuilder, NativePtySystem, PtySize, PtySystem};
19use serde_json::Value;
20
21use crate::events::EventEmitter;
22use crate::sandbox::{is_dir_allowed, safe_lock, validate_path, SharedList};
23
24const MAX_PTY_WRITE: usize = 1024 * 1024;
25
26pub struct PtyConfig {
27 pub allowed_tools: Vec<String>,
28}
29
30pub struct PtySession {
31 writer: Box<dyn Write + Send>,
32 master: Box<dyn portable_pty::MasterPty + Send>,
33 child: Box<dyn portable_pty::Child + Send + Sync>,
34}
35
36pub struct PtySessions {
37 pub inner: Arc<Mutex<HashMap<u32, PtySession>>>,
38 pub next_id: AtomicU32,
39 pub config: PtyConfig,
40}
41
42impl PtySessions {
43 pub fn new(config: PtyConfig) -> Self {
44 Self {
45 inner: Arc::new(Mutex::new(HashMap::new())),
46 next_id: AtomicU32::new(1),
47 config,
48 }
49 }
50}
51
52fn validate_pty_tool<'a>(sessions: &'a PtySessions, tool: &str) -> Result<&'a str, String> {
53 sessions
54 .config
55 .allowed_tools
56 .iter()
57 .find(|t| t.as_str() == tool)
58 .map(|s| s.as_str())
59 .ok_or_else(|| format!("Tool not allowed: {}", tool))
60}
61
62fn trusted_dirs() -> Vec<PathBuf> {
63 let mut dirs: Vec<PathBuf> = vec![
64 "/opt/homebrew/bin".into(),
65 "/usr/local/bin".into(),
66 "/usr/bin".into(),
67 "/bin".into(),
68 ];
69 if let Ok(home) = std::env::var("HOME") {
70 let home = PathBuf::from(home);
71 dirs.push(home.join(".cargo/bin"));
72 dirs.push(home.join(".local/bin"));
73 dirs.push(home.join(".volta/bin"));
74 dirs.push(home.join(".npm-global/bin"));
75 dirs.push(home.join(".bun/bin"));
76 }
77 dirs
78}
79
80fn resolve_pty_tool(tool: &str) -> Result<String, String> {
81 for d in &trusted_dirs() {
82 let candidate = d.join(tool);
83 if candidate.is_file() {
84 return Ok(candidate.to_string_lossy().to_string());
85 }
86 }
87 if let Ok(home) = std::env::var("HOME") {
88 let nvm = PathBuf::from(home).join(".nvm/versions/node");
89 if let Ok(entries) = fs::read_dir(&nvm) {
90 for e in entries.flatten() {
91 let p = e.path().join("bin").join(tool);
92 if p.is_file() {
93 return Ok(p.to_string_lossy().to_string());
94 }
95 }
96 }
97 }
98 Err(format!("Tool `{}` not found in trusted install dirs", tool))
99}
100
101pub fn spawn(
102 sessions: &PtySessions,
103 allowed_dirs: &SharedList,
104 emitter: &EventEmitter,
105 args: &Value,
106) -> Result<Value, String> {
107 let tool = args
108 .get("tool")
109 .and_then(|v| v.as_str())
110 .ok_or("missing tool")?;
111 let cwd = args
112 .get("cwd")
113 .and_then(|v| v.as_str())
114 .ok_or("missing cwd")?;
115 let cols = args.get("cols").and_then(|v| v.as_u64()).unwrap_or(80) as u16;
116 let rows = args.get("rows").and_then(|v| v.as_u64()).unwrap_or(24) as u16;
117
118 let tool = validate_pty_tool(sessions, tool)?.to_string();
119 validate_path(cwd)?;
120 let canonical_cwd = is_dir_allowed(cwd, allowed_dirs)?;
121
122 let pty_system = NativePtySystem::default();
123 let pair = pty_system
124 .openpty(PtySize {
125 rows: rows.max(1),
126 cols: cols.max(1),
127 pixel_width: 0,
128 pixel_height: 0,
129 })
130 .map_err(|e| format!("Cannot open pty: {}", e))?;
131
132 let tool_abs = resolve_pty_tool(&tool)?;
133 let mut cmd = CommandBuilder::new(&tool_abs);
134 cmd.cwd(&canonical_cwd);
135 cmd.env("TERM", "xterm-256color");
136 let mut safe_path = String::from("/opt/homebrew/bin:/usr/local/bin:/usr/bin:/bin");
137 if let Ok(home) = std::env::var("HOME") {
138 safe_path.push_str(&format!(
139 ":{h}/.cargo/bin:{h}/.local/bin:{h}/.volta/bin:{h}/.npm-global/bin:{h}/.bun/bin",
140 h = home
141 ));
142 cmd.env("HOME", &home);
143 }
144 cmd.env("PATH", safe_path);
145
146 let child = pair
147 .slave
148 .spawn_command(cmd)
149 .map_err(|e| format!("Cannot spawn {}: {}", tool, e))?;
150 let mut reader = pair
151 .master
152 .try_clone_reader()
153 .map_err(|e| format!("Cannot clone pty reader: {}", e))?;
154 let writer = pair
155 .master
156 .take_writer()
157 .map_err(|e| format!("Cannot take pty writer: {}", e))?;
158
159 let id = sessions.next_id.fetch_add(1, Ordering::SeqCst);
160 {
161 let mut map = safe_lock(&sessions.inner);
162 map.insert(
163 id,
164 PtySession {
165 writer,
166 master: pair.master,
167 child,
168 },
169 );
170 }
171
172 let emitter_reader = emitter.clone();
173 let sessions_for_reader = Arc::clone(&sessions.inner);
174 let _ = thread::Builder::new()
175 .name(format!("pty-reader-{}", id))
176 .spawn(move || {
177 let mut buf = [0u8; 8192];
178 loop {
179 match reader.read(&mut buf) {
180 Ok(0) => break,
181 Ok(n) => {
182 let encoded = base64::engine::general_purpose::STANDARD.encode(&buf[..n]);
183 emitter_reader
184 .emit("pty:data", serde_json::json!({ "id": id, "data": encoded }));
185 }
186 Err(_) => break,
187 }
188 }
189 emitter_reader.emit("pty:exit", serde_json::json!({ "id": id }));
190 let mut map = safe_lock(&sessions_for_reader);
191 map.remove(&id);
192 });
193
194 Ok(Value::from(id))
195}
196
197pub fn write(sessions: &PtySessions, args: &Value) -> Result<Value, String> {
198 let id = args
199 .get("id")
200 .and_then(|v| v.as_u64())
201 .ok_or("missing id")? as u32;
202 let data = args
203 .get("data")
204 .and_then(|v| v.as_str())
205 .ok_or("missing data")?;
206 if data.len() > MAX_PTY_WRITE {
207 return Err("Input too large".to_string());
208 }
209 let mut map = safe_lock(&sessions.inner);
210 let session = map.get_mut(&id).ok_or("Session not found")?;
211 session
212 .writer
213 .write_all(data.as_bytes())
214 .map_err(|e| format!("Write failed: {}", e))?;
215 Ok(Value::Null)
216}
217
218pub fn resize(sessions: &PtySessions, args: &Value) -> Result<Value, String> {
219 let id = args
220 .get("id")
221 .and_then(|v| v.as_u64())
222 .ok_or("missing id")? as u32;
223 let cols = args.get("cols").and_then(|v| v.as_u64()).unwrap_or(80) as u16;
224 let rows = args.get("rows").and_then(|v| v.as_u64()).unwrap_or(24) as u16;
225 let map = safe_lock(&sessions.inner);
226 let session = map.get(&id).ok_or("Session not found")?;
227 session
228 .master
229 .resize(PtySize {
230 rows: rows.max(1),
231 cols: cols.max(1),
232 pixel_width: 0,
233 pixel_height: 0,
234 })
235 .map_err(|e| format!("Resize failed: {}", e))?;
236 Ok(Value::Null)
237}
238
239pub fn kill(sessions: &PtySessions, args: &Value) -> Result<Value, String> {
240 let id = args
241 .get("id")
242 .and_then(|v| v.as_u64())
243 .ok_or("missing id")? as u32;
244 let mut map = safe_lock(&sessions.inner);
245 if let Some(mut session) = map.remove(&id) {
246 let _ = session.child.kill();
247 }
248 Ok(Value::Null)
249}
250
251#[cfg(test)]
252mod tests {
253 use super::*;
254
255 fn test_sessions() -> PtySessions {
256 PtySessions::new(PtyConfig {
257 allowed_tools: vec!["claude".into(), "codex".into()],
258 })
259 }
260
261 #[test]
262 fn allows_configured_tools() {
263 let s = test_sessions();
264 assert_eq!(validate_pty_tool(&s, "claude").unwrap(), "claude");
265 assert_eq!(validate_pty_tool(&s, "codex").unwrap(), "codex");
266 }
267
268 #[test]
269 fn rejects_unlisted_binaries() {
270 let s = test_sessions();
271 assert!(validate_pty_tool(&s, "sh").is_err());
272 assert!(validate_pty_tool(&s, "bash").is_err());
273 assert!(validate_pty_tool(&s, "claude; rm -rf /").is_err());
274 assert!(validate_pty_tool(&s, "/usr/bin/claude").is_err());
275 assert!(validate_pty_tool(&s, "").is_err());
276 }
277
278 #[test]
279 fn rejects_path_like_tool_names() {
280 let s = test_sessions();
281 assert!(validate_pty_tool(&s, "./claude").is_err());
283 assert!(validate_pty_tool(&s, "../claude").is_err());
284 }
285
286 #[test]
287 fn trusted_dirs_include_standard_unix_bins() {
288 let dirs = trusted_dirs();
289 let as_str: Vec<String> = dirs.iter().map(|p| p.to_string_lossy().into()).collect();
290 assert!(as_str.iter().any(|d| d == "/usr/bin"));
291 assert!(as_str.iter().any(|d| d == "/bin"));
292 assert!(as_str.iter().any(|d| d == "/opt/homebrew/bin"));
293 assert!(as_str.iter().any(|d| d == "/usr/local/bin"));
294 }
295
296 #[test]
297 fn trusted_dirs_include_home_managers_when_home_set() {
298 if std::env::var("HOME").is_err() {
299 return;
300 }
301 let dirs = trusted_dirs();
302 let as_str: Vec<String> = dirs.iter().map(|p| p.to_string_lossy().into()).collect();
303 assert!(as_str.iter().any(|d| d.ends_with("/.cargo/bin")));
304 assert!(as_str.iter().any(|d| d.ends_with("/.local/bin")));
305 assert!(as_str.iter().any(|d| d.ends_with("/.bun/bin")));
306 }
307
308 #[test]
309 fn resolve_pty_tool_errors_for_nonexistent() {
310 let err = resolve_pty_tool("definitely-not-installed-anywhere-xyz").unwrap_err();
311 assert!(err.contains("not found"), "got: {err}");
312 }
313
314 #[test]
315 fn write_rejects_oversize_payload() {
316 let s = test_sessions();
317 let big = "a".repeat(MAX_PTY_WRITE + 1);
318 let err = write(&s, &serde_json::json!({ "id": 1u32, "data": big })).unwrap_err();
319 assert!(err.contains("too large"), "got: {err}");
320 }
321
322 #[test]
323 fn spawn_rejects_unknown_tool() {
324 let s = test_sessions();
325 let allowed_dirs = crate::sandbox::new_list();
328 let (tx, _rx) = std::sync::mpsc::channel::<()>();
329 drop(tx);
330 assert!(validate_pty_tool(&s, "rogue").is_err());
333 let v = serde_json::json!({ "tool": "rogue" });
335 let _ = (s, allowed_dirs, v); }
338
339 #[test]
340 fn resize_and_kill_on_missing_session() {
341 let s = test_sessions();
342 assert!(resize(
343 &s,
344 &serde_json::json!({ "id": 999u32, "cols": 80, "rows": 24 })
345 )
346 .is_err());
347 assert!(kill(&s, &serde_json::json!({ "id": 999u32 })).is_ok());
349 }
350}