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, None, &env);
187 Self::spawn_prepared(command, working_dir, rows, cols, prepared)
188 }
189
190 pub(crate) fn prepare_scope(
192 working_dir: Option<&str>, base: Option<&[(String, String)]>, env: &HashMap<String, String>,
193 ) -> PreparedScope {
194 let mut effective_env = match base {
197 Some(base) => base.iter().map(|(k, v)| (k.into(), v.into())).collect(),
198 None => crate::process_scope::inherited_env(),
199 };
200 crate::process_scope::set_env(&mut effective_env, "TERM", "xterm-256color");
201 for (k, v) in env { crate::process_scope::set_env(&mut effective_env, k, v); }
202 let scope_cwd = working_dir.map(std::path::PathBuf::from)
203 .filter(|p| p.is_dir())
204 .or_else(|| effective_env.iter().find(|(k, _)| k == "HOME")
205 .map(|(_, v)| std::path::PathBuf::from(v)));
206 let launcher = scope_cwd.as_deref().and_then(|cwd|
207 crate::process_scope::launcher_for_effective_env(
208 &effective_env, Some(cwd), crate::process_scope::ScopeKind::Tool));
209 (effective_env, scope_cwd, launcher)
210 }
211
212 pub(crate) fn spawn_prepared(
213 command: &str, working_dir: Option<&str>, rows: u16, cols: u16,
214 prepared: PreparedScope,
215 ) -> Result<Self> {
216 let pty_system = native_pty_system();
218 let pair = pty_system
219 .openpty(PtySize {
220 rows,
221 cols,
222 pixel_width: 0,
223 pixel_height: 0,
224 })
225 .map_err(|e| RuntimeError::Tool(format!("Failed to open PTY: {e}")))?;
226
227 let parts: Vec<&str> = command.split_whitespace().collect();
230 let program = parts
231 .first()
232 .ok_or_else(|| RuntimeError::Tool("Empty command string".to_string()))?;
233 let (effective_env, scope_cwd, launcher) = prepared;
234 let (program, args): (std::ffi::OsString, Vec<std::ffi::OsString>) = {
235 let program: std::ffi::OsString = (*program).into();
236 let args = parts
237 .iter()
238 .skip(1)
239 .map(std::ffi::OsString::from)
240 .collect::<Vec<_>>();
241 match launcher {
242 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),
243 None => (program, args),
244 }
245 };
246 let mut cmd = CommandBuilder::new(program);
247 for arg in &args {
248 cmd.arg(arg);
249 }
250
251 if let Some(dir) = working_dir {
253 cmd.cwd(dir);
254 }
255
256 cmd.env_clear();
257 for (k, v) in &effective_env { cmd.env(k, v); }
258
259 let child = pair
261 .slave
262 .spawn_command(cmd)
263 .map_err(|e| RuntimeError::Tool(format!("Failed to spawn command: {e}")))?;
264
265 drop(pair.slave);
267
268 let writer = pair
270 .master
271 .take_writer()
272 .map_err(|e| RuntimeError::Tool(format!("Failed to take PTY writer: {e}")))?;
273
274 let mut reader = pair
275 .master
276 .try_clone_reader()
277 .map_err(|e| RuntimeError::Tool(format!("Failed to clone PTY reader: {e}")))?;
278
279 let (output_tx, output_rx) = mpsc::channel::<Vec<u8>>(PTY_OUTPUT_CHANNEL_CAPACITY);
281 let accounting = PtyAccounting::new();
282 let reader_accounting = Arc::clone(&accounting);
283 let alive = Arc::new(AtomicBool::new(true));
284 let reader_alive = alive.clone();
285
286 PTY_ACTIVE_READERS.fetch_add(1, Ordering::SeqCst);
287 let reader_task = tokio::task::spawn_blocking(move || {
288 struct ReaderGauge(Arc<PtyAccounting>);
292 impl Drop for ReaderGauge {
293 fn drop(&mut self) {
294 self.0.finish_party();
295 PTY_ACTIVE_READERS.fetch_sub(1, Ordering::SeqCst);
296 }
297 }
298 let _gauge = ReaderGauge(Arc::clone(&reader_accounting));
299 let mut buf = [0u8; PTY_OUTPUT_MAX_CHUNK_BYTES];
300 loop {
301 match reader.read(&mut buf) {
302 Ok(0) => {
303 break;
305 }
306 Ok(n) => {
307 reader_accounting.add_produced(n as u64);
308 match output_tx.blocking_send(buf[..n].to_vec()) {
309 Ok(()) => {
310 #[cfg(test)]
311 {
312 let ms = TEST_PAUSE_AFTER_SEND_MS.load(Ordering::Relaxed);
313 if ms > 0 {
314 std::thread::sleep(Duration::from_millis(ms));
315 }
316 }
317 PTY_ACCEPTED_BYTES.fetch_add(n as u64, Ordering::Relaxed);
318 }
319 Err(_) => {
320 reader_accounting.add_dropped(n as u64);
321 break;
322 }
323 }
324 }
325 Err(_) => {
326 break;
328 }
329 }
330 }
331 reader_alive.store(false, Ordering::SeqCst);
332 });
333
334 let killer = child.clone_killer();
336
337 Ok(PtyHandle {
338 master: pair.master,
339 writer,
340 _reader_task: reader_task,
341 output_rx,
342 accounting,
343 child,
344 alive,
345 killer,
346 })
347 }
348
349 pub fn write(&mut self, input: &[u8]) -> Result<()> {
351 self.writer
352 .write_all(input)
353 .map_err(|e| RuntimeError::Tool(format!("PTY write failed: {e}")))?;
354 self.writer
355 .flush()
356 .map_err(|e| RuntimeError::Tool(format!("PTY flush failed: {e}")))?;
357 Ok(())
358 }
359
360 pub async fn try_read_output(&mut self, timeout: Duration) -> Vec<u8> {
369 let mut collected = Vec::new();
370
371 while let Ok(chunk) = self.output_rx.try_recv() {
373 self.accounting.add_consumed(chunk.len() as u64);
374 collected.extend_from_slice(&chunk);
375 }
376
377 if collected.is_empty() {
379 match tokio::time::timeout(timeout, self.output_rx.recv()).await {
380 Ok(Some(chunk)) => {
381 self.accounting.add_consumed(chunk.len() as u64);
382 collected.extend_from_slice(&chunk);
383 }
384 Ok(None) | Err(_) => {
385 return collected;
387 }
388 }
389
390 while let Ok(chunk) = self.output_rx.try_recv() {
392 self.accounting.add_consumed(chunk.len() as u64);
393 collected.extend_from_slice(&chunk);
394 }
395 }
396
397 collected
398 }
399
400 pub fn resize(&self, rows: u16, cols: u16) -> Result<()> {
402 self.master
403 .resize(PtySize {
404 rows,
405 cols,
406 pixel_width: 0,
407 pixel_height: 0,
408 })
409 .map_err(|e| RuntimeError::Tool(format!("PTY resize failed: {e}")))
410 }
411
412 pub fn is_alive(&mut self) -> bool {
416 if !self.alive.load(Ordering::SeqCst) {
417 return false;
418 }
419 match self.child.try_wait() {
420 Ok(Some(_status)) => {
421 self.alive.store(false, Ordering::SeqCst);
423 false
424 }
425 Ok(None) => true,
426 Err(_) => {
427 self.alive.store(false, Ordering::SeqCst);
429 false
430 }
431 }
432 }
433}
434
435impl Drop for PtyHandle {
436 fn drop(&mut self) {
437 self.accounting.finish_party();
444 if self.alive.load(Ordering::SeqCst) {
445 let _ = self.killer.kill();
446 for _ in 0..5 {
451 match self.child.try_wait() {
452 Ok(Some(_)) | Err(_) => break, Ok(None) => std::thread::sleep(std::time::Duration::from_millis(10)),
454 }
455 }
456 }
457 }
458}
459
460#[cfg(test)]
461mod tests {
462 use super::*;
463 use serial_test::serial;
468 use std::collections::HashMap;
469
470 #[test]
473 fn a_pty_shell_starts_from_the_session_env() {
474 let base = [("SDO3_SESSION_ONLY".to_string(), "from-session".to_string()), ("PATH".into(), "/bin".into())];
475 let extra = HashMap::from([("EXTRA".to_string(), "x".to_string())]);
476 let (env, _, _) = PtyHandle::prepare_scope(None, Some(&base), &extra);
477 let get = |k: &str| env.iter().find(|(n, _)| n == k).map(|(_, v)| v.to_string_lossy().into_owned());
478 assert_eq!(get("SDO3_SESSION_ONLY").as_deref(), Some("from-session"));
479 assert_eq!(get("TERM").as_deref(), Some("xterm-256color"));
480 assert_eq!(get("EXTRA").as_deref(), Some("x"));
481 assert_eq!(env.len(), 4, "nothing inherited from the process: {env:?}");
482 }
483
484 #[cfg(target_os = "linux")]
487 #[test]
488 #[serial_test::serial(synaps_base_dir)]
489 fn a_pty_shell_is_scoped_in_the_tool_slice() {
490 use crate::process_scope::test_support::{fake_env, fake_systemd_run, recorded};
491 let base = crate::test_env::BaseDirGuard::new();
492 let _scopes = crate::test_env::EnvVarGuard::remove("SYNAPS_PROCESS_SCOPES");
493 let dir = tempfile::tempdir().unwrap();
494 let fake = fake_systemd_run(dir.path());
495 let env: HashMap<String, String> = fake_env(&fake).into_iter().collect();
496 let cwd = dir.path().to_str().unwrap();
497 let slices = || {
498 let (_, _, launcher) = PtyHandle::prepare_scope(Some(cwd), None, &env);
499 let (_, args) = launcher.expect("the fake probe succeeds").wrap("synaps tool: shell", "/bin/sh", []);
500 args.iter()
501 .map(|a| a.to_string_lossy().into_owned())
502 .filter(|a| a.starts_with("--slice"))
503 .collect::<Vec<_>>()
504 };
505 assert!(slices().is_empty());
506 std::fs::write(
507 base.path().join("config"),
508 "tool_scope_slice = t_tools.slice\ndaemon_scope_slice = d_daemon.slice\n",
509 )
510 .unwrap();
511 assert_eq!(slices(), ["--slice=t_tools.slice"]);
512 assert!(recorded(dir.path(), "probe.args").contains(&"--slice=t_tools.slice".to_string()));
513 }
514
515 async fn read_until(handle: &mut PtyHandle, needle: &str) -> String {
519 let deadline = tokio::time::Instant::now() + Duration::from_secs(3);
520 let mut out = Vec::new();
521 while !String::from_utf8_lossy(&out).contains(needle)
522 && tokio::time::Instant::now() < deadline
523 {
524 out.extend(handle.try_read_output(Duration::from_millis(200)).await);
525 }
526 String::from_utf8_lossy(&out).into_owned()
527 }
528
529 async fn wait_released(accounting: &PtyAccounting, what: &str) {
535 let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
536 while !accounting.released() {
537 assert!(
538 tokio::time::Instant::now() < deadline,
539 "{what}: reader still running or bytes leaked \
540 (parties {}, retained {})",
541 accounting.parties.load(Ordering::Acquire),
542 accounting.retained()
543 );
544 tokio::time::sleep(Duration::from_millis(10)).await;
545 }
546 }
547
548 #[cfg(target_os = "linux")]
554 #[tokio::test]
555 #[serial(synaps_base_dir)]
556 async fn pty_shell_gets_its_own_scope_when_available() {
557 if std::env::var("SYNAPS_TEST_REAL_SCOPES").as_deref() != Ok("1") { return; }
558 let own = std::fs::read_to_string("/proc/self/cgroup").unwrap_or_default();
559 let mut handle = PtyHandle::spawn("bash --norc --noprofile", None, HashMap::new(), 24, 80)
562 .expect("spawn");
563 handle
564 .write(b"cat /proc/$$/cgroup; test -t 0 && echo tty=yes; exit\n")
565 .expect("write");
566 let mut out = Vec::new();
567 let deadline = std::time::Instant::now() + Duration::from_secs(5);
568 while std::time::Instant::now() < deadline
569 && !String::from_utf8_lossy(&out).contains("tty=yes\r")
570 {
571 out.extend(handle.try_read_output(Duration::from_millis(200)).await);
572 }
573 let accounting = Arc::clone(&handle.accounting);
574 drop(handle);
575 wait_released(&accounting, "scoped shell").await;
576 let text = String::from_utf8_lossy(&out).replace('\r', "");
577 assert!(
578 text.contains("tty=yes"),
579 "the shell must still have a terminal: {text}"
580 );
581 let child = text
584 .lines()
585 .filter(|l| !l.contains("cat /proc/")) .find_map(|l| l.find("0::/").map(|i| l[i..].trim().to_string()))
587 .unwrap_or_default();
588 assert!(child.starts_with("0::/"), "no cgroup line in: {text:?}");
589 if crate::process_scope::launcher_for_env(None).is_some() {
590 assert_ne!(child, own.trim(), "{text}");
591 assert!(child.ends_with(".scope"), "{child}");
592 } else {
593 assert_eq!(child, own.trim(), "unwrapped fallback");
594 }
595 }
596
597 #[tokio::test]
598 #[serial(synaps_base_dir)]
599 async fn pty_slow_consumer_retention_is_bounded_and_drop_releases_reader() {
600 let handle = PtyHandle::spawn(
601 "bash -c 'yes x | head -c 10485760'",
602 None,
603 HashMap::new(),
604 24,
605 80,
606 )
607 .expect("spawn producer");
608 let accounting = Arc::clone(&handle.accounting);
609 let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
613 while accounting.retained() == 0 {
614 assert!(
615 tokio::time::Instant::now() < deadline,
616 "stalled consumer must retain queued PTY bytes"
617 );
618 tokio::time::sleep(Duration::from_millis(10)).await;
619 }
620 tokio::time::sleep(Duration::from_millis(100)).await;
622 let retained = accounting.retained();
623 assert!(
624 retained <= ((PTY_OUTPUT_CHANNEL_CAPACITY + 1) * PTY_OUTPUT_MAX_CHUNK_BYTES) as u64,
625 "retained {retained}"
626 );
627 drop(handle);
628 wait_released(&accounting, "drop must release the reader and its queued bytes").await;
629 }
630
631 #[tokio::test]
637 #[serial(synaps_base_dir)]
638 async fn pty_drop_races_reader_and_always_conserves_accounting() {
639 let mut released = Vec::new();
640 for _ in 0..5 {
641 let mut handle = PtyHandle::spawn(
642 "bash -c 'yes x | head -c 33554432'",
643 None,
644 HashMap::new(),
645 24,
646 80,
647 )
648 .expect("spawn producer");
649 let _ = handle.try_read_output(Duration::from_millis(50)).await;
654 released.push(Arc::clone(&handle.accounting));
655 drop(handle);
656 }
657 for accounting in &released {
660 wait_released(accounting, "drop/reader race must conserve accounting").await;
661 }
662 }
663
664 #[tokio::test]
669 #[serial(synaps_base_dir)]
670 async fn pty_preempted_accounting_across_drop_is_released_exactly_once() {
671 TEST_PAUSE_AFTER_SEND_MS.store(150, Ordering::Relaxed);
672 let handle = PtyHandle::spawn("bash -c 'echo leak-window'", None, HashMap::new(), 24, 80)
673 .expect("spawn producer");
674 let accounting = Arc::clone(&handle.accounting);
675 tokio::time::sleep(Duration::from_millis(40)).await;
677 drop(handle); TEST_PAUSE_AFTER_SEND_MS.store(0, Ordering::Relaxed);
679
680 wait_released(&accounting, "a drop inside the accounting gap must not leak").await;
681 }
682
683 #[tokio::test]
688 #[serial(synaps_base_dir)]
689 async fn pty_concurrent_handles_release_exactly_once_per_handle() {
690 let mut tasks = Vec::new();
691 for i in 0..8 {
692 tasks.push(tokio::spawn(async move {
693 let mut handle = PtyHandle::spawn(
694 "bash -c 'yes x | head -c 262144'",
695 None,
696 HashMap::new(),
697 24,
698 80,
699 )
700 .expect("spawn producer");
701 match i % 3 {
702 0 => {} 1 => {
704 tokio::time::sleep(Duration::from_millis(5)).await;
705 }
706 _ => {
707 let _ = handle.try_read_output(Duration::from_millis(10)).await;
709 }
710 }
711 let accounting = Arc::clone(&handle.accounting);
712 drop(handle);
713 accounting
714 }));
715 }
716 for t in tasks {
717 let accounting = t.await.unwrap();
718 wait_released(&accounting, "concurrent handles must each release exactly once").await;
719 }
720 }
721
722 #[tokio::test]
723 #[serial(synaps_base_dir)]
724 async fn test_spawn_echo_hello() {
725 let mut handle = PtyHandle::spawn("echo hello", None, HashMap::new(), 24, 80)
726 .expect("failed to spawn echo");
727
728 let text = read_until(&mut handle, "hello").await;
729 assert!(
730 text.contains("hello"),
731 "expected 'hello' in output, got: {text:?}"
732 );
733 }
734
735 #[tokio::test]
736 #[serial(synaps_base_dir)]
737 async fn test_cat_echo_back() {
738 let mut handle =
739 PtyHandle::spawn("cat", None, HashMap::new(), 24, 80).expect("failed to spawn cat");
740
741 handle.write(b"test\n").expect("write failed");
743
744 let text = read_until(&mut handle, "test").await;
745 assert!(
746 text.contains("test"),
747 "expected 'test' in output, got: {text:?}"
748 );
749 }
750
751 #[tokio::test]
752 #[serial(synaps_base_dir)]
753 async fn test_exit_code_detection() {
754 let mut handle = PtyHandle::spawn("bash -c exit 42", None, HashMap::new(), 24, 80)
755 .expect("failed to spawn bash exit");
756
757 let _ = handle.try_read_output(Duration::from_secs(3)).await;
759
760 tokio::time::sleep(Duration::from_millis(200)).await;
762
763 assert!(!handle.is_alive(), "expected process to have exited");
764 }
765}