1use std::collections::HashMap;
11use std::io::Write;
12use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
13use std::sync::Arc;
14use std::time::Duration;
15
16use portable_pty::{native_pty_system, Child, ChildKiller, CommandBuilder, MasterPty, PtySize};
17use tokio::sync::mpsc;
18use tokio::task::JoinHandle;
19
20use crate::{Result, RuntimeError};
21
22const PTY_OUTPUT_CHANNEL_CAPACITY: usize = 64;
23const PTY_OUTPUT_MAX_CHUNK_BYTES: usize = 4096;
24static PTY_PRODUCED_BYTES: AtomicU64 = AtomicU64::new(0);
25static PTY_ACCEPTED_BYTES: AtomicU64 = AtomicU64::new(0);
26static PTY_CONSUMED_BYTES: AtomicU64 = AtomicU64::new(0);
27static PTY_DROPPED_BYTES: AtomicU64 = AtomicU64::new(0);
28static PTY_ACTIVE_READERS: AtomicU64 = AtomicU64::new(0);
29
30#[cfg(test)]
35static TEST_PAUSE_AFTER_SEND_MS: AtomicU64 = AtomicU64::new(0);
36
37struct PtyAccounting {
49 produced: AtomicU64,
50 consumed: AtomicU64,
51 dropped: AtomicU64,
52 parties: AtomicU64,
54}
55
56impl PtyAccounting {
57 fn new() -> Arc<Self> {
58 Arc::new(Self {
59 produced: AtomicU64::new(0),
60 consumed: AtomicU64::new(0),
61 dropped: AtomicU64::new(0),
62 parties: AtomicU64::new(2),
63 })
64 }
65
66 fn add_produced(&self, n: u64) {
67 self.produced.fetch_add(n, Ordering::Relaxed);
68 PTY_PRODUCED_BYTES.fetch_add(n, Ordering::Relaxed);
69 }
70
71 fn add_consumed(&self, n: u64) {
72 self.consumed.fetch_add(n, Ordering::Relaxed);
73 PTY_CONSUMED_BYTES.fetch_add(n, Ordering::Relaxed);
74 }
75
76 fn add_dropped(&self, n: u64) {
77 self.dropped.fetch_add(n, Ordering::Relaxed);
78 PTY_DROPPED_BYTES.fetch_add(n, Ordering::Relaxed);
79 }
80
81 #[cfg(test)]
85 fn released(&self) -> bool {
86 self.parties.load(Ordering::Acquire) == 0
87 && self.produced.load(Ordering::Acquire)
88 == self.consumed.load(Ordering::Acquire) + self.dropped.load(Ordering::Acquire)
89 }
90
91 #[cfg(test)]
92 fn retained(&self) -> u64 {
93 let produced = self.produced.load(Ordering::Acquire);
94 let consumed = self.consumed.load(Ordering::Acquire);
95 let dropped = self.dropped.load(Ordering::Acquire);
96 produced.saturating_sub(consumed).saturating_sub(dropped)
97 }
98
99 fn finish_party(&self) {
103 if self.parties.fetch_sub(1, Ordering::AcqRel) == 1 {
104 let produced = self.produced.load(Ordering::Acquire);
105 let consumed = self.consumed.load(Ordering::Acquire);
106 let dropped = self.dropped.load(Ordering::Acquire);
107 let remainder = produced.saturating_sub(consumed).saturating_sub(dropped);
108 if remainder > 0 {
109 self.add_dropped(remainder);
110 }
111 }
112 }
113}
114
115#[derive(Debug, Clone, Copy)]
116pub struct PtyOutputSnapshot {
117 pub produced_bytes: u64,
118 pub accepted_bytes: u64,
121 pub dropped_bytes: u64,
122 pub retained_bytes: u64,
123 pub active_readers: u64,
124}
125
126pub fn pty_output_snapshot() -> PtyOutputSnapshot {
127 let produced = PTY_PRODUCED_BYTES.load(Ordering::Relaxed);
128 let accepted = PTY_ACCEPTED_BYTES.load(Ordering::Relaxed);
129 let consumed = PTY_CONSUMED_BYTES.load(Ordering::Relaxed);
130 let dropped = PTY_DROPPED_BYTES.load(Ordering::Relaxed);
131 PtyOutputSnapshot {
132 produced_bytes: produced,
133 accepted_bytes: accepted,
134 dropped_bytes: dropped,
135 retained_bytes: produced.saturating_sub(consumed).saturating_sub(dropped),
136 active_readers: PTY_ACTIVE_READERS.load(Ordering::SeqCst),
137 }
138}
139
140pub struct PtyHandle {
146 master: Box<dyn MasterPty + Send>,
148 writer: Box<dyn Write + Send>,
150 _reader_task: JoinHandle<()>,
152 output_rx: mpsc::Receiver<Vec<u8>>,
154 accounting: Arc<PtyAccounting>,
157 child: Box<dyn Child + Send + Sync>,
159 alive: Arc<AtomicBool>,
161 killer: Box<dyn ChildKiller + Send + Sync>,
163}
164
165type PreparedScope = (
166 Vec<(std::ffi::OsString, std::ffi::OsString)>,
167 Option<std::path::PathBuf>,
168 Option<crate::process_scope::ScopeLauncher>,
169);
170
171impl PtyHandle {
172 pub fn spawn(
180 command: &str,
181 working_dir: Option<&str>,
182 env: HashMap<String, String>,
183 rows: u16,
184 cols: u16,
185 ) -> Result<Self> {
186 let prepared = Self::prepare_scope(working_dir, &env);
187 Self::spawn_prepared(command, working_dir, rows, cols, prepared)
188 }
189
190 pub(crate) fn prepare_scope(working_dir: Option<&str>, env: &HashMap<String, String>) -> PreparedScope {
191 let mut effective_env = crate::process_scope::inherited_env();
194 crate::process_scope::set_env(&mut effective_env, "TERM", "xterm-256color");
195 for (k, v) in env { crate::process_scope::set_env(&mut effective_env, k, v); }
196 let scope_cwd = working_dir.map(std::path::PathBuf::from)
197 .filter(|p| p.is_dir())
198 .or_else(|| effective_env.iter().find(|(k, _)| k == "HOME")
199 .map(|(_, v)| std::path::PathBuf::from(v)));
200 let launcher = scope_cwd.as_deref().and_then(|cwd|
201 crate::process_scope::launcher_for_effective_env(
202 &effective_env, Some(cwd), crate::process_scope::ScopeKind::Tool));
203 (effective_env, scope_cwd, launcher)
204 }
205
206 pub(crate) fn spawn_prepared(
207 command: &str, working_dir: Option<&str>, rows: u16, cols: u16,
208 prepared: PreparedScope,
209 ) -> Result<Self> {
210 let pty_system = native_pty_system();
212 let pair = pty_system
213 .openpty(PtySize {
214 rows,
215 cols,
216 pixel_width: 0,
217 pixel_height: 0,
218 })
219 .map_err(|e| RuntimeError::Tool(format!("Failed to open PTY: {e}")))?;
220
221 let parts: Vec<&str> = command.split_whitespace().collect();
224 let program = parts
225 .first()
226 .ok_or_else(|| RuntimeError::Tool("Empty command string".to_string()))?;
227 let (effective_env, scope_cwd, launcher) = prepared;
228 let (program, args): (std::ffi::OsString, Vec<std::ffi::OsString>) = {
229 let program: std::ffi::OsString = (*program).into();
230 let args = parts
231 .iter()
232 .skip(1)
233 .map(std::ffi::OsString::from)
234 .collect::<Vec<_>>();
235 match launcher {
236 Some(launcher) => launcher.wrap("synaps tool: shell", crate::process_scope::resolve_program(program, &effective_env, scope_cwd.as_deref()).map_err(|e| RuntimeError::Tool(e.to_string()))?, args),
237 None => (program, args),
238 }
239 };
240 let mut cmd = CommandBuilder::new(program);
241 for arg in &args {
242 cmd.arg(arg);
243 }
244
245 if let Some(dir) = working_dir {
247 cmd.cwd(dir);
248 }
249
250 cmd.env_clear();
251 for (k, v) in &effective_env { cmd.env(k, v); }
252
253 let child = pair
255 .slave
256 .spawn_command(cmd)
257 .map_err(|e| RuntimeError::Tool(format!("Failed to spawn command: {e}")))?;
258
259 drop(pair.slave);
261
262 let writer = pair
264 .master
265 .take_writer()
266 .map_err(|e| RuntimeError::Tool(format!("Failed to take PTY writer: {e}")))?;
267
268 let mut reader = pair
269 .master
270 .try_clone_reader()
271 .map_err(|e| RuntimeError::Tool(format!("Failed to clone PTY reader: {e}")))?;
272
273 let (output_tx, output_rx) = mpsc::channel::<Vec<u8>>(PTY_OUTPUT_CHANNEL_CAPACITY);
275 let accounting = PtyAccounting::new();
276 let reader_accounting = Arc::clone(&accounting);
277 let alive = Arc::new(AtomicBool::new(true));
278 let reader_alive = alive.clone();
279
280 PTY_ACTIVE_READERS.fetch_add(1, Ordering::SeqCst);
281 let reader_task = tokio::task::spawn_blocking(move || {
282 struct ReaderGauge(Arc<PtyAccounting>);
286 impl Drop for ReaderGauge {
287 fn drop(&mut self) {
288 self.0.finish_party();
289 PTY_ACTIVE_READERS.fetch_sub(1, Ordering::SeqCst);
290 }
291 }
292 let _gauge = ReaderGauge(Arc::clone(&reader_accounting));
293 let mut buf = [0u8; PTY_OUTPUT_MAX_CHUNK_BYTES];
294 loop {
295 match reader.read(&mut buf) {
296 Ok(0) => {
297 break;
299 }
300 Ok(n) => {
301 reader_accounting.add_produced(n as u64);
302 match output_tx.blocking_send(buf[..n].to_vec()) {
303 Ok(()) => {
304 #[cfg(test)]
305 {
306 let ms = TEST_PAUSE_AFTER_SEND_MS.load(Ordering::Relaxed);
307 if ms > 0 {
308 std::thread::sleep(Duration::from_millis(ms));
309 }
310 }
311 PTY_ACCEPTED_BYTES.fetch_add(n as u64, Ordering::Relaxed);
312 }
313 Err(_) => {
314 reader_accounting.add_dropped(n as u64);
315 break;
316 }
317 }
318 }
319 Err(_) => {
320 break;
322 }
323 }
324 }
325 reader_alive.store(false, Ordering::SeqCst);
326 });
327
328 let killer = child.clone_killer();
330
331 Ok(PtyHandle {
332 master: pair.master,
333 writer,
334 _reader_task: reader_task,
335 output_rx,
336 accounting,
337 child,
338 alive,
339 killer,
340 })
341 }
342
343 pub fn write(&mut self, input: &[u8]) -> Result<()> {
345 self.writer
346 .write_all(input)
347 .map_err(|e| RuntimeError::Tool(format!("PTY write failed: {e}")))?;
348 self.writer
349 .flush()
350 .map_err(|e| RuntimeError::Tool(format!("PTY flush failed: {e}")))?;
351 Ok(())
352 }
353
354 pub async fn try_read_output(&mut self, timeout: Duration) -> Vec<u8> {
363 let mut collected = Vec::new();
364
365 while let Ok(chunk) = self.output_rx.try_recv() {
367 self.accounting.add_consumed(chunk.len() as u64);
368 collected.extend_from_slice(&chunk);
369 }
370
371 if collected.is_empty() {
373 match tokio::time::timeout(timeout, self.output_rx.recv()).await {
374 Ok(Some(chunk)) => {
375 self.accounting.add_consumed(chunk.len() as u64);
376 collected.extend_from_slice(&chunk);
377 }
378 Ok(None) | Err(_) => {
379 return collected;
381 }
382 }
383
384 while let Ok(chunk) = self.output_rx.try_recv() {
386 self.accounting.add_consumed(chunk.len() as u64);
387 collected.extend_from_slice(&chunk);
388 }
389 }
390
391 collected
392 }
393
394 pub fn resize(&self, rows: u16, cols: u16) -> Result<()> {
396 self.master
397 .resize(PtySize {
398 rows,
399 cols,
400 pixel_width: 0,
401 pixel_height: 0,
402 })
403 .map_err(|e| RuntimeError::Tool(format!("PTY resize failed: {e}")))
404 }
405
406 pub fn is_alive(&mut self) -> bool {
410 if !self.alive.load(Ordering::SeqCst) {
411 return false;
412 }
413 match self.child.try_wait() {
414 Ok(Some(_status)) => {
415 self.alive.store(false, Ordering::SeqCst);
417 false
418 }
419 Ok(None) => true,
420 Err(_) => {
421 self.alive.store(false, Ordering::SeqCst);
423 false
424 }
425 }
426 }
427}
428
429impl Drop for PtyHandle {
430 fn drop(&mut self) {
431 self.accounting.finish_party();
438 if self.alive.load(Ordering::SeqCst) {
439 let _ = self.killer.kill();
440 for _ in 0..5 {
445 match self.child.try_wait() {
446 Ok(Some(_)) | Err(_) => break, Ok(None) => std::thread::sleep(std::time::Duration::from_millis(10)),
448 }
449 }
450 }
451 }
452}
453
454#[cfg(test)]
455mod tests {
456 use super::*;
457 use serial_test::serial;
462 use std::collections::HashMap;
463
464 #[cfg(target_os = "linux")]
467 #[test]
468 #[serial_test::serial(synaps_base_dir)]
469 fn a_pty_shell_is_scoped_in_the_tool_slice() {
470 use crate::process_scope::test_support::{fake_env, fake_systemd_run, recorded};
471 let base = crate::test_env::BaseDirGuard::new();
472 let _scopes = crate::test_env::EnvVarGuard::remove("SYNAPS_PROCESS_SCOPES");
473 let dir = tempfile::tempdir().unwrap();
474 let fake = fake_systemd_run(dir.path());
475 let env: HashMap<String, String> = fake_env(&fake).into_iter().collect();
476 let cwd = dir.path().to_str().unwrap();
477 let slices = || {
478 let (_, _, launcher) = PtyHandle::prepare_scope(Some(cwd), &env);
479 let (_, args) = launcher.expect("the fake probe succeeds").wrap("synaps tool: shell", "/bin/sh", []);
480 args.iter()
481 .map(|a| a.to_string_lossy().into_owned())
482 .filter(|a| a.starts_with("--slice"))
483 .collect::<Vec<_>>()
484 };
485 assert!(slices().is_empty());
486 std::fs::write(
487 base.path().join("config"),
488 "tool_scope_slice = t_tools.slice\ndaemon_scope_slice = d_daemon.slice\n",
489 )
490 .unwrap();
491 assert_eq!(slices(), ["--slice=t_tools.slice"]);
492 assert!(recorded(dir.path(), "probe.args").contains(&"--slice=t_tools.slice".to_string()));
493 }
494
495 async fn read_until(handle: &mut PtyHandle, needle: &str) -> String {
499 let deadline = tokio::time::Instant::now() + Duration::from_secs(3);
500 let mut out = Vec::new();
501 while !String::from_utf8_lossy(&out).contains(needle)
502 && tokio::time::Instant::now() < deadline
503 {
504 out.extend(handle.try_read_output(Duration::from_millis(200)).await);
505 }
506 String::from_utf8_lossy(&out).into_owned()
507 }
508
509 async fn wait_released(accounting: &PtyAccounting, what: &str) {
515 let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
516 while !accounting.released() {
517 assert!(
518 tokio::time::Instant::now() < deadline,
519 "{what}: reader still running or bytes leaked \
520 (parties {}, retained {})",
521 accounting.parties.load(Ordering::Acquire),
522 accounting.retained()
523 );
524 tokio::time::sleep(Duration::from_millis(10)).await;
525 }
526 }
527
528 #[cfg(target_os = "linux")]
534 #[tokio::test]
535 #[serial(synaps_base_dir)]
536 async fn pty_shell_gets_its_own_scope_when_available() {
537 if std::env::var("SYNAPS_TEST_REAL_SCOPES").as_deref() != Ok("1") { return; }
538 let own = std::fs::read_to_string("/proc/self/cgroup").unwrap_or_default();
539 let mut handle = PtyHandle::spawn("bash --norc --noprofile", None, HashMap::new(), 24, 80)
542 .expect("spawn");
543 handle
544 .write(b"cat /proc/$$/cgroup; test -t 0 && echo tty=yes; exit\n")
545 .expect("write");
546 let mut out = Vec::new();
547 let deadline = std::time::Instant::now() + Duration::from_secs(5);
548 while std::time::Instant::now() < deadline
549 && !String::from_utf8_lossy(&out).contains("tty=yes\r")
550 {
551 out.extend(handle.try_read_output(Duration::from_millis(200)).await);
552 }
553 let accounting = Arc::clone(&handle.accounting);
554 drop(handle);
555 wait_released(&accounting, "scoped shell").await;
556 let text = String::from_utf8_lossy(&out).replace('\r', "");
557 assert!(
558 text.contains("tty=yes"),
559 "the shell must still have a terminal: {text}"
560 );
561 let child = text
564 .lines()
565 .filter(|l| !l.contains("cat /proc/")) .find_map(|l| l.find("0::/").map(|i| l[i..].trim().to_string()))
567 .unwrap_or_default();
568 assert!(child.starts_with("0::/"), "no cgroup line in: {text:?}");
569 if crate::process_scope::launcher_for_env(None).is_some() {
570 assert_ne!(child, own.trim(), "{text}");
571 assert!(child.ends_with(".scope"), "{child}");
572 } else {
573 assert_eq!(child, own.trim(), "unwrapped fallback");
574 }
575 }
576
577 #[tokio::test]
578 #[serial(synaps_base_dir)]
579 async fn pty_slow_consumer_retention_is_bounded_and_drop_releases_reader() {
580 let handle = PtyHandle::spawn(
581 "bash -c 'yes x | head -c 10485760'",
582 None,
583 HashMap::new(),
584 24,
585 80,
586 )
587 .expect("spawn producer");
588 let accounting = Arc::clone(&handle.accounting);
589 let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
593 while accounting.retained() == 0 {
594 assert!(
595 tokio::time::Instant::now() < deadline,
596 "stalled consumer must retain queued PTY bytes"
597 );
598 tokio::time::sleep(Duration::from_millis(10)).await;
599 }
600 tokio::time::sleep(Duration::from_millis(100)).await;
602 let retained = accounting.retained();
603 assert!(
604 retained <= ((PTY_OUTPUT_CHANNEL_CAPACITY + 1) * PTY_OUTPUT_MAX_CHUNK_BYTES) as u64,
605 "retained {retained}"
606 );
607 drop(handle);
608 wait_released(&accounting, "drop must release the reader and its queued bytes").await;
609 }
610
611 #[tokio::test]
617 #[serial(synaps_base_dir)]
618 async fn pty_drop_races_reader_and_always_conserves_accounting() {
619 let mut released = Vec::new();
620 for _ in 0..5 {
621 let mut handle = PtyHandle::spawn(
622 "bash -c 'yes x | head -c 33554432'",
623 None,
624 HashMap::new(),
625 24,
626 80,
627 )
628 .expect("spawn producer");
629 let _ = handle.try_read_output(Duration::from_millis(50)).await;
634 released.push(Arc::clone(&handle.accounting));
635 drop(handle);
636 }
637 for accounting in &released {
640 wait_released(accounting, "drop/reader race must conserve accounting").await;
641 }
642 }
643
644 #[tokio::test]
649 #[serial(synaps_base_dir)]
650 async fn pty_preempted_accounting_across_drop_is_released_exactly_once() {
651 TEST_PAUSE_AFTER_SEND_MS.store(150, Ordering::Relaxed);
652 let handle = PtyHandle::spawn("bash -c 'echo leak-window'", None, HashMap::new(), 24, 80)
653 .expect("spawn producer");
654 let accounting = Arc::clone(&handle.accounting);
655 tokio::time::sleep(Duration::from_millis(40)).await;
657 drop(handle); TEST_PAUSE_AFTER_SEND_MS.store(0, Ordering::Relaxed);
659
660 wait_released(&accounting, "a drop inside the accounting gap must not leak").await;
661 }
662
663 #[tokio::test]
668 #[serial(synaps_base_dir)]
669 async fn pty_concurrent_handles_release_exactly_once_per_handle() {
670 let mut tasks = Vec::new();
671 for i in 0..8 {
672 tasks.push(tokio::spawn(async move {
673 let mut handle = PtyHandle::spawn(
674 "bash -c 'yes x | head -c 262144'",
675 None,
676 HashMap::new(),
677 24,
678 80,
679 )
680 .expect("spawn producer");
681 match i % 3 {
682 0 => {} 1 => {
684 tokio::time::sleep(Duration::from_millis(5)).await;
685 }
686 _ => {
687 let _ = handle.try_read_output(Duration::from_millis(10)).await;
689 }
690 }
691 let accounting = Arc::clone(&handle.accounting);
692 drop(handle);
693 accounting
694 }));
695 }
696 for t in tasks {
697 let accounting = t.await.unwrap();
698 wait_released(&accounting, "concurrent handles must each release exactly once").await;
699 }
700 }
701
702 #[tokio::test]
703 #[serial(synaps_base_dir)]
704 async fn test_spawn_echo_hello() {
705 let mut handle = PtyHandle::spawn("echo hello", None, HashMap::new(), 24, 80)
706 .expect("failed to spawn echo");
707
708 let text = read_until(&mut handle, "hello").await;
709 assert!(
710 text.contains("hello"),
711 "expected 'hello' in output, got: {text:?}"
712 );
713 }
714
715 #[tokio::test]
716 #[serial(synaps_base_dir)]
717 async fn test_cat_echo_back() {
718 let mut handle =
719 PtyHandle::spawn("cat", None, HashMap::new(), 24, 80).expect("failed to spawn cat");
720
721 handle.write(b"test\n").expect("write failed");
723
724 let text = read_until(&mut handle, "test").await;
725 assert!(
726 text.contains("test"),
727 "expected 'test' in output, got: {text:?}"
728 );
729 }
730
731 #[tokio::test]
732 #[serial(synaps_base_dir)]
733 async fn test_exit_code_detection() {
734 let mut handle = PtyHandle::spawn("bash -c exit 42", None, HashMap::new(), 24, 80)
735 .expect("failed to spawn bash exit");
736
737 let _ = handle.try_read_output(Duration::from_secs(3)).await;
739
740 tokio::time::sleep(Duration::from_millis(200)).await;
742
743 assert!(!handle.is_alive(), "expected process to have exited");
744 }
745}