fix(ssh): harden keepalives and tunnel shutdown

Author: drew

Cargo.lock

@@ -82,6 +82,7 @@ pin-project-lite = "0.2"
 tokio-stream = "0.1"
 protobuf-src = "1.1.0"
 url = "2"
+socket2 = "0.6"
 
 # Database
 sqlx = { version = "0.8", features = ["runtime-tokio-rustls", "postgres", "sqlite", "migrate"] }

Cargo.toml

@@ -141,11 +141,22 @@ sequenceDiagram
 5. When SSH starts, it spawns the `ssh-proxy` subprocess as its `ProxyCommand`.
 6. `crates/openshell-cli/src/ssh.rs` -- `sandbox_ssh_proxy()`:
    - Parses the gateway URL, connects via TCP (plain) or TLS (mTLS)
+   - Enables TCP keepalive on the gateway socket
    - Sends a raw HTTP CONNECT request with `X-Sandbox-Id` and `X-Sandbox-Token` headers
    - Reads the response status line; proceeds if 200
    - Spawns two `tokio::spawn` tasks for bidirectional copy between stdin/stdout and the gateway stream
    - When the remote-to-stdout direction completes, aborts the stdin-to-remote task (SSH has all the data it needs)
 
+### Connection stability
+
+Recent SSH stability hardening is split across the client, gateway, sandbox, and edge tunnel paths:
+
+- **OpenSSH keepalives**: the CLI now sets `ServerAliveInterval=30` and `ServerAliveCountMax=3` on every SSH invocation so idle sessions still emit SSH traffic.
+- **TCP keepalive**: the CLI-to-gateway and gateway-to-sandbox TCP sockets enable 30-second keepalive probes to reduce drops from NAT, load balancers, and other idle-sensitive middleboxes.
+- **Sandbox SSH daemon**: the embedded `russh` server disables its default 10-minute inactivity timeout and instead sends protocol keepalives every 30 seconds. This prevents quiet shells from being garbage-collected while still detecting dead peers.
+- **Edge WebSocket tunnel**: the WebSocket bridge now lets both copy directions observe shutdown instead of aborting the peer task immediately, which reduces abrupt closes and truncated tail data.
+- **Limit diagnostics**: when the gateway rejects a connection because the per-session or per-sandbox cap is reached, it now logs the active count and configured limit to make 429s easier to diagnose.
+
 ### Command Execution (CLI)
 
 The `sandbox exec` path is identical to interactive connect except:

architecture/sandbox-connect.md

@@ -126,24 +126,18 @@ async fn handle_connection(tcp_stream: TcpStream, config: &TunnelConfig) -> Resu
     let (ws_sink, ws_source) = ws_stream.split();
     let (tcp_read, tcp_write) = tokio::io::split(tcp_stream);
 
