fix(cli): validate background forward cleanup pid

Signed-off-by: Shiju <shiju@nvidia.com>

Author: shiju-nv

crates/openshell-cli/src/ssh.rs

@@ -10,6 +10,8 @@ use nix::sys::signal::{SaFlags, SigAction, SigHandler, SigSet, Signal, sigaction
 #[cfg(unix)]
 use nix::unistd::Pid;
 use openshell_core::ObjectId;
+#[cfg(unix)]
+use openshell_core::forward::pid_matches_forward;
 use openshell_core::forward::{
     ForwardSpec, build_proxy_command, find_ssh_forward_pid, format_gateway_url,
     resolve_ssh_gateway, shell_escape, validate_ssh_session_response, write_forward_pid,
@@ -31,11 +33,15 @@ use tokio::net::TcpStream;
 use tokio::process::{Child, Command as TokioCommand};
 use tokio_stream::wrappers::ReceiverStream;
 
-/// Time budget for a forward to become healthy after `ssh` starts: covers both
-/// backgrounded-PID discovery and listener readiness, in foreground and
-/// background.
-const FORWARD_STARTUP_GRACE_PERIOD: Duration = Duration::from_secs(2);
-/// Delay between listener/PID probes within the grace period.
+/// Time budget for finding the OpenSSH background process after `ssh -f`
+/// returns. PID discovery is separate from listener readiness so missing
+/// process tracking still fails quickly.
+const FORWARD_PID_DISCOVERY_TIMEOUT: Duration = Duration::from_secs(2);
+/// Time budget for the local listener to become reachable after `ssh` starts.
+/// This is a user-visible readiness deadline for both foreground and background
+/// forwards, not a soft cleanup grace period.
+const FORWARD_LISTENER_READINESS_TIMEOUT: Duration = Duration::from_secs(10);
+/// Delay between listener/PID probes within the configured timeout.
 const FORWARD_LISTENER_PROBE_INTERVAL: Duration = Duration::from_millis(50);
 /// Per-attempt connect timeout, so one hung probe cannot consume the whole
 /// grace period.
@@ -378,11 +384,11 @@ pub async fn sandbox_forward(
                 )
             })?;
 
-        if let Err(err) = wait_for_forward_listener(spec, FORWARD_STARTUP_GRACE_PERIOD)
+        if let Err(err) = wait_for_forward_listener(spec, FORWARD_LISTENER_READINESS_TIMEOUT)
             .await
             .wrap_err("ssh process started but local forward listener was not reachable")
         {
-            terminate_forward_pid(pid);
+            terminate_forward_pid(pid, port, &session.sandbox_id);
             return Err(err);
         }
 
