From f4c25725c5259f7f68fa4a2940c4ad0cc701eefa Mon Sep 17 00:00:00 2001 From: l0ng-ai <24760907+l0ng-ai@users.noreply.github.com> Date: Tue, 14 Jul 2026 02:52:36 +0800 Subject: [PATCH] fix(ssh): self-heal reuse of a connection whose transport silently died mark_dead() only runs from Drop, but a parked forward/loopback accept loop holds an Arc, so a dead connection's Drop never runs and is_alive() stayed true. A reconnect for the same ConnectionKey reused the dead russh handle, the first channel-open errored, and the whole reconnect failed until forwards were torn down. Two complementary fixes: - is_alive() now also consults the russh handle's own liveness via a non-blocking try_lock + handle.is_closed() (the session task ending closes its command sender), catching the stale-flag case cheaply. - run_session treats the first shell-channel open on a *reused* connection as a liveness probe: on failure it marks the connection dead, evicts its registry slot, and reconnects fresh once (a fresh connection failing there is a real error). Preconfigured forwards now establish after this probe, on the confirmed-live connection. open_connection returns a `reused` flag to drive this. Adds a unit test that evicting a key from the registry map clears its slot. The end-to-end reuse-after-death path needs a live server, so it stays covered by E2E. --- src/daemon/ssh/mod.rs | 90 +++++++++++++++++++++++++++++++++------ src/daemon/ssh/session.rs | 24 +++++++++-- 2 files changed, 97 insertions(+), 17 deletions(-) diff --git a/src/daemon/ssh/mod.rs b/src/daemon/ssh/mod.rs index c81c520d..be492045 100644 --- a/src/daemon/ssh/mod.rs +++ b/src/daemon/ssh/mod.rs @@ -241,7 +241,7 @@ impl SshManager { // the transport + SSH handshake only — never around interactive auth, // which the user may reasonably take a while to complete (the broker // enforces its own per-prompt timeout). - let conn = self + let (mut conn, reused) = self .open_connection(spec, broker) .await .map_err(|e| format!("{e}"))?; @@ -253,19 +253,43 @@ impl SshManager { broker.status(SshPhase::Connected); + // Open the shell channel on the (possibly shared) connection. This is also + // the first liveness probe of a *reused* connection: if its transport died + // silently — a parked forward/loopback accept loop holds an `Arc`, so the + // dead connection's `Drop` (and `mark_dead`) never ran — the first channel + // open errors. Self-heal: mark it dead, evict its registry slot, and + // reconnect fresh once. A fresh connection that fails here is a real error. + let channel = match conn.open_session_channel().await { + Ok(channel) => channel, + Err(e) if reused => { + log::info!( + "reused ssh connection to {}:{} was dead ({e}); reconnecting", + spec.host, + spec.port + ); + conn.mark_dead(); + self.evict_connection(conn.key()); + let (fresh, _) = self + .open_connection(spec, broker) + .await + .map_err(|e| format!("{e}"))?; + conn = fresh; + *conn_slot.lock().unwrap() = Arc::downgrade(&conn); + conn.open_session_channel() + .await + .map_err(|e| format!("open shell channel failed: {e}"))? + } + Err(e) => return Err(format!("open shell channel failed: {e}")), + }; + // Establish the profile's preconfigured forwards (FR-F2) now that the - // connection is authenticated. Failures are non-fatal — each surfaces as a - // `ForwardStatus::Error` on the forward row, never a killed session. + // connection is authenticated *and* confirmed live. Failures are non-fatal — + // each surfaces as a `ForwardStatus::Error` on the forward row, never a + // killed session. for rule in &spec.forwards { self.forwards.establish(pane_id, conn.clone(), rule).await; } - // Open the shell channel on the (possibly shared) connection. - let channel = conn - .open_session_channel() - .await - .map_err(|e| format!("open shell channel failed: {e}"))?; - let (pw, ph) = ( u32::from(size.cols).saturating_mul(u32::from(size.cell_w)), u32::from(size.rows).saturating_mul(u32::from(size.cell_h)), @@ -306,13 +330,24 @@ impl SshManager { Ok(()) } + /// Drop a connection key's registry slot so the next `open_connection` for it + /// establishes a fresh connection instead of upgrading a stale `Weak`. Called + /// by the self-healing reuse path when a reused connection turns out dead. + fn evict_connection(&self, key: &ConnectionKey) { + self.conns.lock().unwrap().remove(key); + } + /// Establish (or reuse) the connection for `spec`, recursing through the jump - /// chain. Boxed because it is `async`-recursive. + /// chain. Boxed because it is `async`-recursive. The returned `bool` is `true` + /// when an existing connection was reused (no fresh authentication) — the + /// caller uses it to self-heal: a reused connection whose transport silently + /// died errors on its first channel open, and only then is it worth evicting + /// and reconnecting. fn open_connection<'a>( &'a self, spec: &'a NativeSshSpec, broker: &'a Arc, - ) -> Pin>> + Send + 'a>> { + ) -> Pin, bool)>> + Send + 'a>> { Box::pin(async move { let key = ConnectionKey::from_spec(spec); let slot: ConnSlot = { @@ -325,14 +360,14 @@ impl SshManager { if let Some(conn) = guard.upgrade() { if conn.is_alive() { // Reuse: a new channel on the existing authenticated connection. - return Ok(conn); + return Ok((conn, true)); } } // Establish the jump connection first (recursively) so its // `direct-tcpip` channel can be this connection's transport. let jump = match &spec.jump { - Some(jump_spec) => Some(self.open_connection(jump_spec, broker).await?), + Some(jump_spec) => Some(self.open_connection(jump_spec, broker).await?.0), None => None, }; @@ -374,7 +409,7 @@ impl SshManager { let conn = SshConnection::new(handle, key, remote_forwards); *guard = Arc::downgrade(&conn); - Ok(conn) + Ok((conn, false)) }) } } @@ -447,6 +482,33 @@ mod tests { assert_eq!(a, ConnectionKey::from_spec(&base_spec())); } + #[test] + fn evict_connection_clears_the_registry_slot() { + // The self-heal path evicts a dead connection's key so the next + // `open_connection` establishes fresh instead of upgrading a stale `Weak`. + // Exercise just the registry map — no live server needed. + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .expect("build test runtime"); + let mgr = SshManager { + runtime, + conns: Mutex::new(HashMap::new()), + forwards: SshForwardRegistry::default(), + }; + let key = ConnectionKey::from_spec(&base_spec()); + mgr.conns + .lock() + .unwrap() + .insert(key.clone(), Arc::new(tokio::sync::Mutex::new(Weak::new()))); + assert!(mgr.conns.lock().unwrap().contains_key(&key)); + + mgr.evict_connection(&key); + assert!( + !mgr.conns.lock().unwrap().contains_key(&key), + "evicted key must be gone so the next open creates a new entry" + ); + } + #[test] fn connection_key_includes_jump_chain() { let mut with_jump = base_spec(); diff --git a/src/daemon/ssh/session.rs b/src/daemon/ssh/session.rs index 49401d14..703ac78a 100644 --- a/src/daemon/ssh/session.rs +++ b/src/daemon/ssh/session.rs @@ -286,12 +286,30 @@ impl SshConnection { &self.key } - /// Whether this connection is still usable for reuse. Set false once a channel - /// driver observes the session drop. + /// Whether this connection is still usable for reuse. + /// + /// Two signals: the `alive` flag (cleared by [`mark_dead`](Self::mark_dead) on + /// teardown or when a reuse attempt finds the transport dead) **and** the russh + /// handle's own liveness — when russh's session task ends (transport dropped), + /// its command sender closes, so `handle.is_closed()` flips to true. The flag + /// alone is unreliable: `mark_dead` only runs from `Drop`, but a parked + /// forward/loopback accept loop holds an `Arc`, so a dead + /// connection's `Drop` never runs and the flag stays true. Consulting + /// `is_closed()` (via a non-blocking `try_lock`; a contended lock means an open + /// is in flight, so assume alive) catches that case cheaply. The self-healing + /// reconnect in `SshManager::run_session` is the belt-and-suspenders backstop. pub fn is_alive(&self) -> bool { - self.alive.load(Ordering::SeqCst) + if !self.alive.load(Ordering::SeqCst) { + return false; + } + match self.handle.try_lock() { + Ok(handle) => !handle.is_closed(), + Err(_) => true, + } } + /// Mark this connection unusable for reuse (teardown, or a reuse attempt that + /// found the transport dead). Idempotent. pub(super) fn mark_dead(&self) { self.alive.store(false, Ordering::SeqCst); }