-    // Two tasks: TCP->WS and WS->TCP. Abort the peer task once either
-    // direction finishes so the connection tears down promptly.
-    let mut tcp_to_ws = tokio::spawn(copy_tcp_to_ws(tcp_read, ws_sink));
-    let mut ws_to_tcp = tokio::spawn(copy_ws_to_tcp(ws_source, tcp_write));
-
-    tokio::select! {
-        res = &mut tcp_to_ws => {
-            if let Err(e) = res {
-                debug!(error = %e, "tcp->ws task panicked");
-            }
-            ws_to_tcp.abort();
-        }
-        res = &mut ws_to_tcp => {
-            if let Err(e) = res {
-                debug!(error = %e, "ws->tcp task panicked");
-            }
-            tcp_to_ws.abort();
-        }
+    // Keep both directions alive until each side has observed shutdown. This
+    // avoids cutting off trailing bytes when one half closes slightly earlier
+    // than the other.
+    let tcp_to_ws = tokio::spawn(copy_tcp_to_ws(tcp_read, ws_sink));
+    let ws_to_tcp = tokio::spawn(copy_ws_to_tcp(ws_source, tcp_write));
+
+    let (tcp_to_ws_res, ws_to_tcp_res) = tokio::join!(tcp_to_ws, ws_to_tcp);
+    if let Err(e) = tcp_to_ws_res {
+        debug!(error = %e, "tcp->ws task panicked");
+    }
+    if let Err(e) = ws_to_tcp_res {
+        debug!(error = %e, "ws->tcp task panicked");
     }
 
     Ok(())

crates/openshell-cli/src/edge_tunnel.rs

@@ -10,6 +10,7 @@ use nix::sys::signal::{SaFlags, SigAction, SigHandler, SigSet, Signal, sigaction
 use openshell_core::forward::{
     find_ssh_forward_pid, resolve_ssh_gateway, shell_escape, write_forward_pid,
 };
+use openshell_core::net::enable_tcp_keepalive;
 use openshell_core::proto::{CreateSshSessionRequest, GetSandboxRequest};
 use owo_colors::OwoColorize;
 use rustls::pki_types::ServerName;
@@ -56,6 +57,8 @@ impl Editor {
 
 const SSH_PROXY_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
 const SSH_PROXY_STATUS_TIMEOUT: Duration = Duration::from_secs(10);
+const SSH_SERVER_ALIVE_INTERVAL_SECS: u64 = 30;
+const SSH_SERVER_ALIVE_COUNT_MAX: u64 = 3;
 
 struct SshSessionConfig {
     proxy_command: String,
@@ -145,6 +148,12 @@ fn ssh_base_command(proxy_command: &str) -> Command {
         .arg("-o")
         .arg("GlobalKnownHostsFile=/dev/null")
         .arg("-o")
+        .arg(format!(
+            "ServerAliveInterval={SSH_SERVER_ALIVE_INTERVAL_SECS}"
+        ))
+        .arg("-o")
+        .arg(format!("ServerAliveCountMax={SSH_SERVER_ALIVE_COUNT_MAX}"))
+        .arg("-o")
         .arg("LogLevel=ERROR");
     command
 }
@@ -722,8 +731,13 @@ pub async fn sandbox_ssh_proxy(
     .await
     .map_err(|_| miette::miette!("timed out waiting for gateway CONNECT response"))??;
     if status != 200 {
+        let reason = match status {
+            401 => " (SSH session expired, was revoked, or is invalid)",
+            429 => " (too many concurrent SSH connections for this session or sandbox)",
+            _ => "",
+        };
         return Err(miette::miette!(
-            "gateway CONNECT failed with status {status}"
+            "gateway CONNECT failed with status {status}{reason}"
         ));
     }
 
@@ -776,7 +790,7 @@ fn render_ssh_config(gateway: &str, name: &str) -> String {
     let proxy_cmd = format!("{exe} ssh-proxy --gateway-name {gateway} --name {name}");
     let host_alias = host_alias(name);
     format!(
-        "Host {host_alias}\n    User sandbox\n    StrictHostKeyChecking no\n    UserKnownHostsFile /dev/null\n    GlobalKnownHostsFile /dev/null\n    LogLevel ERROR\n    ProxyCommand {proxy_cmd}\n"
+        "Host {host_alias}\n    User sandbox\n    StrictHostKeyChecking no\n    UserKnownHostsFile /dev/null\n    GlobalKnownHostsFile /dev/null\n    ServerAliveInterval {SSH_SERVER_ALIVE_INTERVAL_SECS}\n    ServerAliveCountMax {SSH_SERVER_ALIVE_COUNT_MAX}\n    LogLevel ERROR\n    ProxyCommand {proxy_cmd}\n"
     )
 }
 
@@ -1000,6 +1014,7 @@ async fn connect_gateway(
         .map_err(|_| miette::miette!("timed out connecting to edge tunnel proxy"))?
         .into_diagnostic()?;
         tcp.set_nodelay(true).into_diagnostic()?;
+        let _ = enable_tcp_keepalive(&tcp);
         return Ok(Box::new(tcp));
     }
 
@@ -1008,6 +1023,7 @@ async fn connect_gateway(
         .map_err(|_| miette::miette!("timed out connecting to SSH gateway"))?
         .into_diagnostic()?;
     tcp.set_nodelay(true).into_diagnostic()?;
+    let _ = enable_tcp_keepalive(&tcp);
     if scheme.eq_ignore_ascii_case("https") {
         let materials = require_tls_materials(&format!("https://{host}:{port}"), tls)?;
         let config = build_rustls_config(&materials)?;
@@ -1167,4 +1183,19 @@ mod tests {
         assert!(message.contains("Forwarding port 3000 to sandbox demo"));
         assert!(message.contains("Access at: http://localhost:3000/"));
     }
+
+    #[test]
+    fn ssh_base_command_enables_server_keepalives() {
+        let command = ssh_base_command("openshell ssh-proxy");
+        let args = command
+            .get_args()
+            .map(|arg| arg.to_string_lossy().into_owned())
+            .collect::<Vec<_>>();
+
+        assert!(args.contains(&format!(
+            "ServerAliveInterval={SSH_SERVER_ALIVE_INTERVAL_SECS}"
+        )));
+        assert!(args.contains(&format!("ServerAliveCountMax={SSH_SERVER_ALIVE_COUNT_MAX}")));
+    }
+}
 }

crates/openshell-cli/src/ssh.rs

@@ -19,6 +19,7 @@ miette = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 url = { workspace = true }
+socket2 = { workspace = true }
 
 [build-dependencies]
 tonic-build = { workspace = true }

crates/openshell-core/Cargo.toml

@@ -13,6 +13,7 @@ pub mod config;
 pub mod error;
 pub mod forward;
 pub mod inference;
+pub mod net;
 pub mod paths;
 pub mod proto;
 

crates/openshell-core/src/lib.rs

@@ -0,0 +1,60 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Shared socket configuration helpers.
+
+use socket2::{SockRef, TcpKeepalive};
+use std::io;
+use std::time::Duration;
+
+/// Idle time before TCP keepalive probes start.
+pub const TCP_KEEPALIVE_IDLE: Duration = Duration::from_secs(30);
+
+/// Interval between TCP keepalive probes on supported platforms.
+pub const TCP_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30);
+
+fn default_keepalive() -> TcpKeepalive {
+    let keepalive = TcpKeepalive::new().with_time(TCP_KEEPALIVE_IDLE);
+    #[cfg(any(
+        target_os = "android",
+        target_os = "dragonfly",
+        target_os = "freebsd",
+        target_os = "fuchsia",
+        target_os = "illumos",
+        target_os = "ios",
+        target_os = "visionos",
+        target_os = "linux",
+        target_os = "macos",
+        target_os = "netbsd",
+        target_os = "tvos",
+        target_os = "watchos",
+        target_os = "windows",
+        target_os = "cygwin",
+    ))]
+    let keepalive = keepalive.with_interval(TCP_KEEPALIVE_INTERVAL);
+    keepalive
+}
+
+/// Enable aggressive TCP keepalive on a socket.
+#[cfg(unix)]
+pub fn enable_tcp_keepalive<S>(socket: &S) -> io::Result<()>
+where
+    S: std::os::fd::AsFd,
+{
+    SockRef::from(socket).set_tcp_keepalive(&default_keepalive())
+}
+
+/// Enable aggressive TCP keepalive on a socket.
+#[cfg(windows)]
+pub fn enable_tcp_keepalive<S>(socket: &S) -> io::Result<()>
+where
+    S: std::os::windows::io::AsSocket,
+{
+    SockRef::from(socket).set_tcp_keepalive(&default_keepalive())
+}
+
+/// Enable aggressive TCP keepalive on a socket.
+#[cfg(not(any(unix, windows)))]
+pub fn enable_tcp_keepalive<S>(_socket: &S) -> io::Result<()> {
+    Ok(())
+}

crates/openshell-core/src/net.rs

@@ -10,6 +10,7 @@ use crate::sandbox;
 #[cfg(target_os = "linux")]
 use crate::{register_managed_child, unregister_managed_child};
 use miette::{IntoDiagnostic, Result};
+use openshell_core::net::enable_tcp_keepalive;
 use nix::pty::{Winsize, openpty};
 use nix::unistd::setsid;
 use rand_core::OsRng;
@@ -33,6 +34,20 @@ const PREFACE_MAGIC: &str = "NSSH1";
 const SSH_HANDSHAKE_SECRET_ENV: &str = "OPENSHELL_SSH_HANDSHAKE_SECRET";
 
 const SSH_PREFACE_TIMEOUT: Duration = Duration::from_secs(10);
+const SSH_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30);
+const SSH_KEEPALIVE_MAX_MISSES: usize = 3;
+
+fn build_ssh_server_config(host_key: PrivateKey) -> russh::server::Config {
+    russh::server::Config {
+        auth_rejection_time: Duration::from_secs(1),
+        inactivity_timeout: None,
+        keepalive_interval: Some(SSH_KEEPALIVE_INTERVAL),
+        keepalive_max: SSH_KEEPALIVE_MAX_MISSES,
+        nodelay: true,
+        keys: vec![host_key],
+        ..Default::default()
+    }
+}
 
 /// A time-bounded set of nonces used to detect replayed NSSH1 handshakes.
 /// Each entry records the `Instant` it was inserted; a background reaper task