@@ -412,7 +418,7 @@ pub async fn sandbox_forward(
 /// session) means forwarding never came up, so it errors instead of waiting
 /// out the grace period.
 async fn wait_for_foreground_forward_start(child: &mut Child, spec: &ForwardSpec) -> Result<()> {
-    let listener = wait_for_forward_listener(spec, FORWARD_STARTUP_GRACE_PERIOD);
+    let listener = wait_for_forward_listener(spec, FORWARD_LISTENER_READINESS_TIMEOUT);
     tokio::pin!(listener);
     tokio::select! {
         result = &mut listener => result,
@@ -439,7 +445,7 @@ async fn wait_for_foreground_forward_start(child: &mut Child, spec: &ForwardSpec
 /// so the PID is unknown when `command.status()` returns and must be discovered
 /// afterward. Returns `None` if it never appears within the grace period.
 async fn wait_for_ssh_forward_pid(sandbox_id: &str, port: u16) -> Option<u32> {
-    let deadline = tokio::time::Instant::now() + FORWARD_STARTUP_GRACE_PERIOD;
+    let deadline = tokio::time::Instant::now() + FORWARD_PID_DISCOVERY_TIMEOUT;
     loop {
         if let Some(pid) = find_ssh_forward_pid(sandbox_id, port) {
             return Some(pid);
@@ -513,20 +519,26 @@ fn forward_probe_host(spec: &ForwardSpec) -> &str {
 /// are ignored: the process may already be exiting, and the caller surfaces the
 /// original listener error regardless.
 #[cfg(unix)]
-fn terminate_forward_pid(pid: u32) {
+fn terminate_forward_pid(pid: u32, port: u16, sandbox_id: &str) {
     let Ok(raw_pid) = i32::try_from(pid) else {
         return;
     };
     if raw_pid <= 0 {
         return;
     }
+    if !pid_matches_forward(pid, port, Some(sandbox_id)) {
+        // The PID came from a process-table scan, not a file we own. Re-check
+        // immediately before signaling so a stale or spoofed match is left
+        // untouched instead of terminating an unrelated process.
+        return;
+    }
 
     let _ = nix::sys::signal::kill(Pid::from_raw(raw_pid), Signal::SIGTERM);
 }
 
 /// Non-Unix builds cannot manage OpenSSH process IDs with Unix signals.
 #[cfg(not(unix))]
-fn terminate_forward_pid(_pid: u32) {}
+fn terminate_forward_pid(_pid: u32, _port: u16, _sandbox_id: &str) {}
 
 fn foreground_forward_started_message(name: &str, spec: &ForwardSpec) -> String {
     format!(
@@ -1793,6 +1805,48 @@ mod tests {
         assert!(text.contains("local forward listener did not open"));
     }
 
+    #[cfg(unix)]
+    #[test]
+    fn terminate_forward_pid_skips_process_that_no_longer_matches_forward() {
+        let dir = tempfile::tempdir().unwrap();
+        let terminated_path = dir.path().join("terminated");
+        let mut child = Command::new("python3")
+            .arg("-c")
+            .arg(
+                r#"
+import pathlib
+import signal
+import sys
+import time
+
+terminated_path = pathlib.Path(sys.argv[1])
+
+def stop(_signum, _frame):
+    terminated_path.write_text("terminated")
+    raise SystemExit(0)
+
+signal.signal(signal.SIGTERM, stop)
+
+while True:
+    time.sleep(1)
+"#,
+            )
+            .arg(&terminated_path)
+            .spawn()
+            .unwrap();
+        std::thread::sleep(Duration::from_millis(100));
+
+        terminate_forward_pid(child.id(), 43210, "id-spoofed-forward");
+        std::thread::sleep(Duration::from_millis(200));
+
+        assert!(
+            !terminated_path.exists(),
+            "mismatched process should not receive SIGTERM"
+        );
+        let _ = child.kill();
+        let _ = child.wait();
+    }
+
     #[test]
     fn split_sandbox_path_separates_parent_and_basename() {
         assert_eq!(

crates/openshell-cli/tests/sandbox_create_lifecycle_integration.rs

@@ -727,10 +727,27 @@ fn install_fake_pgrep_no_match(dir: &TempDir) -> std::path::PathBuf {
     install_executable_script(dir, "pgrep", "#!/bin/sh\nexit 1\n")
 }
 
-async fn wait_for_file(path: &std::path::Path, timeout: Duration) -> bool {
+async fn wait_for_process_exit(pid: u32, timeout: Duration) -> bool {
     let deadline = Instant::now() + timeout;
     loop {
-        if path.exists() {
+        let output = std::process::Command::new("ps")
+            .arg("-o")
+            .arg("stat=")
+            .arg("-p")
+            .arg(pid.to_string())
+            .stderr(std::process::Stdio::null())
+            .output();
+        let alive = output.is_ok_and(|output| {
+            if !output.status.success() {
+                return false;
+            }
+            let stat = String::from_utf8_lossy(&output.stdout);
+            // Linux can leave the orphaned fake forward as a short-lived zombie
+            // until the container's init process reaps it. A zombie has already
+            // exited, so it satisfies this cleanup assertion.
+            !stat.trim_start().starts_with('Z')
+        });
+        if !alive {
             return true;
         }
         if Instant::now() >= deadline {
@@ -844,39 +861,87 @@ exit 1
     ssh_path
 }
 
-fn install_fake_unreachable_forwarding_ssh(dir: &TempDir) -> std::path::PathBuf {
+struct FakeUnreachableForward {
+    log_path: std::path::PathBuf,
+    pid_path: std::path::PathBuf,
+}
+
+fn install_fake_unreachable_forwarding_ssh(dir: &TempDir) -> FakeUnreachableForward {
+    let log_path = dir.path().join("fake-forward.log");
     let pid_path = dir.path().join("fake-forward.pid");
-    let terminated_path = dir.path().join("fake-forward.terminated");
+    let ready_path = dir.path().join("fake-forward.ready");
     install_executable_script(
         dir,
         "ssh",
         r#"#!/bin/sh
 set -eu
 
-nohup python3 -c '
+forward=""
+sandbox_id=""
+previous=""
+
+for arg in "$@"; do
+  if [ "$previous" = "-L" ]; then
+    forward="$arg"
+    previous=""
+    continue
+  fi
+
+  if [ "$previous" = "-o" ]; then
+    case "$arg" in
+      ProxyCommand=*)
+        sandbox_id="$(printf '%s\n' "$arg" | sed -n 's/.*--sandbox-id \([^ ]*\).*/\1/p')"
+        ;;
+    esac
+    previous=""
+    continue
+  fi
+
+  case "$arg" in
+    -L|-o)
+      previous="$arg"
+      ;;
+  esac
+done
+
+if [ -z "$forward" ] || [ -z "$sandbox_id" ]; then
+  exit 1
+fi
+
+trap '' HUP
+python3 -c '
 import pathlib
 import signal
 import sys
 import time
 
-terminated_path = pathlib.Path(sys.argv[1])
-
-def stop(_signum, _frame):
-    terminated_path.write_text("terminated")
-    raise SystemExit(0)
+ready_path = pathlib.Path(sys.argv[1])
 
-signal.signal(signal.SIGTERM, stop)
-signal.signal(signal.SIGINT, stop)
 signal.signal(signal.SIGHUP, signal.SIG_IGN)
 
+ready_path.write_text("ready")
+
 while True:
     time.sleep(1)
-' '@TERMINATED_PATH@' >/dev/null 2>&1 &
-echo $! > '@PID_PATH@'
+' '@READY_PATH@' ssh ssh-proxy --sandbox-id "$sandbox_id" -L "$forward" >'@LOG_PATH@' 2>&1 &
+pid="$!"
+i=0
+while [ "$i" -lt 100 ]; do
+  if [ -e '@READY_PATH@' ]; then
+    break
+  fi
+  i=$((i + 1))
+  sleep 0.05
+done
+if [ ! -e '@READY_PATH@' ]; then
+  exit 1
+fi
+echo "$pid" > '@PID_PATH@'
 
 exit 0
 "#
-        .replace("@TERMINATED_PATH@", &terminated_path.display().to_string())
+        .replace("@LOG_PATH@", &log_path.display().to_string())
+        .replace("@READY_PATH@", &ready_path.display().to_string())
         .replace("@PID_PATH@", &pid_path.display().to_string()),
     );
 
@@ -896,7 +961,7 @@ exit 1
         ),
     );
 
-    terminated_path
+    FakeUnreachableForward { log_path, pid_path }
 }
 
 fn test_env(fake_ssh_dir: &TempDir, xdg_dir: &TempDir) -> EnvVarGuard {
@@ -1553,14 +1618,37 @@ async fn sandbox_forward_background_fails_when_pid_is_not_discoverable() {
     );
 }
 
+#[tokio::test]
+async fn sandbox_forward_foreground_fails_when_ssh_exits_before_listener_opens() {
+    let server = run_server().await;
+    let fake_ssh_dir = tempfile::tempdir().unwrap();
+    let xdg_dir = tempfile::tempdir().unwrap();
+    let _env = test_env(&fake_ssh_dir, &xdg_dir);
+    let tls = test_tls(&server);
+    install_fake_ssh(&fake_ssh_dir);
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let forward_port = listener.local_addr().unwrap().port();
+    drop(listener);
+
+    let spec = openshell_core::forward::ForwardSpec::new(forward_port);
+    let err = run::sandbox_forward(&server.endpoint, "foreground-forward", &spec, false, &tls)
+        .await
+        .expect_err("foreground forward should fail when ssh exits before listener readiness");
+    let msg = format!("{err}");
+    assert!(
+        msg.contains("ssh exited before local forward listener opened"),
+        "error should explain that ssh exited before listener readiness, got: {msg}",
+    );
+}
+
 #[tokio::test]
 async fn sandbox_forward_background_terminates_discovered_pid_when_listener_never_opens() {
     let server = run_server().await;
     let fake_ssh_dir = tempfile::tempdir().unwrap();
     let xdg_dir = tempfile::tempdir().unwrap();
     let _env = test_env(&fake_ssh_dir, &xdg_dir);
     let tls = test_tls(&server);
-    let terminated_path = install_fake_unreachable_forwarding_ssh(&fake_ssh_dir);
+    let fake_forward = install_fake_unreachable_forwarding_ssh(&fake_ssh_dir);
     let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
     let forward_port = listener.local_addr().unwrap().port();
     drop(listener);
@@ -1578,10 +1666,30 @@ async fn sandbox_forward_background_terminates_discovered_pid_when_listener_neve
         openshell_core::forward::read_forward_pid("unreachable-forward", forward_port).is_none(),
         "unreachable background forwards must not write a PID file",
     );
-    assert!(
-        wait_for_file(&terminated_path, Duration::from_secs(2)).await,
-        "discovered background SSH process should be terminated after listener failure",
-    );
+    let pid = fs::read_to_string(&fake_forward.pid_path)
+        .expect("fake forward should record a PID")
+        .trim()
+        .parse::<u32>()
+        .expect("fake forward PID should be numeric");
+    if !wait_for_process_exit(pid, Duration::from_secs(2)).await {
+        let log = fs::read_to_string(&fake_forward.log_path).unwrap_or_default();
+        let command = std::process::Command::new("ps")
+            .arg("-ww")
+            .arg("-o")
+            .arg("command=")
+            .arg("-p")
+            .arg(pid.to_string())
+            .output()
+            .ok()
+            .map(|output| String::from_utf8_lossy(&output.stdout).to_string())
+            .unwrap_or_default();
+        panic!(
+            "discovered background SSH process should exit after listener failure cleanup; pid={}, command={}, log={}",
+            pid,
+            command.trim(),
+            log.trim(),
+        );
+    }
 }
 
 #[tokio::test]