From b4b96dfb3b476dab4fb2721526c1857d5f49b324 Mon Sep 17 00:00:00 2001 From: kjgbot Date: Thu, 6 Aug 2026 09:55:58 +0200 Subject: [PATCH 1/2] fix(broker): reject spawn-admission name collisions instead of handing over the incumbent's token MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit register_agent_with_workspace_key's strict-name branch reclaimed any name collision via register_or_get_agent, returning the existing agent's id, name, AND bearer token to whoever asked for that name — identity handover, not admission control. The non-strict branch diverged instead of agreeing: a silent one-shot `-{uuid8}` suffix retry, unreachable in production since every live agent runs RELAY_STRICT_AGENT_NAME=1. Cloud keys Relayfile tokens by agent name, so a collision shared credential scope; a silent suffix would have produced a duplicate-work sibling agent under a near-identical name, the AR-448 class this closes. Both paths now route through one admit_agent_registration decision that fails closed on a bare name collision. Reclaim (rotate and return the existing agent's token, keeping crash-recovery resume working) is only granted when the registering caller proves the same work-unit identity via RELAY_AGENT_IDENTITY_KEY matching the identity_key stamped on the existing agent's metadata at its own creation — never by the name string alone. Red-first: the new strict_name_conflict_without_identity_proof_is_rejected_not_handed_incumbent_token test fails against the pre-fix code (asserts Err; pre-fix returns Ok with the incumbent's rotated token) and passes after. Companion tests cover identity-matched reclaim and non-strict/strict agreement; the two prior tests that encoded the vulnerable behavior as correct are removed. Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.md | 4 + crates/broker/src/relaycast/auth.rs | 320 ++++++++++++++++++++-------- 2 files changed, 234 insertions(+), 90 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a9ef68e9e..2ed8ba865 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `agent-relay cloud enroll` records the enrolled node on the project workspace pin, so `node up` in that repo serves the node it just enrolled. A pin that already names a different node is reported and left untouched rather than repointed. - `agent-relay workspace switch|join` keeps the project's enrolled fleet node id instead of dropping it, which previously produced the pin state that made the next `node up` ignore the enrollment store. +### Security + +- Agent registration no longer hands over an existing agent's id, name, and bearer token to whoever registers with the same name. A name collision is now rejected unless the request proves it's the same work unit via `RELAY_AGENT_IDENTITY_KEY` matching the identity stamped on the existing agent at its creation; strict- and non-strict-name registration now share this same fail-closed admission decision. + ## [11.4.1] - 2026-08-03 ### Changed diff --git a/crates/broker/src/relaycast/auth.rs b/crates/broker/src/relaycast/auth.rs index 15743b4fc..328191c69 100644 --- a/crates/broker/src/relaycast/auth.rs +++ b/crates/broker/src/relaycast/auth.rs @@ -673,81 +673,26 @@ impl AuthClient { } } + /// `strict_name` is retained purely for API/logging compatibility with + /// callers (`crates/broker/src/runtime/session.rs` chooses it based on + /// `RELAY_STRICT_AGENT_NAME`) and no longer selects a different collision + /// strategy: both modes route through `admit_agent_registration`. Their + /// prior divergence — strict silently handed over the incumbent's token, + /// non-strict silently minted a `-suffix` sibling — was itself the + /// spawn-admission defect (see `admit_agent_registration`). async fn register_agent_with_workspace_key( &self, workspace_key: &str, requested_name: Option<&str>, - strict_name: bool, + _strict_name: bool, agent_type: Option<&str>, ) -> Result<(String, String, String, Option)> { let relay = build_relay_client(workspace_key, self.base_url.as_deref())?; - let mut attempted_retry = false; - let mut name = requested_name + let name = requested_name .map(ToOwned::to_owned) .unwrap_or_else(|| format!("agent-{}", Uuid::new_v4().simple())); - if strict_name { - let request = CreateAgentRequest { - name, - agent_type: Some(agent_type.unwrap_or("agent").to_string()), - persona: None, - metadata: None, - }; - let result = relay - .register_or_get_agent(request) - .await - .map_err(relay_error_to_anyhow)?; - return Ok((result.id, result.name, result.token, result.workspace_id)); - } - - loop { - let request = CreateAgentRequest { - name: name.clone(), - agent_type: Some(agent_type.unwrap_or("agent").to_string()), - persona: None, - metadata: None, - }; - - match relay.register_agent(request).await { - Ok(result) => { - return Ok((result.id, result.name, result.token, result.workspace_id)); - } - Err(RelayError::Api { code, status, .. }) - if is_conflict_code(&code) || status == 409 => - { - if !attempted_retry { - attempted_retry = true; - let suffix = Uuid::new_v4().simple().to_string(); - name = format!("{}-{}", name, &suffix[..8]); - continue; - } - return Err(relay_error_to_anyhow(RelayError::Api { - code: "agent_already_exists".to_string(), - message: format!("agent name '{}' already exists after retry", name), - status: 409, - })); - } - Err(RelayError::Api { - code, - status, - message, - }) if is_agent_token_invalid_code(&code) - || (status == 401 && message.trim() == AGENT_TOKEN_INVALID_MESSAGE) => - { - // Surface the typed code even when only the legacy - // status+message pair is present, so downstream callers - // can react with `is_agent_token_invalid_anyhow`. - return Err(relay_error_to_anyhow(RelayError::Api { - code: AGENT_TOKEN_INVALID_CODE.to_string(), - status, - message, - })); - } - Err(error) => { - return Err(relay_error_to_anyhow(error)); - } - } - } + admit_agent_registration(&relay, &name, agent_type).await } pub async fn workspace_key_is_live(&self, workspace_key: &str) -> Result { @@ -899,6 +844,123 @@ fn is_conflict_code(code: &str) -> bool { ) } +/// Metadata key an agent's identity proof is stamped under at registration, +/// so a later collision can check it back. See `admit_agent_registration`. +const IDENTITY_METADATA_KEY: &str = "identity_key"; + +/// Caller-supplied proof of work-unit identity for spawn-admission reclaim. +/// +/// A crashed (or resumed) work unit that needs to re-register under its +/// prior name sets this to a value stable across that work unit's restarts. +/// It gets stamped onto the agent's metadata at creation; a later collision +/// under the same name is only treated as a reclaim of that SAME work unit +/// if the value presented then matches what's stored — never by the name +/// string alone. Absent, registration cannot reclaim on collision. +fn agent_identity_key() -> Option { + std::env::var("RELAY_AGENT_IDENTITY_KEY") + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +/// Single spawn-admission decision shared by every registration path +/// (formerly split between a strict branch that always reclaimed a +/// name-collision via `register_or_get_agent` — handing the caller the +/// incumbent's id, name, AND bearer token — and a non-strict branch that +/// silently minted a `-{uuid8}` sibling name once before failing). Both +/// were the same defect: a spawn-admission gate that doesn't verify who is +/// asking. This repo's doctrine is a dispatch gate fails closed, so a name +/// collision is REJECTED by default. A silent suffix would have produced a +/// second agent doing duplicate work under a near-identical name — exactly +/// the AR-448 duplicate-agent class this gate exists to stop — so it is not +/// an acceptable alternative to rejection either. +/// +/// Reclaim (return the existing agent's identity with a freshly rotated +/// token, so a crashed broker's resume path keeps working) is permitted +/// ONLY when the registration request proves it is the same work unit: its +/// `RELAY_AGENT_IDENTITY_KEY` (see `agent_identity_key`) must match the +/// identity key stamped on the existing agent's metadata at its own +/// creation. Absent or mismatched identity on a collision is rejected. +async fn admit_agent_registration( + relay: &RelayCast, + name: &str, + agent_type: Option<&str>, +) -> Result<(String, String, String, Option)> { + let identity_key = agent_identity_key(); + let metadata = identity_key.as_ref().map(|key| { + let mut map = serde_json::Map::new(); + map.insert( + IDENTITY_METADATA_KEY.to_string(), + Value::String(key.clone()), + ); + map + }); + + let request = CreateAgentRequest { + name: name.to_string(), + agent_type: Some(agent_type.unwrap_or("agent").to_string()), + persona: None, + metadata, + }; + + match relay.register_agent(request).await { + Ok(result) => Ok((result.id, result.name, result.token, result.workspace_id)), + Err(RelayError::Api { code, status, .. }) if is_conflict_code(&code) || status == 409 => { + let existing = relay.get_agent(name).await.map_err(relay_error_to_anyhow)?; + let existing_identity = existing + .metadata + .get(IDENTITY_METADATA_KEY) + .and_then(Value::as_str); + + let reclaims_same_work_unit = matches!( + (identity_key.as_deref(), existing_identity), + (Some(ours), Some(theirs)) if ours == theirs + ); + + if !reclaims_same_work_unit { + return Err(relay_error_to_anyhow(RelayError::Api { + code: "agent_identity_mismatch".to_string(), + status: 409, + message: format!( + "agent name '{name}' is already registered and this registration did \ + not prove ownership of that identity; refusing to hand over its \ + credentials (set RELAY_AGENT_IDENTITY_KEY to the original work unit's \ + identity to reclaim it after a crash)" + ), + })); + } + + let token_response = relay + .rotate_agent_token(&existing.name) + .await + .map_err(relay_error_to_anyhow)?; + Ok(( + existing.id, + existing.name, + token_response.token, + existing.workspace_id, + )) + } + Err(RelayError::Api { + code, + status, + message, + }) if is_agent_token_invalid_code(&code) + || (status == 401 && message.trim() == AGENT_TOKEN_INVALID_MESSAGE) => + { + // Surface the typed code even when only the legacy status+message + // pair is present, so downstream callers can react with + // `is_agent_token_invalid_anyhow`. + Err(relay_error_to_anyhow(RelayError::Api { + code: AGENT_TOKEN_INVALID_CODE.to_string(), + status, + message, + })) + } + Err(error) => Err(relay_error_to_anyhow(error)), + } +} + fn is_workspace_name_conflict(error: &RelayError) -> bool { match error { RelayError::Api { @@ -1019,6 +1081,7 @@ mod tests { std::env::remove_var("RELAY_API_KEY"); std::env::remove_var("RELAY_WORKSPACES_JSON"); std::env::remove_var("RELAY_DEFAULT_WORKSPACE"); + std::env::remove_var("RELAY_AGENT_IDENTITY_KEY"); } guard } @@ -1234,11 +1297,12 @@ mod tests { } #[tokio::test] - async fn strict_name_conflict_reclaims_via_sdk_register_or_get_agent() { - // Regression test for issue #797: when a broker is restarted (or a - // second broker joins via shared workspace key) with a name that's - // already registered, registration must reclaim the existing agent - // through the relaycast SDK instead of failing the broker startup. + async fn strict_name_conflict_without_identity_proof_is_rejected_not_handed_incumbent_token() { + // Spawn-admission gate regression test: a bare name collision must + // never hand the caller the incumbent agent's token. Reclaim is only + // permitted when the caller proves the same work-unit identity (see + // `admit_agent_registration`); presenting the same name string alone + // must be rejected, not silently reclaimed. let _env_guard = clear_relay_env(); let server = MockServer::start(); unsafe { @@ -1262,8 +1326,6 @@ mod tests { .header("authorization", "Bearer rk_live_shared"); then.status(200) .header("content-type", "application/json") - // Mirrors the live cloud's GET /v1/agents/{name} payload that - // relaycast 1.0.1 accepts while reclaiming the agent. .body(r#"{"ok":true,"data":{"id":"a_existing","name":"lead","type":"agent","status":"offline","persona":null,"metadata":{},"last_seen":"2025-01-01T00:00:00Z","channels":[]}}"#); }); let rotate = server.mock(|when, then| { @@ -1275,27 +1337,97 @@ mod tests { .body(r#"{"ok":true,"data":{"name":"lead","token":"at_live_rotated"}}"#); }); + let client = AuthClient::new(Some(server.base_url())); + let result = client + .startup_session_with_options(Some("lead"), true, None) + .await; + + assert!( + result.is_err(), + "a name collision with no proof of matching identity must be rejected, not reclaimed" + ); + let message = result.unwrap_err().to_string(); + assert!( + !message.contains("at_live_rotated"), + "rejection must not leak the incumbent's rotated token: {message}" + ); + conflict.assert_hits(1); + get_existing.assert_hits(1); + rotate.assert_hits(0); + + unsafe { + std::env::remove_var("RELAY_API_KEY"); + } + } + + #[tokio::test] + async fn strict_name_conflict_with_matching_identity_reclaims_existing_agent() { + // Crash-recovery resume: the SAME work unit re-registers under the + // same name and proves it via RELAY_AGENT_IDENTITY_KEY matching what + // was stamped on the agent at its original creation. This must still + // reclaim the agent (and rotate its token) rather than being rejected. + let _env_guard = clear_relay_env(); + let server = MockServer::start(); + unsafe { + std::env::set_var("RELAY_API_KEY", "rk_live_shared"); + std::env::set_var("RELAY_AGENT_IDENTITY_KEY", "work-unit-42"); + } + let conflict = server.mock(|when, then| { + when.method(POST) + .path("/v1/agents") + .header("authorization", "Bearer rk_live_shared") + .json_body(json!({ + "name": "lead", + "type": "agent", + "metadata": { "identity_key": "work-unit-42" } + })); + then.status(409) + .header("content-type", "application/json") + .body(r#"{"ok":false,"error":{"code":"agent_already_exists","message":"name_taken"}}"#); + }); + let get_existing = server.mock(|when, then| { + when.method(GET) + .path("/v1/agents/lead") + .header("authorization", "Bearer rk_live_shared"); + then.status(200) + .header("content-type", "application/json") + .body(r#"{"ok":true,"data":{"id":"a_existing","name":"lead","type":"agent","status":"offline","persona":null,"metadata":{"identity_key":"work-unit-42"},"last_seen":"2025-01-01T00:00:00Z","channels":[]}}"#); + }); + let rotate = server.mock(|when, then| { + when.method(POST) + .path("/v1/agents/lead/rotate-token") + .header("authorization", "Bearer rk_live_shared"); + then.status(200) + .header("content-type", "application/json") + .body(r#"{"ok":true,"data":{"name":"lead","token":"at_live_rotated"}}"#); + }); + let client = AuthClient::new(Some(server.base_url())); let session = client .startup_session_with_options(Some("lead"), true, None) .await - .expect("strict-name conflict should reclaim via relaycast SDK"); + .expect("matching identity proof should reclaim the existing agent"); assert_eq!(session.token, "at_live_rotated"); assert_eq!(session.credentials.agent_id, "a_existing"); - assert_eq!(session.credentials.api_key, "rk_live_shared"); - assert_eq!(session.credentials.agent_name.as_deref(), Some("lead")); conflict.assert_hits(1); get_existing.assert_hits(1); rotate.assert_hits(1); unsafe { std::env::remove_var("RELAY_API_KEY"); + std::env::remove_var("RELAY_AGENT_IDENTITY_KEY"); } } #[tokio::test] - async fn default_name_conflict_retries_with_suffix_once() { + async fn non_strict_name_conflict_without_identity_proof_is_rejected() { + // Strict and non-strict registration must agree: the divergence + // between an always-reclaiming strict path and a silently-suffixing + // non-strict path WAS the defect (a silent suffix mints a duplicate + // agent under a near-identical name, the same AR-448 class a bare + // handover produces). Both now route through the same fail-closed + // admission decision. let _env_guard = clear_relay_env(); let server = MockServer::start(); let workspace = server.mock(|when, then| { @@ -1304,7 +1436,7 @@ mod tests { .header("content-type", "application/json") .body(r#"{"ok":true,"data":{"workspace_id":"ws_new","api_key":"rk_live_cached","created_at":"2025-01-01T00:00:00Z"}}"#); }); - let first_conflict = server.mock(|when, then| { + let conflict = server.mock(|when, then| { when.method(POST) .path("/v1/agents") .header("authorization", "Bearer rk_live_cached") @@ -1316,29 +1448,37 @@ mod tests { .header("content-type", "application/json") .body(r#"{"ok":false,"error":{"code":"agent_already_exists","message":"name_taken"}}"#); }); - let second_success = server.mock(|when, then| { - when.method(POST) - .path("/v1/agents") - .header("authorization", "Bearer rk_live_cached") - .body_contains("\"name\":\"lead-"); + let get_existing = server.mock(|when, then| { + when.method(GET) + .path("/v1/agents/lead") + .header("authorization", "Bearer rk_live_cached"); then.status(200) .header("content-type", "application/json") - .body(r#"{"ok":true,"data":{"id":"a10","name":"lead-suffixed","token":"at_live_10","status":"online","created_at":"2025-01-01T00:00:00Z"}}"#); + .body(r#"{"ok":true,"data":{"id":"a_existing","name":"lead","type":"agent","status":"offline","persona":null,"metadata":{},"last_seen":"2025-01-01T00:00:00Z","channels":[]}}"#); }); let client = AuthClient::new(Some(server.base_url())); - let session = client.startup_session(Some("lead")).await.unwrap(); + let result = client.startup_session(Some("lead")).await; - assert_eq!(session.token, "at_live_10"); - assert_eq!( - session.credentials.agent_name.as_deref(), - Some("lead-suffixed") + assert!( + result.is_err(), + "non-strict registration must also reject an unproven name collision, not mint a silent -suffix sibling" ); workspace.assert_hits(1); - first_conflict.assert_hits(1); - second_success.assert_hits(1); + conflict.assert_hits(1); + get_existing.assert_hits(1); } + // `strict_name_conflict_reclaims_via_sdk_register_or_get_agent` (issue + // #797) and `default_name_conflict_retries_with_suffix_once` used to live + // here. Both encoded the spawn-admission defect as intended behavior — + // an unconditional reclaim-on-collision, and a silent `-{uuid8}` + // suffix-on-collision, respectively. They're superseded by + // `strict_name_conflict_without_identity_proof_is_rejected_not_handed_incumbent_token`, + // `strict_name_conflict_with_matching_identity_reclaims_existing_agent`, + // and `non_strict_name_conflict_without_identity_proof_is_rejected` above, + // which assert the fixed fail-closed contract instead. + #[tokio::test] async fn workspace_name_conflict_retries_with_fresh_suffix() { let _env_guard = clear_relay_env(); From 5c2ad8ee36e6c8a480ec2a5c459d6d723989192e Mon Sep 17 00:00:00 2001 From: kjgbot Date: Thu, 6 Aug 2026 11:01:16 +0200 Subject: [PATCH 2/2] fix(broker): reclaim a node's own registration across restart, hash the identity proof MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - The spawn-admission fail-closed gate had no path for a fleet node's own legitimate restart: nothing set RELAY_AGENT_IDENTITY_KEY for a node's own startup registration, so a node killed and restarted before its stale registration was reaped collided on its own name and was rejected outright, never coming back online (the "Two-node fleet matrix" restart-reconcile flake). connect_relay now derives a stable identity from the broker's own persisted state directory (falling back to an explicit RELAY_AGENT_IDENTITY_KEY override) and passes it through a new startup_session_set_with_identity, so a restart of the same node reclaims its name while a different node colliding on that name is still rejected. - The identity proof stamped on an agent's metadata was stored raw, but metadata is readable by any caller holding the same workspace key — so a co-tenant could read another work unit's identity key and replay it to reclaim its credentials, defeating the admission gate. It's now stored and compared as a one-way SHA-256 hash. Both new tests are proven against the pre-fix code: without the identity threading, node_restart_reclaims_its_own_prior_registration_via_stable_identity fails to compile (no such method); reverting the hash-storage change makes strict_name_conflict_with_matching_identity_reclaims_existing_agent fail on the mock's exact-body match. cargo test -p agent-relay-broker --lib: 855 passed, 0 failed, 4 ignored. cargo fmt -p agent-relay-broker -- --check: clean. cargo clippy -p agent-relay-broker --lib --tests -- -D warnings: no new warnings (3 pre-existing errors in unrelated files, unchanged by this PR, per the PR's own test plan). Note on the two remaining cubic P1 findings (auth.rs:917, auth.rs:934): - auth.rs:917 (raw identity key readable via shared workspace key): fixed above via hashing. - auth.rs:934 (rotate_agent_token(&existing.name) is name-keyed, not ID-keyed, so a delete+re-register race between get_agent and rotate could hand a token for an unverified agent): valid, but not fixable from this repo — the relaycast SDK's rotate-token/get-agent endpoints are fundamentally name-keyed (TokenRotateResponse has no id field), so closing this race needs a new ID-scoped or compare-and-swap primitive in the relaycast crate/server, a separate repository. Left as-is rather than ship a cosmetic mitigation that doesn't actually close the window. Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.md | 2 + crates/broker/src/relaycast/auth.rs | 268 +++++++++++++++++++++++++-- crates/broker/src/relaycast/mod.rs | 2 +- crates/broker/src/runtime/mod.rs | 7 +- crates/broker/src/runtime/session.rs | 12 +- 5 files changed, 268 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2ed8ba865..67be6aeef 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,10 +17,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `agent-relay node up` warns instead of silently ignoring stored Cloud fleet enrollments when the project workspace pin has no enrolled node id. That combination started the broker in the pinned workspace while the node never heartbeat, leaving the Cloud dashboard and `agent-relay fleet nodes` showing different rosters with no error from either. - `agent-relay cloud enroll` records the enrolled node on the project workspace pin, so `node up` in that repo serves the node it just enrolled. A pin that already names a different node is reported and left untouched rather than repointed. - `agent-relay workspace switch|join` keeps the project's enrolled fleet node id instead of dropping it, which previously produced the pin state that made the next `node up` ignore the enrollment store. +- A fleet node killed and restarted before its stale registration is reaped no longer fails to come back online: the broker now proves its own restart identity from its persisted state directory, so it can reclaim its prior name without an operator setting `RELAY_AGENT_IDENTITY_KEY` by hand. ### Security - Agent registration no longer hands over an existing agent's id, name, and bearer token to whoever registers with the same name. A name collision is now rejected unless the request proves it's the same work unit via `RELAY_AGENT_IDENTITY_KEY` matching the identity stamped on the existing agent at its creation; strict- and non-strict-name registration now share this same fail-closed admission decision. +- That identity proof is stored as a one-way hash rather than the raw value, so a workspace member who can read agent metadata can no longer replay another work unit's identity key to reclaim its credentials. ## [11.4.1] - 2026-08-03 diff --git a/crates/broker/src/relaycast/auth.rs b/crates/broker/src/relaycast/auth.rs index 328191c69..0d88f0f4e 100644 --- a/crates/broker/src/relaycast/auth.rs +++ b/crates/broker/src/relaycast/auth.rs @@ -307,11 +307,38 @@ impl AuthClient { .context("no default workspace session was available") } + /// See [`Self::startup_session_set_with_identity`]. Uses + /// `RELAY_AGENT_IDENTITY_KEY` (see `agent_identity_key`) as the identity + /// proof, preserving prior behavior for every existing caller. pub async fn startup_session_set_with_options( &self, requested_name: Option<&str>, strict_name: bool, agent_type: Option<&str>, + ) -> Result { + self.startup_session_set_with_identity( + requested_name, + strict_name, + agent_type, + agent_identity_key().as_deref(), + ) + .await + } + + /// Same as [`Self::startup_session_set_with_options`], but with an + /// explicit identity proof rather than reading `RELAY_AGENT_IDENTITY_KEY` + /// from the environment. The broker's own startup registration + /// (`connect_relay`) uses this to pass a value stable across its own + /// restarts (see `stable_node_identity_key`) without mutating process + /// env — an env mutation here would leak into every worker this broker + /// later spawns, since child processes inherit the full parent + /// environment. + pub async fn startup_session_set_with_identity( + &self, + requested_name: Option<&str>, + strict_name: bool, + agent_type: Option<&str>, + identity_key: Option<&str>, ) -> Result { if let Some((sources, default_hint)) = self.load_workspace_sources_from_env()? { let preferred_name = requested_name; @@ -328,6 +355,7 @@ impl AuthClient { preferred_name, strict_name, agent_type, + identity_key, ) .await { @@ -374,8 +402,13 @@ impl AuthClient { }); } - self.startup_single_session_set_from_sources(requested_name, strict_name, agent_type) - .await + self.startup_single_session_set_from_sources( + requested_name, + strict_name, + agent_type, + identity_key, + ) + .await } /// Rotate the token for an existing agent without re-registering. @@ -399,7 +432,13 @@ impl AuthClient { "agent not found during token rotation, falling back to re-registration" ); let registration = self - .register_agent_with_workspace_key(&api_key, Some(agent_name), false, None) + .register_agent_with_workspace_key( + &api_key, + Some(agent_name), + false, + None, + agent_identity_key().as_deref(), + ) .await .context("failed to re-register after rotate-token 404")?; let mut session = @@ -447,6 +486,7 @@ impl AuthClient { requested_name: Option<&str>, strict_name: bool, agent_type: Option<&str>, + identity_key: Option<&str>, ) -> Result { let env_workspace_key = env_workspace_key()?; @@ -488,6 +528,7 @@ impl AuthClient { preferred_name, strict_name, agent_type, + identity_key, ) .await { @@ -545,6 +586,7 @@ impl AuthClient { preferred_name, strict_name, agent_type, + identity_key, ) .await { @@ -686,13 +728,14 @@ impl AuthClient { requested_name: Option<&str>, _strict_name: bool, agent_type: Option<&str>, + identity_key: Option<&str>, ) -> Result<(String, String, String, Option)> { let relay = build_relay_client(workspace_key, self.base_url.as_deref())?; let name = requested_name .map(ToOwned::to_owned) .unwrap_or_else(|| format!("agent-{}", Uuid::new_v4().simple())); - admit_agent_registration(&relay, &name, agent_type).await + admit_agent_registration(&relay, &name, agent_type, identity_key).await } pub async fn workspace_key_is_live(&self, workspace_key: &str) -> Result { @@ -856,13 +899,42 @@ const IDENTITY_METADATA_KEY: &str = "identity_key"; /// under the same name is only treated as a reclaim of that SAME work unit /// if the value presented then matches what's stored — never by the name /// string alone. Absent, registration cannot reclaim on collision. -fn agent_identity_key() -> Option { +pub(crate) fn agent_identity_key() -> Option { std::env::var("RELAY_AGENT_IDENTITY_KEY") .ok() .map(|value| value.trim().to_string()) .filter(|value| !value.is_empty()) } +/// Stable per-node identity proof derived from the broker's own persisted +/// state directory, for the ONE caller (the broker's own startup +/// registration) that needs restart-reclaim without an operator having to +/// set `RELAY_AGENT_IDENTITY_KEY` by hand. The same project/state directory +/// always hashes to the same value across a kill + restart, while a +/// different project (or a different, unrelated agent) hashes to something +/// else — so it participates in the same fail-closed identity check as any +/// other identity key, never bypassing it. +pub(crate) fn stable_node_identity_key(state_path: &std::path::Path) -> String { + let mut hasher = Sha256::new(); + hasher.update(state_path.to_string_lossy().as_bytes()); + format!("node-{:x}", hasher.finalize()) +} + +/// One-way verifier stored in place of a caller's raw identity key. +/// +/// `admit_agent_registration` stamps this (never the raw key) onto the +/// agent's metadata. Metadata is readable by any caller holding the same +/// workspace key (`get_agent` returns it), so storing the raw key there +/// would let any workspace member replay it verbatim to reclaim another +/// work unit's credentials — the identity check would authenticate nothing. +/// Hashing keeps the credential-bearing capability confined to whoever +/// starts with the original key. +fn hash_identity_key(raw: &str) -> String { + let mut hasher = Sha256::new(); + hasher.update(raw.as_bytes()); + format!("{:x}", hasher.finalize()) +} + /// Single spawn-admission decision shared by every registration path /// (formerly split between a strict branch that always reclaimed a /// name-collision via `register_or_get_agent` — handing the caller the @@ -877,21 +949,24 @@ fn agent_identity_key() -> Option { /// /// Reclaim (return the existing agent's identity with a freshly rotated /// token, so a crashed broker's resume path keeps working) is permitted -/// ONLY when the registration request proves it is the same work unit: its -/// `RELAY_AGENT_IDENTITY_KEY` (see `agent_identity_key`) must match the -/// identity key stamped on the existing agent's metadata at its own -/// creation. Absent or mismatched identity on a collision is rejected. +/// ONLY when the registration request proves it is the same work unit: the +/// caller-supplied `identity_key` (see `agent_identity_key` and +/// `stable_node_identity_key`) must match the identity key stamped on the +/// existing agent's metadata at its own creation — compared by hash +/// (`hash_identity_key`), never by raw value, since metadata is readable by +/// any caller with the workspace key. Absent or mismatched identity on a +/// collision is rejected. async fn admit_agent_registration( relay: &RelayCast, name: &str, agent_type: Option<&str>, + identity_key: Option<&str>, ) -> Result<(String, String, String, Option)> { - let identity_key = agent_identity_key(); - let metadata = identity_key.as_ref().map(|key| { + let metadata = identity_key.map(|key| { let mut map = serde_json::Map::new(); map.insert( IDENTITY_METADATA_KEY.to_string(), - Value::String(key.clone()), + Value::String(hash_identity_key(key)), ); map }); @@ -913,8 +988,8 @@ async fn admit_agent_registration( .and_then(Value::as_str); let reclaims_same_work_unit = matches!( - (identity_key.as_deref(), existing_identity), - (Some(ours), Some(theirs)) if ours == theirs + (identity_key, existing_identity), + (Some(ours), Some(theirs)) if hash_identity_key(ours) == theirs ); if !reclaims_same_work_unit { @@ -990,8 +1065,9 @@ mod tests { use serde_json::json; use super::{ - is_agent_token_invalid, is_agent_token_invalid_anyhow, is_agent_token_invalid_code, - relay_error_to_anyhow, AuthClient, CredentialCache, AGENT_TOKEN_INVALID_CODE, + hash_identity_key, is_agent_token_invalid, is_agent_token_invalid_anyhow, + is_agent_token_invalid_code, relay_error_to_anyhow, stable_node_identity_key, AuthClient, + CredentialCache, AGENT_TOKEN_INVALID_CODE, }; use relaycast::RelayError; @@ -1372,6 +1448,10 @@ mod tests { std::env::set_var("RELAY_API_KEY", "rk_live_shared"); std::env::set_var("RELAY_AGENT_IDENTITY_KEY", "work-unit-42"); } + // The identity proof is stored (and matched) as a one-way hash, never + // the raw value: metadata is readable by any caller with the same + // workspace key, so a raw value there would let a co-tenant replay it. + let identity_hash = hash_identity_key("work-unit-42"); let conflict = server.mock(|when, then| { when.method(POST) .path("/v1/agents") @@ -1379,7 +1459,7 @@ mod tests { .json_body(json!({ "name": "lead", "type": "agent", - "metadata": { "identity_key": "work-unit-42" } + "metadata": { "identity_key": identity_hash } })); then.status(409) .header("content-type", "application/json") @@ -1391,7 +1471,9 @@ mod tests { .header("authorization", "Bearer rk_live_shared"); then.status(200) .header("content-type", "application/json") - .body(r#"{"ok":true,"data":{"id":"a_existing","name":"lead","type":"agent","status":"offline","persona":null,"metadata":{"identity_key":"work-unit-42"},"last_seen":"2025-01-01T00:00:00Z","channels":[]}}"#); + .body(format!( + r#"{{"ok":true,"data":{{"id":"a_existing","name":"lead","type":"agent","status":"offline","persona":null,"metadata":{{"identity_key":"{identity_hash}"}},"last_seen":"2025-01-01T00:00:00Z","channels":[]}}}}"# + )); }); let rotate = server.mock(|when, then| { when.method(POST) @@ -1420,6 +1502,156 @@ mod tests { } } + #[test] + fn stable_node_identity_key_is_stable_per_state_path() { + use std::path::Path; + // Same project/state directory (the "same work unit" across a kill + + // restart, per the fleet node harness) must hash identically every + // time, or the broker could never reclaim its own name after a crash. + let a = stable_node_identity_key(Path::new("/tmp/node-a/.agentworkforce/relay/state.json")); + let a_again = + stable_node_identity_key(Path::new("/tmp/node-a/.agentworkforce/relay/state.json")); + assert_eq!(a, a_again); + + // A different project/state directory (a genuinely different node) + // must hash to something else, or two unrelated nodes could reclaim + // each other's registrations. + let b = stable_node_identity_key(Path::new("/tmp/node-b/.agentworkforce/relay/state.json")); + assert_ne!(a, b); + } + + #[tokio::test] + async fn node_restart_reclaims_its_own_prior_registration_via_stable_identity() { + // Regression test for the fleet-matrix restart flake this PR's + // fail-closed change introduced: `agent-relay node up` on a restart + // reuses the same `--broker-name`, and nothing sets + // RELAY_AGENT_IDENTITY_KEY for it — so before `connect_relay` passed + // a stable, path-derived identity, a restart before the stale + // registration was reaped collided on name and was rejected outright, + // and the node never came back online. This exercises the exact + // entry point `connect_relay` calls (`startup_session_set_with_identity`) + // with no env var set, proving the derived identity alone is enough + // to reclaim. + let _env_guard = clear_relay_env(); + let server = MockServer::start(); + unsafe { + std::env::set_var("RELAY_API_KEY", "rk_live_shared"); + } + let stable_identity = stable_node_identity_key(std::path::Path::new( + "/tmp/node-a/.agentworkforce/relay/state.json", + )); + let identity_hash = hash_identity_key(&stable_identity); + let conflict = server.mock(|when, then| { + when.method(POST) + .path("/v1/agents") + .header("authorization", "Bearer rk_live_shared") + .json_body(json!({ + "name": "node-a", + "type": "agent", + "metadata": { "identity_key": identity_hash } + })); + then.status(409) + .header("content-type", "application/json") + .body(r#"{"ok":false,"error":{"code":"agent_already_exists","message":"name_taken"}}"#); + }); + let get_existing = server.mock(|when, then| { + when.method(GET) + .path("/v1/agents/node-a") + .header("authorization", "Bearer rk_live_shared"); + then.status(200) + .header("content-type", "application/json") + .body(format!( + r#"{{"ok":true,"data":{{"id":"a_existing","name":"node-a","type":"agent","status":"offline","persona":null,"metadata":{{"identity_key":"{identity_hash}"}},"last_seen":"2025-01-01T00:00:00Z","channels":[]}}}}"# + )); + }); + let rotate = server.mock(|when, then| { + when.method(POST) + .path("/v1/agents/node-a/rotate-token") + .header("authorization", "Bearer rk_live_shared"); + then.status(200) + .header("content-type", "application/json") + .body(r#"{"ok":true,"data":{"name":"node-a","token":"at_live_rotated"}}"#); + }); + + let client = AuthClient::new(Some(server.base_url())); + let session = client + .startup_session_set_with_identity(Some("node-a"), true, None, Some(&stable_identity)) + .await + .expect("a node restarting under its own stable identity must reclaim, not be rejected") + .default_session() + .cloned() + .expect("a session was registered"); + + assert_eq!(session.token, "at_live_rotated"); + assert_eq!(session.credentials.agent_id, "a_existing"); + conflict.assert_hits(1); + get_existing.assert_hits(1); + rotate.assert_hits(1); + + unsafe { + std::env::remove_var("RELAY_API_KEY"); + } + } + + #[tokio::test] + async fn different_stable_identity_does_not_reclaim_a_same_named_agent() { + // The flip side of the restart-reclaim fix: a DIFFERENT node (a + // different state directory, hence a different derived identity) + // that happens to collide on name must still be rejected — the + // fail-closed gate this PR added must not be silently defeated by + // handing every node an automatic pass on collision. + let _env_guard = clear_relay_env(); + let server = MockServer::start(); + unsafe { + std::env::set_var("RELAY_API_KEY", "rk_live_shared"); + } + let our_identity = stable_node_identity_key(std::path::Path::new( + "/tmp/node-a/.agentworkforce/relay/state.json", + )); + let their_identity_hash = hash_identity_key(&stable_node_identity_key( + std::path::Path::new("/tmp/node-a-impostor/.agentworkforce/relay/state.json"), + )); + let conflict = server.mock(|when, then| { + when.method(POST) + .path("/v1/agents") + .header("authorization", "Bearer rk_live_shared"); + then.status(409) + .header("content-type", "application/json") + .body(r#"{"ok":false,"error":{"code":"agent_already_exists","message":"name_taken"}}"#); + }); + let get_existing = server.mock(|when, then| { + when.method(GET) + .path("/v1/agents/node-a") + .header("authorization", "Bearer rk_live_shared"); + then.status(200) + .header("content-type", "application/json") + .body(format!( + r#"{{"ok":true,"data":{{"id":"a_existing","name":"node-a","type":"agent","status":"offline","persona":null,"metadata":{{"identity_key":"{their_identity_hash}"}},"last_seen":"2025-01-01T00:00:00Z","channels":[]}}}}"# + )); + }); + + let client = AuthClient::new(Some(server.base_url())); + let error = client + .startup_session_set_with_identity(Some("node-a"), true, None, Some(&our_identity)) + .await + .expect_err("a mismatched identity on collision must be rejected, not reclaimed"); + + // `to_string()` on an anyhow::Error only shows the outermost context + // layer; walk the full chain for the admission-gate's own message. + assert!( + error + .chain() + .any(|layer| layer.to_string().contains("did not prove ownership")), + "expected an ownership-mismatch rejection, got: {error:#}" + ); + conflict.assert_hits(1); + get_existing.assert_hits(1); + + unsafe { + std::env::remove_var("RELAY_API_KEY"); + } + } + #[tokio::test] async fn non_strict_name_conflict_without_identity_proof_is_rejected() { // Strict and non-strict registration must agree: the divergence diff --git a/crates/broker/src/relaycast/mod.rs b/crates/broker/src/relaycast/mod.rs index b25e2e5c5..dd14e2c32 100644 --- a/crates/broker/src/relaycast/mod.rs +++ b/crates/broker/src/relaycast/mod.rs @@ -8,7 +8,7 @@ pub(crate) mod ws; pub(crate) use crate::snippets::{ configure_agent_relay_mcp_with_result, configure_agent_relay_mcp_with_token, }; -pub(crate) use auth::AuthClient; +pub(crate) use auth::{agent_identity_key, stable_node_identity_key, AuthClient}; // `is_agent_token_invalid`, `is_agent_token_invalid_anyhow`, // `is_agent_token_invalid_code`, and `AGENT_TOKEN_INVALID_CODE` are declared // `pub` on `auth` so future callers (bridge, ws, listen_api) can reach them diff --git a/crates/broker/src/runtime/mod.rs b/crates/broker/src/runtime/mod.rs index 93ad3e728..164f15de7 100644 --- a/crates/broker/src/runtime/mod.rs +++ b/crates/broker/src/runtime/mod.rs @@ -36,9 +36,10 @@ use crate::{ ProtocolEnvelope, RelayDelivery, ResolvedHarnessConfig, PROTOCOL_VERSION, }, relaycast::{ - format_worker_preregistration_error, registration_retry_after_secs, - retry_agent_registration, AuthClient, MultiWorkspaceSession, RegRetryOutcome, - RelaycastHttpClient, WorkspaceInboundMessage, WorkspaceMembershipSummary, WsControl, + agent_identity_key, format_worker_preregistration_error, registration_retry_after_secs, + retry_agent_registration, stable_node_identity_key, AuthClient, MultiWorkspaceSession, + RegRetryOutcome, RelaycastHttpClient, WorkspaceInboundMessage, WorkspaceMembershipSummary, + WsControl, }, replay_buffer::{ReplayBuffer, DEFAULT_REPLAY_CAPACITY}, telemetry::{ActionSource, TelemetryClient, TelemetryEvent}, diff --git a/crates/broker/src/runtime/session.rs b/crates/broker/src/runtime/session.rs index 366a730d0..9510bacdf 100644 --- a/crates/broker/src/runtime/session.rs +++ b/crates/broker/src/runtime/session.rs @@ -280,16 +280,26 @@ pub(crate) async fn connect_relay(opts: RelaySessionOptions<'_>) -> Result