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<SshConnection>, 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.
This commit is contained in:
l0ng-ai
2026-07-14 02:52:36 +08:00
parent 8f72b79ff2
commit f4c25725c5
2 changed files with 97 additions and 17 deletions
+76 -14
View File
@@ -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<PromptBroker>,
) -> Pin<Box<dyn Future<Output = anyhow::Result<Arc<SshConnection>>> + Send + 'a>> {
) -> Pin<Box<dyn Future<Output = anyhow::Result<(Arc<SshConnection>, 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();
+21 -3
View File
@@ -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<SshConnection>`, 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);
}