@@ -52,14 +67,7 @@ async fn ssh_server_init(
 )> {
     let mut rng = OsRng;
     let host_key = PrivateKey::random(&mut rng, Algorithm::Ed25519).into_diagnostic()?;
-
-    let mut config = russh::server::Config {
-        auth_rejection_time: Duration::from_secs(1),
-        ..Default::default()
-    };
-    config.keys.push(host_key);
-
-    let config = Arc::new(config);
+    let config = Arc::new(build_ssh_server_config(host_key));
     let ca_paths = ca_file_paths.as_ref().map(|p| Arc::new(p.clone()));
     let listener = TcpListener::bind(listen_addr).await.into_diagnostic()?;
     info!(addr = %listen_addr, "SSH server listening");
@@ -163,6 +171,9 @@ async fn handle_connection(
     nonce_cache: &NonceCache,
 ) -> Result<()> {
     info!(peer = %peer, "SSH connection: reading handshake preface");
+    if let Err(err) = enable_tcp_keepalive(&stream) {
+        warn!(peer = %peer, error = %err, "SSH connection: failed to enable TCP keepalive");
+    }
     let mut line = String::new();
     tokio::time::timeout(SSH_PREFACE_TIMEOUT, read_line(&mut stream, &mut line))
         .await
@@ -1466,4 +1477,17 @@ mod tests {
         assert!(!stdout.contains(SSH_HANDSHAKE_SECRET_ENV));
         assert!(stdout.contains("ANTHROPIC_API_KEY=openshell:resolve:env:ANTHROPIC_API_KEY"));
     }
+
+    #[test]
+    fn ssh_server_config_keeps_idle_sessions_alive() {
+        let mut rng = OsRng;
+        let host_key = PrivateKey::random(&mut rng, Algorithm::Ed25519).unwrap();
+        let config = build_ssh_server_config(host_key);
+
+        assert_eq!(config.inactivity_timeout, None);
+        assert_eq!(config.keepalive_interval, Some(SSH_KEEPALIVE_INTERVAL));
+        assert_eq!(config.keepalive_max, SSH_KEEPALIVE_MAX_MISSES);
+        assert!(config.nodelay);
+        assert_eq!(config.keys.len(), 1);
+    }
 }

crates/openshell-sandbox/src/ssh.rs

@@ -8,6 +8,7 @@ use http::StatusCode;
 use hyper::Request;
 use hyper::upgrade::Upgraded;
 use hyper_util::rt::TokioIo;
+use openshell_core::net::enable_tcp_keepalive;
 use openshell_core::proto::{Sandbox, SandboxPhase, SshSession};
 use prost::Message;
 use std::net::SocketAddr;
@@ -118,7 +119,7 @@ async fn ssh_connect(
         let mut counts = state.ssh_connections_by_token.lock().unwrap();
         let count = counts.entry(token.clone()).or_insert(0);
         if *count >= MAX_CONNECTIONS_PER_TOKEN {
-            warn!(token = %token, "SSH tunnel: per-token connection limit reached");
+            warn!(sandbox_id = %sandbox_id, active_connections = *count, limit = MAX_CONNECTIONS_PER_TOKEN, "SSH tunnel: per-token connection limit reached");
             return StatusCode::TOO_MANY_REQUESTS.into_response();
         }
         *count += 1;
@@ -137,7 +138,7 @@ async fn ssh_connect(
                     token_counts.remove(&token);
                 }
             }
-            warn!(sandbox_id = %sandbox_id, "SSH tunnel: per-sandbox connection limit reached");
+            warn!(sandbox_id = %sandbox_id, active_connections = *count, limit = MAX_CONNECTIONS_PER_SANDBOX, "SSH tunnel: per-sandbox connection limit reached");
             return StatusCode::TOO_MANY_REQUESTS.into_response();
         }
         *count += 1;
