diff --git a/crates/maestro-registry/src/handler.rs b/crates/maestro-registry/src/handler.rs index 5c8fe00..ac08eb7 100644 --- a/crates/maestro-registry/src/handler.rs +++ b/crates/maestro-registry/src/handler.rs @@ -58,7 +58,9 @@ pub async fn handle_socket( } let (out_tx, mut out_rx) = mpsc::channel::(OUTBOUND_BUFFER); - registry + // Hold the fencing token for the life of this connection; only a deregister + // carrying it may tear down the entry we just installed (NAP-258). + let connection_token = registry .register( daemon_id.clone(), hostname, @@ -88,7 +90,7 @@ pub async fn handle_socket( .await .is_err() { - registry.deregister(&daemon_id).await; + registry.deregister(&daemon_id, connection_token).await; return Some(daemon_id); } } @@ -144,7 +146,7 @@ pub async fn handle_socket( } catalog_task.abort(); - registry.deregister(&daemon_id).await; + registry.deregister(&daemon_id, connection_token).await; Some(daemon_id) } diff --git a/crates/maestro-registry/src/state.rs b/crates/maestro-registry/src/state.rs index 3ce7fa0..cb53b06 100644 --- a/crates/maestro-registry/src/state.rs +++ b/crates/maestro-registry/src/state.rs @@ -102,6 +102,11 @@ struct DaemonEntry { perimeter_id: Option, /// Latest private address advertised on the declared perimeter. private_addr: Option, + /// Fencing token minted per WS connection in [`Registry::register`]. Only + /// the handler holding the matching token may [`Registry::deregister`] this + /// entry, so a stale handler (old socket timing out after the daemon + /// reconnected on a new one) can't wipe the live entry. (NAP-258) + connection_token: Uuid, } impl DaemonEntry { @@ -153,6 +158,10 @@ impl Registry { /// Register a freshly-authenticated daemon. Replaces any previous entry /// with the same `daemon_id` (and drops its sessions), matching the /// spec's "re-registration restores catalog" expectation. + /// + /// Returns a per-connection fencing token the caller must pass back to + /// [`Registry::deregister`] when the connection ends, so a stale handler + /// can't wipe a newer live entry (NAP-258). pub async fn register( &self, daemon_id: String, @@ -160,7 +169,8 @@ impl Registry { version: String, daemon_url: String, outbound: mpsc::Sender, - ) { + ) -> Uuid { + let connection_token = Uuid::new_v4(); let mut state = self.inner.lock().await; let removed_sessions = state.drop_daemon_sessions(&daemon_id); state.daemons.insert( @@ -172,6 +182,7 @@ impl Registry { outbound, perimeter_id: None, private_addr: None, + connection_token, }, ); drop(state); @@ -189,29 +200,34 @@ impl Registry { daemon_url, }, }); + connection_token } - /// Remove the daemon and any sessions it owned. Also fails any pending - /// RPCs that were sent to it. - pub async fn deregister(&self, daemon_id: &str) { - let removed_sessions; - let had_daemon; - { + /// Remove the daemon and any sessions it owned, but only if `token` matches + /// the entry's current [`DaemonEntry::connection_token`]. A stale handler + /// whose token no longer matches (the daemon reconnected on a new socket) + /// is a no-op, so it can't wipe the live entry or its sessions (NAP-258). + pub async fn deregister(&self, daemon_id: &str, token: Uuid) { + let removed_sessions = { let mut state = self.inner.lock().await; - had_daemon = state.daemons.remove(daemon_id).is_some(); - removed_sessions = state.drop_daemon_sessions(daemon_id); - } + match state.daemons.get(daemon_id) { + Some(entry) if entry.connection_token == token => {} + // Either superseded by a newer connection or already gone — + // leave the live entry (and its sessions) untouched. + _ => return, + } + state.daemons.remove(daemon_id); + state.drop_daemon_sessions(daemon_id) + }; for id in removed_sessions { self.broadcast(RegistryEvent::SessionRemoved { daemon_id: daemon_id.to_string(), id, }); } - if had_daemon { - self.broadcast(RegistryEvent::DaemonRemoved { - daemon_id: daemon_id.to_string(), - }); - } + self.broadcast(RegistryEvent::DaemonRemoved { + daemon_id: daemon_id.to_string(), + }); } // ----------------------------------------------------------------------- @@ -576,3 +592,86 @@ fn rpc_into_result(rpc: RpcResult) -> Result { RpcResult::Err { error } => Err(RpcError(error)), } } + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Utc; + use maestro_protocol::SessionStatus; + use std::path::PathBuf; + + fn sample_session(id: Uuid, daemon_id: &str) -> Session { + Session { + id, + daemon_id: daemon_id.to_string(), + cwd: PathBuf::from("/tmp"), + command: "bash".into(), + env: HashMap::new(), + metadata: HashMap::new(), + status: SessionStatus::Running, + created_at: Utc::now(), + } + } + + async fn register(reg: &Registry, daemon_id: &str) -> (Uuid, mpsc::Receiver) { + let (tx, rx) = mpsc::channel::(8); + let token = reg + .register( + daemon_id.to_string(), + "host".into(), + "0.1.0".into(), + format!("ws://{daemon_id}:7070"), + tx, + ) + .await; + (token, rx) + } + + /// NAP-258: a stale handler's deregister (old socket timing out after the + /// daemon reconnected on a new one) must NOT wipe the live entry or its + /// sessions. Only the current connection's token may tear it down. + #[tokio::test] + async fn stale_deregister_is_a_noop_after_reconnect() { + let reg = Registry::new(); + + // First connection. + let (token_old, _rx_old) = register(®, "X").await; + // Daemon reconnects on a new socket; register overwrites the entry and + // mints a fresh token. + let (token_new, _rx_new) = register(®, "X").await; + assert_ne!(token_old, token_new); + + // The live connection advertises a session. + let sid = Uuid::new_v4(); + reg.add_session("X", sample_session(sid, "X")).await; + + // The old handler finally times out and deregisters with its stale + // token — this must be a no-op. + reg.deregister("X", token_old).await; + assert!( + reg.daemons().await.iter().any(|d| d.daemon_id == "X"), + "stale deregister wiped the live daemon entry" + ); + assert_eq!( + reg.sessions_for("X").await.len(), + 1, + "stale deregister wiped the live daemon's sessions" + ); + + // The current handler's deregister (matching token) tears it down. + reg.deregister("X", token_new).await; + assert!( + !reg.daemons().await.iter().any(|d| d.daemon_id == "X"), + "matching-token deregister should remove the daemon" + ); + assert!(reg.sessions_for("X").await.is_empty()); + } + + /// Deregistering an unknown daemon (already gone) is a harmless no-op. + #[tokio::test] + async fn deregister_unknown_daemon_is_noop() { + let reg = Registry::new(); + reg.deregister("ghost", Uuid::new_v4()).await; + assert!(reg.daemons().await.is_empty()); + } +}