diff --git a/crates/maestro-registry/src/handler.rs b/crates/maestro-registry/src/handler.rs index da0920a..6277cf8 100644 --- a/crates/maestro-registry/src/handler.rs +++ b/crates/maestro-registry/src/handler.rs @@ -127,7 +127,7 @@ pub async fn handle_socket( } Ok(Some(Ok(Message::Close(_)))) => break, Ok(Some(Ok(Message::Text(text)))) => { - if !dispatch_text(®istry, &daemon_id, &text).await { + if !dispatch_text(®istry, &daemon_id, expected_token, &text).await { break; } } @@ -158,7 +158,12 @@ async fn read_hello( } } -async fn dispatch_text(registry: &Registry, daemon_id: &str, text: &str) -> bool { +async fn dispatch_text( + registry: &Registry, + daemon_id: &str, + shared_secret: &str, + text: &str, +) -> bool { let parsed: DaemonOutbound = match serde_json::from_str(text) { Ok(p) => p, Err(e) => { @@ -182,10 +187,16 @@ async fn dispatch_text(registry: &Registry, daemon_id: &str, text: &str) -> bool DaemonOutbound::SessionUpdated { session } => { registry.update_session(daemon_id, session).await; } - DaemonOutbound::Heartbeat { .. } => { - // Per-read timeout above is the heartbeat watchdog; no extra work. - // perimeter_id / private_addr advertised here are persisted by - // NAP-243 for the T2 transport. + DaemonOutbound::Heartbeat { + perimeter_id, + private_addr, + } => { + // Per-read timeout above is the heartbeat watchdog; capture the + // perimeter advertisement so the broker can hand it back as a T2 + // candidate (NAP-247). + registry + .update_perimeter(daemon_id, perimeter_id, private_addr) + .await; } DaemonOutbound::CreateSessionResult { reply_to, result } => { registry.complete_create(reply_to, result).await; @@ -193,11 +204,18 @@ async fn dispatch_text(registry: &Registry, daemon_id: &str, text: &str) -> bool DaemonOutbound::KillSessionResult { reply_to, result } => { registry.complete_kill(reply_to, result).await; } - DaemonOutbound::ResolveAttach { reply_to, .. } - | DaemonOutbound::RelayAck { reply_to, .. } => { - // Handled by NAP-243 (broker mode); type already lands here so old - // daemons sending it don't kill the connection on parse failure. - tracing::debug!(%daemon_id, ?reply_to, "broker message received but unhandled (NAP-243 pending)"); + DaemonOutbound::ResolveAttach { + reply_to, + session_id, + } => { + registry + .resolve_attach(daemon_id, reply_to, session_id, shared_secret) + .await; + } + DaemonOutbound::RelayAck { reply_to, .. } => { + // T3 (relay) is gated on NAP-238. Land the variant so newer daemons + // don't crash the connection; the registry doesn't act on it yet. + tracing::debug!(%daemon_id, ?reply_to, "RelayAck received but T3 not implemented (NAP-238 pending)"); } } true diff --git a/crates/maestro-registry/src/jwt.rs b/crates/maestro-registry/src/jwt.rs index a92ceed..d96c8af 100644 --- a/crates/maestro-registry/src/jwt.rs +++ b/crates/maestro-registry/src/jwt.rs @@ -6,11 +6,17 @@ use chrono::{Duration, Utc}; use jsonwebtoken::{encode, Algorithm, EncodingKey, Header}; -use maestro_protocol::JwtClaims; +use maestro_protocol::{JwtClaims, PairAttachClaims}; use uuid::Uuid; /// Spec §Auth: ~5 minutes. pub const ATTACH_JWT_TTL: Duration = Duration::minutes(5); +/// Pair-scoped attach JWT TTL — short-lived per the connection-orchestration +/// spec (60 s default; ample window for the home daemon to dial T2). +pub const PAIR_ATTACH_JWT_TTL: Duration = Duration::seconds(60); +/// `iss` claim baked into every registry-issued pair JWT. Plain string for +/// now — once the registry gets a per-instance identity, plumb that in here. +pub const REGISTRY_ISSUER: &str = "maestro-registry"; /// Mint a short-lived attach JWT for `(daemon_id, session_id)`. pub fn issue_attach_jwt( @@ -30,6 +36,29 @@ pub fn issue_attach_jwt( encode(&header, &claims, &key) } +/// Mint a pair-scoped attach JWT used by the home-daemon broker (NAP-247). +/// `sub` is the home daemon-id; `aud` is the target daemon-id. Target +/// validates signature, `aud == self.daemon_id`, `exp > now`, and +/// `sub in allow_inbound` before accepting the broker-attach. +pub fn issue_pair_attach_jwt( + secret: &str, + home_daemon_id: &str, + target_daemon_id: &str, + session_id: Uuid, +) -> Result { + let now = Utc::now(); + let claims = PairAttachClaims { + iss: REGISTRY_ISSUER.into(), + sub: home_daemon_id.to_string(), + aud: target_daemon_id.to_string(), + sid: session_id, + exp: (now + PAIR_ATTACH_JWT_TTL).timestamp(), + }; + let header = Header::new(Algorithm::HS256); + let key = EncodingKey::from_secret(secret.as_bytes()); + encode(&header, &claims, &key) +} + #[cfg(test)] mod tests { use super::*; @@ -51,4 +80,21 @@ mod tests { assert_eq!(decoded.claims.session_id, session); assert!(decoded.claims.exp > Utc::now().timestamp()); } + + #[test] + fn pair_jwt_carries_iss_sub_aud_sid() { + let secret = "shared-secret"; + let session = Uuid::new_v4(); + let token = issue_pair_attach_jwt(secret, "laptop", "konishi", session).unwrap(); + let mut v = Validation::new(Algorithm::HS256); + v.set_audience(&["konishi"]); + let decoded = + decode::(&token, &DecodingKey::from_secret(secret.as_bytes()), &v) + .unwrap(); + assert_eq!(decoded.claims.iss, REGISTRY_ISSUER); + assert_eq!(decoded.claims.sub, "laptop"); + assert_eq!(decoded.claims.aud, "konishi"); + assert_eq!(decoded.claims.sid, session); + assert!(decoded.claims.exp > Utc::now().timestamp()); + } } diff --git a/crates/maestro-registry/src/state.rs b/crates/maestro-registry/src/state.rs index c1add52..3ce7fa0 100644 --- a/crates/maestro-registry/src/state.rs +++ b/crates/maestro-registry/src/state.rs @@ -14,12 +14,14 @@ use std::sync::Arc; use std::time::Duration; use maestro_protocol::{ - CreateSessionRequest, DaemonInbound, DaemonInfo, ErrorInfo, RegistryEvent, ReplyTo, RpcResult, - Session, + AttachCandidates, CreateSessionRequest, DaemonInbound, DaemonInfo, ErrorInfo, + PerimeterCandidate, RegistryEvent, ReplyTo, RpcResult, Session, }; use tokio::sync::{broadcast, mpsc, oneshot, Mutex}; use uuid::Uuid; +use crate::jwt::issue_pair_attach_jwt; + /// Capacity of the registry-wide broadcast channel that feeds /// `WS /api/events` subscribers. Subscribers that fall this far behind get a /// `Lagged` event and are expected to resync via `GET /api/sessions`. @@ -95,6 +97,11 @@ struct DaemonEntry { version: String, daemon_url: String, outbound: mpsc::Sender, + /// Latest perimeter advertised by this daemon in its heartbeats. + /// `None` until the daemon advertises one (NAP-247 T2 dial target). + perimeter_id: Option, + /// Latest private address advertised on the declared perimeter. + private_addr: Option, } impl DaemonEntry { @@ -163,6 +170,8 @@ impl Registry { version: version.clone(), daemon_url: daemon_url.clone(), outbound, + perimeter_id: None, + private_addr: None, }, ); drop(state); @@ -281,6 +290,118 @@ impl Registry { state.daemons.get(daemon_id).map(|e| e.daemon_url.clone()) } + /// Update a daemon's advertised perimeter from its heartbeat. Called by + /// the WS handler on every `Heartbeat` frame (NAP-247). + pub async fn update_perimeter( + &self, + daemon_id: &str, + perimeter_id: Option, + private_addr: Option, + ) { + let mut state = self.inner.lock().await; + if let Some(entry) = state.daemons.get_mut(daemon_id) { + entry.perimeter_id = perimeter_id; + entry.private_addr = private_addr; + } + } + + /// `(perimeter_id, private_addr)` last advertised by this daemon, when + /// both are set. The broker uses this to assemble `AttachCandidates` for + /// the T2 transport (NAP-247). + pub async fn perimeter_for(&self, daemon_id: &str) -> Option<(String, String)> { + let state = self.inner.lock().await; + state.daemons.get(daemon_id).and_then(|e| { + match (e.perimeter_id.clone(), e.private_addr.clone()) { + (Some(pid), Some(addr)) => Some((pid, addr)), + _ => None, + } + }) + } + + /// Handle a `ResolveAttach` from a home daemon (NAP-247). Look up which + /// daemon owns `session_id`, mint a pair-scoped JWT for the home → target + /// pair, and reply with `AttachCandidates`. Errors when the session is + /// unknown or its owning daemon has gone offline. + pub async fn resolve_attach( + &self, + home_daemon_id: &str, + reply_to: ReplyTo, + session_id: Uuid, + shared_secret: &str, + ) { + let result = self + .build_attach_candidates(home_daemon_id, session_id, shared_secret) + .await; + let frame = DaemonInbound::ResolveAttachResult { reply_to, result }; + let outbound = { + let state = self.inner.lock().await; + state + .daemons + .get(home_daemon_id) + .map(|e| e.outbound.clone()) + }; + // If the home daemon has gone away mid-flight, drop the response. + if let Some(outbound) = outbound { + let _ = outbound.send(frame).await; + } + } + + async fn build_attach_candidates( + &self, + home_daemon_id: &str, + session_id: Uuid, + shared_secret: &str, + ) -> RpcResult { + let Some(target) = self.daemon_for_session(session_id).await else { + return RpcResult::Err { + error: ErrorInfo { + code: "unknown_session".into(), + message: format!("no daemon owns session {session_id}"), + }, + }; + }; + // Self-attach: home daemon already short-circuits this in its own + // broker. But if it asks anyway, refuse — `aud != self` won't pass + // anyway and we'd just be wasting a JWT. + if target == home_daemon_id { + return RpcResult::Err { + error: ErrorInfo { + code: "self_target".into(), + message: "home daemon owns this session; short-circuit locally".into(), + }, + }; + } + let jwt = match issue_pair_attach_jwt(shared_secret, home_daemon_id, &target, session_id) { + Ok(t) => t, + Err(e) => { + return RpcResult::Err { + error: ErrorInfo { + code: "jwt_failed".into(), + message: e.to_string(), + }, + }; + } + }; + // T1 (direct WSS) stays None — NAP-235 was canceled. + // T2: hand back the target's perimeter so the home daemon can match. + // T3 (relay) stays None — NAP-238 spike pending. + let perimeter = self + .perimeter_for(&target) + .await + .map(|(pid, addr)| PerimeterCandidate { + perimeter_id: pid, + private_addr: addr, + }); + RpcResult::Ok { + value: AttachCandidates { + direct_wss: None, + perimeter, + relay: None, + jwt, + }, + } + } + /// Resolve a CreateSession RPC reply. Returns true if a waiter was found. pub async fn complete_create(&self, reply_to: ReplyTo, result: RpcResult) -> bool { let mut state = self.inner.lock().await;