@@ -265,6 +266,9 @@ async fn establish_upstream_with_timeouts(
     upstream
         .set_nodelay(true)
         .map_err(|err| TunnelSetupError::Other(err.to_string()))?;
+    if let Err(err) = enable_tcp_keepalive(&upstream) {
+        warn!(sandbox_id = %sandbox_id, error = %err, "SSH tunnel: failed to enable upstream TCP keepalive");
+    }
     info!(sandbox_id = %sandbox_id, "SSH tunnel: sending NSSH1 handshake preface");
     let preface =
         build_preface(token, secret).map_err(|err| TunnelSetupError::Other(err.to_string()))?;
@@ -296,9 +300,26 @@ async fn bridge_tunnel(
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     info!(sandbox_id = %sandbox_id, "SSH tunnel established");
     let mut upgraded = TokioIo::new(upgraded);
-    // Discard the result entirely – connection-close errors are expected when
-    // the SSH session ends and do not represent a failure worth propagating.
-    let _ = tokio::io::copy_bidirectional(&mut upgraded, &mut upstream).await;
+    match tokio::io::copy_bidirectional(&mut upgraded, &mut upstream).await {
+        Ok((client_to_sandbox, sandbox_to_client)) => {
+            info!(sandbox_id = %sandbox_id, client_to_sandbox_bytes = client_to_sandbox, sandbox_to_client_bytes = sandbox_to_client, "SSH tunnel closed");
+        }
+        Err(err)
+            if matches!(
+                err.kind(),
+                std::io::ErrorKind::BrokenPipe
+                    | std::io::ErrorKind::ConnectionAborted
+                    | std::io::ErrorKind::ConnectionReset
+                    | std::io::ErrorKind::NotConnected
+                    | std::io::ErrorKind::UnexpectedEof
+            ) =>
+        {
+            info!(sandbox_id = %sandbox_id, error = %err, "SSH tunnel closed during shutdown");
+        }
+        Err(err) => {
+            warn!(sandbox_id = %sandbox_id, error = %err, "SSH tunnel I/O failure");
+        }
+    }
     // Gracefully shut down the write-half of the upgraded connection so the
     // client receives a clean EOF instead of a TCP RST.  This gives SSH time
     // to read any remaining protocol data (e.g. exit-status) from its buffer.