mirror of
https://github.com/l0ng-ai/tty7.git
synced 2026-09-22 08:02:24 +00:00
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:
+76
-14
@@ -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();
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user