From f8e9fc22f08537ddf2acaafe03cc4d97a207404b Mon Sep 17 00:00:00 2001 From: sanil-23 Date: Tue, 21 Jul 2026 12:34:15 +0530 Subject: [PATCH 1/3] feat(rust/signal): add idempotent maintain_keys for the pre-key lifecycle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `signal::maintain::maintain_keys` — a health-gated, idempotent reconcile of an agent's published pre-key bundle with the relay. It publishes only when needed: a no-op when the store already holds a signed pre-key and the relay's one-time pool is healthy, a (re)publish on first boot, a wiped store, or a low/depleted pool. Storing private material locally before upload, and skipping the upload when the store already backs a healthy set, is the anti-orphan invariant — blindly republishing on every start leaves the relay advertising one-time pre-keys a rotated/wiped store can no longer back, so peers' X3DH fails with a MAC error. Composes existing pieces (KeysApi health/rotate/upload + signal::keys generation + SessionStore), so callers stop hand-rolling the lifecycle. Adds wiremock coverage for the idempotent no-op and the depleted-pool refill paths. Co-Authored-By: Claude --- sdk/rust/src/signal/maintain.rs | 133 ++++++++++++++++++++++++++++++ sdk/rust/src/signal/mod.rs | 1 + sdk/rust/tests/signal_maintain.rs | 120 +++++++++++++++++++++++++++ 3 files changed, 254 insertions(+) create mode 100644 sdk/rust/src/signal/maintain.rs create mode 100644 sdk/rust/tests/signal_maintain.rs diff --git a/sdk/rust/src/signal/maintain.rs b/sdk/rust/src/signal/maintain.rs new file mode 100644 index 00000000..7b33904b --- /dev/null +++ b/sdk/rust/src/signal/maintain.rs @@ -0,0 +1,133 @@ +//! Idempotent key-bundle maintenance. +//! +//! Reconciles an agent's *published* pre-key bundle (on the relay) with the +//! private material in its [`SessionStore`], generating and uploading keys ONLY +//! when the relay actually needs them. +//! +//! This is the anti-orphan invariant. Blindly republishing a fresh bundle on +//! every start leaves the relay advertising one-time pre-keys whose private +//! halves a rotated or wiped store can no longer back, so a peer's X3DH against +//! that bundle fails with a MAC error and its first message is silently dropped. +//! Storing the private material locally *before* uploading, and skipping the +//! upload entirely when the store already backs a healthy published set, +//! guarantees the relay never serves a key this agent cannot answer. + +use std::time::{SystemTime, UNIX_EPOCH}; + +use crate::api::keys::KeysApi; +use crate::error::Result; +use crate::signal::keys::{generate_pre_keys, generate_signed_pre_key, serialize_pre_key}; +use crate::signal::store::SessionStore; +use crate::signer::Signer; +use crate::types::{PreKeysRequest, SignedPreKeyRequest}; + +/// Tuning for [`maintain_keys`]. +#[derive(Debug, Clone)] +pub struct MaintainPolicy { + /// How many one-time pre-keys to (re)publish in a batch when the pool needs + /// filling. + pub one_time_batch: usize, +} + +impl Default for MaintainPolicy { + fn default() -> Self { + Self { one_time_batch: 20 } + } +} + +/// What [`maintain_keys`] did on a call. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub struct KeyMaintenance { + /// The store already held a signed pre-key and the relay's one-time pool was + /// healthy, so nothing was published — the common restart path. + pub was_healthy: bool, + /// A fresh signed pre-key was generated and rotated on the relay. + pub rotated_signed: bool, + /// How many one-time pre-keys were generated and uploaded. + pub uploaded_one_time: usize, +} + +/// Idempotently reconcile `agent_id`'s published key bundle with the relay. +/// +/// Safe to call on every boot and periodically — it publishes only when the +/// relay actually needs it: +/// +/// - **No-op** when the store already holds a signed pre-key AND the relay +/// reports a non-empty, not-low one-time pool. *Not* republishing here is what +/// prevents orphaning. +/// - **(Re)publish** on first boot, a store that lost its keys, or a low/empty +/// relay pool: generate a signed pre-key plus `policy.one_time_batch` one-time +/// pre-keys, store the private halves locally, then rotate + upload the public +/// parts. +/// +/// `identity_key_base64` is the agent's base64 Ed25519 identity key; the relay +/// verifies each published key's signature against it. +/// +/// # Errors +/// +/// Returns the underlying error if key generation, a store write, the signed +/// pre-key rotation, or the one-time upload fails. A failed *health* check is +/// treated as "cannot confirm healthy" and falls through to a (re)publish rather +/// than erroring, so a transient relay blip never leaves the agent unpublished. +pub async fn maintain_keys( + keys: &KeysApi, + store: &dyn SessionStore, + signer: &dyn Signer, + agent_id: &str, + identity_key_base64: &str, + policy: &MaintainPolicy, +) -> Result { + // A signed pre-key in the store means we've published before and can back + // what the relay serves. If the relay's one-time pool is also healthy, there + // is nothing to do — and NOT republishing is precisely what avoids orphaning. + if store.active_signed_pre_key().await.is_ok() { + if let Ok(health) = keys.health(agent_id).await { + if !health.low_one_time_pre_keys && health.one_time_pre_key_count > 0 { + return Ok(KeyMaintenance { + was_healthy: true, + ..Default::default() + }); + } + } + } + + let now_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0); + + // Signed pre-key: (re)generate, store the private half, then rotate on the + // relay so peers pick up the new public half. + let spk = generate_signed_pre_key(signer, &format!("spk_{now_secs}")).await?; + store.store_signed_pre_key(spk.clone()).await?; + keys.rotate_signed_pre_key( + agent_id, + &SignedPreKeyRequest { + identity_key: Some(identity_key_base64.to_string()), + signed_pre_key: serialize_pre_key(&spk), + }, + ) + .await?; + + // One-time pre-keys: generate a batch, store the private halves BEFORE upload + // so the relay never advertises a key the store cannot back, then upload the + // public parts. + let one_time = generate_pre_keys(signer, now_secs, policy.one_time_batch).await?; + for key in &one_time { + store.store_pre_key(key.clone()).await?; + } + keys.upload_pre_keys( + agent_id, + &PreKeysRequest { + identity_key: Some(identity_key_base64.to_string()), + pre_keys: one_time.iter().map(serialize_pre_key).collect(), + }, + ) + .await?; + + Ok(KeyMaintenance { + was_healthy: false, + rotated_signed: true, + uploaded_one_time: one_time.len(), + }) +} diff --git a/sdk/rust/src/signal/mod.rs b/sdk/rust/src/signal/mod.rs index 8df5ed7e..462ad2f4 100644 --- a/sdk/rust/src/signal/mod.rs +++ b/sdk/rust/src/signal/mod.rs @@ -7,6 +7,7 @@ pub mod crypto; pub mod keys; +pub mod maintain; pub mod memory_store; pub mod ratchet; pub mod sender_key; diff --git a/sdk/rust/tests/signal_maintain.rs b/sdk/rust/tests/signal_maintain.rs new file mode 100644 index 00000000..7c46d404 --- /dev/null +++ b/sdk/rust/tests/signal_maintain.rs @@ -0,0 +1,120 @@ +//! Coverage for idempotent key-bundle maintenance +//! ([`tinyplace::signal::maintain::maintain_keys`]): it (re)publishes on first +//! boot / a depleted pool, and is a no-op when the store already backs a healthy +//! published set — the anti-orphan invariant. + +mod common; + +use common::{all_requests, client_for, test_signer}; +use serde_json::json; +use wiremock::matchers::method; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +use tinyplace::signal::crypto::generate_x25519_keypair; +use tinyplace::signal::maintain::{maintain_keys, MaintainPolicy}; +use tinyplace::signal::memory_store::MemorySessionStore; +use tinyplace::Signer; + +/// A relay mock whose `GET /health` reports `one_time_count` one-time pre-keys, +/// and whose `PUT`s (rotate signed pre-key / upload one-time) return `null`. +async fn relay_with_health(one_time_count: i64) -> MockServer { + let server = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "agentId": "x", + "oneTimePreKeyCount": one_time_count, + "lowOneTimePreKeys": one_time_count < 5, + }))) + .mount(&server) + .await; + Mock::given(method("PUT")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!(null))) + .mount(&server) + .await; + server +} + +#[tokio::test] +async fn maintain_keys_publishes_once_then_no_ops_when_healthy() { + let server = relay_with_health(20).await; + let client = client_for(&server); + let signer = test_signer(); + let store = MemorySessionStore::new(generate_x25519_keypair()); + let agent_id = signer.agent_id(); + let policy = MaintainPolicy::default(); + + // First call: empty store → (re)publish a full bundle. + let first = maintain_keys( + &client.keys, + &store, + signer.as_ref(), + &agent_id, + "identity-key", + &policy, + ) + .await + .unwrap(); + assert!(!first.was_healthy); + assert!(first.rotated_signed); + assert_eq!(first.uploaded_one_time, 20); + + // Second call: the store now holds a signed pre-key AND the relay pool is + // healthy → NO-OP. Not republishing is what avoids orphaning served keys. + let second = maintain_keys( + &client.keys, + &store, + signer.as_ref(), + &agent_id, + "identity-key", + &policy, + ) + .await + .unwrap(); + assert!(second.was_healthy); + assert_eq!(second.uploaded_one_time, 0); + + // On the wire: exactly ONE one-time upload across both calls. + let uploads = all_requests(&server) + .await + .into_iter() + .filter(|r| r.method.as_str() == "PUT" && r.url.path().ends_with("/prekeys")) + .count(); + assert_eq!(uploads, 1, "idempotent maintain uploads only once"); +} + +#[tokio::test] +async fn maintain_keys_refills_a_depleted_pool() { + // The relay reports an empty pool (every one-time key consumed by peers). + let server = relay_with_health(0).await; + let client = client_for(&server); + let signer = test_signer(); + let store = MemorySessionStore::new(generate_x25519_keypair()); + let agent_id = signer.agent_id(); + let policy = MaintainPolicy::default(); + + maintain_keys( + &client.keys, + &store, + signer.as_ref(), + &agent_id, + "identity-key", + &policy, + ) + .await + .unwrap(); + + // Store now holds a signed pre-key, but the relay pool is empty → republish + // rather than no-op, so new peers can still complete a full X3DH. + let again = maintain_keys( + &client.keys, + &store, + signer.as_ref(), + &agent_id, + "identity-key", + &policy, + ) + .await + .unwrap(); + assert!(!again.was_healthy); + assert_eq!(again.uploaded_one_time, 20); +} From b661d8acafc5166c12a18d88e6bbf1a3bc02ff0a Mon Sep 17 00:00:00 2001 From: sanil-23 Date: Tue, 21 Jul 2026 12:56:16 +0530 Subject: [PATCH 2/3] fix(rust/signal): decouple signed-key and one-time maintenance MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses review on #267. The no-op gate previously required BOTH "store has an active signed pre-key" AND "relay pool healthy", so a merely-depleted one-time pool fell through to a full republish that needlessly rotated a perfectly good signed pre-key. Worse, it never checked WHICH signed pre-key the relay advertises: `active_signed_pre_key().is_ok()` is true even when the store's active key is not the one the relay serves, so an inbound PREKEY_BUNDLE naming the relay's id would fail to decrypt and maintenance would never repair it. The two halves are now maintained independently: - Signed pre-key: rotated only when the store cannot answer `KeyHealth::signed_pre_key_key_id` — the exact id an inbound PREKEY_BUNDLE names. A still-backed key is left untouched even when the pool needs a refill. - One-time pre-keys: topped up only on the relay's own low/empty report, without disturbing a healthy signed pre-key. Tests now pin one half healthy and the other degraded to prove neither drags the other into a needless republish, plus the first-boot both-halves case. Co-Authored-By: Claude --- sdk/rust/src/signal/maintain.rs | 131 ++++++++++++++----------- sdk/rust/tests/signal_maintain.rs | 157 ++++++++++++++++++++---------- 2 files changed, 179 insertions(+), 109 deletions(-) diff --git a/sdk/rust/src/signal/maintain.rs b/sdk/rust/src/signal/maintain.rs index 7b33904b..7e173c48 100644 --- a/sdk/rust/src/signal/maintain.rs +++ b/sdk/rust/src/signal/maintain.rs @@ -38,10 +38,11 @@ impl Default for MaintainPolicy { /// What [`maintain_keys`] did on a call. #[derive(Debug, Default, Clone, PartialEq, Eq)] pub struct KeyMaintenance { - /// The store already held a signed pre-key and the relay's one-time pool was - /// healthy, so nothing was published — the common restart path. + /// Nothing needed publishing: the store backed the signed pre-key the relay + /// advertises and its one-time pool was healthy — the common restart path. pub was_healthy: bool, - /// A fresh signed pre-key was generated and rotated on the relay. + /// A fresh signed pre-key was generated and rotated on the relay, because the + /// store could not answer the id the relay was advertising. pub rotated_signed: bool, /// How many one-time pre-keys were generated and uploaded. pub uploaded_one_time: usize, @@ -49,16 +50,22 @@ pub struct KeyMaintenance { /// Idempotently reconcile `agent_id`'s published key bundle with the relay. /// -/// Safe to call on every boot and periodically — it publishes only when the -/// relay actually needs it: +/// Safe to call on every boot and periodically — it publishes only what the +/// relay actually needs. The two halves are maintained **independently**: /// -/// - **No-op** when the store already holds a signed pre-key AND the relay -/// reports a non-empty, not-low one-time pool. *Not* republishing here is what -/// prevents orphaning. -/// - **(Re)publish** on first boot, a store that lost its keys, or a low/empty -/// relay pool: generate a signed pre-key plus `policy.one_time_batch` one-time -/// pre-keys, store the private halves locally, then rotate + upload the public -/// parts. +/// - **Signed pre-key** — rotated only when the store cannot answer the id the +/// relay advertises (`KeyHealth::signed_pre_key_key_id`). An inbound +/// `PREKEY_BUNDLE` names that exact id, so a relay serving a key we no longer +/// hold breaks first contact; a still-backed signed key is left untouched. +/// - **One-time pre-keys** — topped up only when the relay reports the pool low +/// or empty, without disturbing a healthy signed pre-key. +/// +/// So a merely-depleted one-time pool refills without needless signed-key +/// churn, and a stale/unbacked signed key is repaired even while the pool is +/// full. When both are satisfied the call is a **no-op** +/// ([`KeyMaintenance::was_healthy`]) — not republishing is what prevents +/// orphaning. If the health check itself fails, both halves are (re)published +/// rather than risk leaving the agent unpublished. /// /// `identity_key_base64` is the agent's base64 Ed25519 identity key; the relay /// verifies each published key's signature against it. @@ -77,57 +84,69 @@ pub async fn maintain_keys( identity_key_base64: &str, policy: &MaintainPolicy, ) -> Result { - // A signed pre-key in the store means we've published before and can back - // what the relay serves. If the relay's one-time pool is also healthy, there - // is nothing to do — and NOT republishing is precisely what avoids orphaning. - if store.active_signed_pre_key().await.is_ok() { - if let Ok(health) = keys.health(agent_id).await { - if !health.low_one_time_pre_keys && health.one_time_pre_key_count > 0 { - return Ok(KeyMaintenance { - was_healthy: true, - ..Default::default() - }); - } - } - } + // What is the relay actually serving right now? If we cannot reach it we + // cannot confirm anything is healthy, so both halves are (re)published + // rather than risk leaving the agent unpublished. + let health = keys.health(agent_id).await.ok(); let now_secs = SystemTime::now() .duration_since(UNIX_EPOCH) .map(|d| d.as_secs()) .unwrap_or(0); + let mut report = KeyMaintenance::default(); - // Signed pre-key: (re)generate, store the private half, then rotate on the - // relay so peers pick up the new public half. - let spk = generate_signed_pre_key(signer, &format!("spk_{now_secs}")).await?; - store.store_signed_pre_key(spk.clone()).await?; - keys.rotate_signed_pre_key( - agent_id, - &SignedPreKeyRequest { - identity_key: Some(identity_key_base64.to_string()), - signed_pre_key: serialize_pre_key(&spk), - }, - ) - .await?; + // --- Signed pre-key ----------------------------------------------------- + // Rotate ONLY when the store cannot answer the signed pre-key the relay + // advertises. An inbound `PREKEY_BUNDLE` names that exact id, so a relay + // serving an id we no longer hold makes first contact fail; conversely a + // still-backed signed key is left alone even when the one-time pool needs a + // refill (the two halves are maintained independently). + let signed_pre_key_backed = match health + .as_ref() + .and_then(|h| h.signed_pre_key_key_id.as_deref()) + { + Some(key_id) => store.signed_pre_key(key_id).await.ok().flatten().is_some(), + // No health, or the relay advertises no signed pre-key at all. + None => false, + }; + if !signed_pre_key_backed { + let spk = generate_signed_pre_key(signer, &format!("spk_{now_secs}")).await?; + store.store_signed_pre_key(spk.clone()).await?; + keys.rotate_signed_pre_key( + agent_id, + &SignedPreKeyRequest { + identity_key: Some(identity_key_base64.to_string()), + signed_pre_key: serialize_pre_key(&spk), + }, + ) + .await?; + report.rotated_signed = true; + } - // One-time pre-keys: generate a batch, store the private halves BEFORE upload - // so the relay never advertises a key the store cannot back, then upload the - // public parts. - let one_time = generate_pre_keys(signer, now_secs, policy.one_time_batch).await?; - for key in &one_time { - store.store_pre_key(key.clone()).await?; + // --- One-time pre-keys -------------------------------------------------- + // Topped up independently, driven only by the relay's own pool report. The + // private halves are stored BEFORE upload so the relay never advertises a + // key the store cannot back. + let one_time_depleted = match health.as_ref() { + Some(h) => h.low_one_time_pre_keys || h.one_time_pre_key_count <= 0, + None => true, + }; + if one_time_depleted { + let one_time = generate_pre_keys(signer, now_secs, policy.one_time_batch).await?; + for key in &one_time { + store.store_pre_key(key.clone()).await?; + } + keys.upload_pre_keys( + agent_id, + &PreKeysRequest { + identity_key: Some(identity_key_base64.to_string()), + pre_keys: one_time.iter().map(serialize_pre_key).collect(), + }, + ) + .await?; + report.uploaded_one_time = one_time.len(); } - keys.upload_pre_keys( - agent_id, - &PreKeysRequest { - identity_key: Some(identity_key_base64.to_string()), - pre_keys: one_time.iter().map(serialize_pre_key).collect(), - }, - ) - .await?; - Ok(KeyMaintenance { - was_healthy: false, - rotated_signed: true, - uploaded_one_time: one_time.len(), - }) + report.was_healthy = !report.rotated_signed && report.uploaded_one_time == 0; + Ok(report) } diff --git a/sdk/rust/tests/signal_maintain.rs b/sdk/rust/tests/signal_maintain.rs index 7c46d404..e5ef872c 100644 --- a/sdk/rust/tests/signal_maintain.rs +++ b/sdk/rust/tests/signal_maintain.rs @@ -1,29 +1,35 @@ //! Coverage for idempotent key-bundle maintenance -//! ([`tinyplace::signal::maintain::maintain_keys`]): it (re)publishes on first -//! boot / a depleted pool, and is a no-op when the store already backs a healthy -//! published set — the anti-orphan invariant. +//! ([`tinyplace::signal::maintain::maintain_keys`]). +//! +//! The signed pre-key and the one-time pool are maintained **independently**, so +//! each test pins one half healthy and the other degraded to prove neither +//! drags the other into a needless republish. mod common; use common::{all_requests, client_for, test_signer}; -use serde_json::json; +use serde_json::{json, Value}; use wiremock::matchers::method; use wiremock::{Mock, MockServer, ResponseTemplate}; use tinyplace::signal::crypto::generate_x25519_keypair; +use tinyplace::signal::keys::generate_signed_pre_key; use tinyplace::signal::maintain::{maintain_keys, MaintainPolicy}; use tinyplace::signal::memory_store::MemorySessionStore; -use tinyplace::Signer; +use tinyplace::signal::store::SessionStore; +use tinyplace::{LocalSigner, Signer}; -/// A relay mock whose `GET /health` reports `one_time_count` one-time pre-keys, -/// and whose `PUT`s (rotate signed pre-key / upload one-time) return `null`. -async fn relay_with_health(one_time_count: i64) -> MockServer { +/// A relay mock whose `GET /health` reports `one_time_count` one-time pre-keys +/// and advertises `signed_id` as its signed pre-key, and whose `PUT`s (rotate / +/// upload) return `null`. +async fn relay(one_time_count: i64, signed_id: Value) -> MockServer { let server = MockServer::start().await; Mock::given(method("GET")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "agentId": "x", "oneTimePreKeyCount": one_time_count, "lowOneTimePreKeys": one_time_count < 5, + "signedPreKeyKeyId": signed_id, }))) .mount(&server) .await; @@ -34,87 +40,132 @@ async fn relay_with_health(one_time_count: i64) -> MockServer { server } +/// A store already backing a signed pre-key under `key_id`. +async fn store_backing(signer: &LocalSigner, key_id: &str) -> MemorySessionStore { + let store = MemorySessionStore::new(generate_x25519_keypair()); + let spk = generate_signed_pre_key(signer, key_id).await.unwrap(); + store.store_signed_pre_key(spk).await.unwrap(); + store +} + +/// Count the `PUT`s that hit each key route. +async fn puts(server: &MockServer) -> (usize, usize) { + let requests = all_requests(server).await; + let count = |suffix: &str| { + requests + .iter() + .filter(|r| r.method.as_str() == "PUT" && r.url.path().ends_with(suffix)) + .count() + }; + (count("/signed-prekey"), count("/prekeys")) +} + #[tokio::test] -async fn maintain_keys_publishes_once_then_no_ops_when_healthy() { - let server = relay_with_health(20).await; - let client = client_for(&server); +async fn no_ops_when_the_advertised_signed_key_is_backed_and_the_pool_is_healthy() { let signer = test_signer(); - let store = MemorySessionStore::new(generate_x25519_keypair()); - let agent_id = signer.agent_id(); - let policy = MaintainPolicy::default(); + let server = relay(20, json!("spk_known")).await; + let client = client_for(&server); + let store = store_backing(&signer, "spk_known").await; - // First call: empty store → (re)publish a full bundle. - let first = maintain_keys( + let report = maintain_keys( &client.keys, &store, signer.as_ref(), - &agent_id, + &signer.agent_id(), "identity-key", - &policy, + &MaintainPolicy::default(), ) .await .unwrap(); - assert!(!first.was_healthy); - assert!(first.rotated_signed); - assert_eq!(first.uploaded_one_time, 20); - // Second call: the store now holds a signed pre-key AND the relay pool is - // healthy → NO-OP. Not republishing is what avoids orphaning served keys. - let second = maintain_keys( + assert!(report.was_healthy); + assert!(!report.rotated_signed); + assert_eq!(report.uploaded_one_time, 0); + assert_eq!( + puts(&server).await, + (0, 0), + "a healthy agent publishes nothing" + ); +} + +#[tokio::test] +async fn rotates_only_the_signed_key_when_the_relay_advertises_one_we_cannot_back() { + // The relay serves `spk_missing`, which the store does not hold — an inbound + // PREKEY_BUNDLE naming it would fail to decrypt, so maintenance must repair + // it. The one-time pool is healthy and must be left alone. + let signer = test_signer(); + let server = relay(20, json!("spk_missing")).await; + let client = client_for(&server); + let store = store_backing(&signer, "spk_other").await; + + let report = maintain_keys( &client.keys, &store, signer.as_ref(), - &agent_id, + &signer.agent_id(), "identity-key", - &policy, + &MaintainPolicy::default(), ) .await .unwrap(); - assert!(second.was_healthy); - assert_eq!(second.uploaded_one_time, 0); - - // On the wire: exactly ONE one-time upload across both calls. - let uploads = all_requests(&server) - .await - .into_iter() - .filter(|r| r.method.as_str() == "PUT" && r.url.path().ends_with("/prekeys")) - .count(); - assert_eq!(uploads, 1, "idempotent maintain uploads only once"); + + assert!(!report.was_healthy); + assert!(report.rotated_signed); + assert_eq!( + report.uploaded_one_time, 0, + "a healthy pool is not disturbed" + ); + assert_eq!(puts(&server).await, (1, 0)); } #[tokio::test] -async fn maintain_keys_refills_a_depleted_pool() { - // The relay reports an empty pool (every one-time key consumed by peers). - let server = relay_with_health(0).await; - let client = client_for(&server); +async fn tops_up_only_the_pool_when_depleted_leaving_a_good_signed_key_alone() { + // Every one-time key was consumed by peers, but the signed pre-key the relay + // advertises is still backed — refill without needless signed-key churn. let signer = test_signer(); - let store = MemorySessionStore::new(generate_x25519_keypair()); - let agent_id = signer.agent_id(); - let policy = MaintainPolicy::default(); + let server = relay(0, json!("spk_known")).await; + let client = client_for(&server); + let store = store_backing(&signer, "spk_known").await; - maintain_keys( + let report = maintain_keys( &client.keys, &store, signer.as_ref(), - &agent_id, + &signer.agent_id(), "identity-key", - &policy, + &MaintainPolicy::default(), ) .await .unwrap(); - // Store now holds a signed pre-key, but the relay pool is empty → republish - // rather than no-op, so new peers can still complete a full X3DH. - let again = maintain_keys( + assert!(!report.was_healthy); + assert!(!report.rotated_signed, "signed key must not be rotated"); + assert_eq!(report.uploaded_one_time, 20); + assert_eq!(puts(&server).await, (0, 1)); +} + +#[tokio::test] +async fn publishes_both_halves_on_first_boot() { + // Nothing published yet: the relay advertises no signed pre-key and an empty + // pool, and the store is empty. + let signer = test_signer(); + let server = relay(0, json!(null)).await; + let client = client_for(&server); + let store = MemorySessionStore::new(generate_x25519_keypair()); + + let report = maintain_keys( &client.keys, &store, signer.as_ref(), - &agent_id, + &signer.agent_id(), "identity-key", - &policy, + &MaintainPolicy::default(), ) .await .unwrap(); - assert!(!again.was_healthy); - assert_eq!(again.uploaded_one_time, 20); + + assert!(!report.was_healthy); + assert!(report.rotated_signed); + assert_eq!(report.uploaded_one_time, 20); + assert_eq!(puts(&server).await, (1, 1)); } From 531ac772f5882c5393c973703b5a2bf222be6422 Mon Sep 17 00:00:00 2001 From: sanil-23 Date: Tue, 21 Jul 2026 13:05:09 +0530 Subject: [PATCH 3/3] test(rust/signal): pin the health route and the low-pool branch Addresses review on #267. - The GET mock matched any path, so a wrong health route would still satisfy every scenario. Pin it to `^/keys/[^/]+/health$` (the id segment is the agent's base58 key, not a literal). - `tops_up_only_the_pool_when_low` used a zero count, where BOTH disjuncts of `low_one_time_pre_keys || one_time_pre_key_count <= 0` are true, so the low-flag branch was never isolated. Use a positive-but-low count (4) to pin that branch specifically; the empty-pool branch stays covered by the first-boot case. Co-Authored-By: Claude --- sdk/rust/tests/signal_maintain.rs | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/sdk/rust/tests/signal_maintain.rs b/sdk/rust/tests/signal_maintain.rs index e5ef872c..66c41c35 100644 --- a/sdk/rust/tests/signal_maintain.rs +++ b/sdk/rust/tests/signal_maintain.rs @@ -9,7 +9,7 @@ mod common; use common::{all_requests, client_for, test_signer}; use serde_json::{json, Value}; -use wiremock::matchers::method; +use wiremock::matchers::{method, path_regex}; use wiremock::{Mock, MockServer, ResponseTemplate}; use tinyplace::signal::crypto::generate_x25519_keypair; @@ -24,7 +24,10 @@ use tinyplace::{LocalSigner, Signer}; /// upload) return `null`. async fn relay(one_time_count: i64, signed_id: Value) -> MockServer { let server = MockServer::start().await; + // Pinned to the real health route (the id segment is the agent's base58 key), + // so a wrong path would fail the scenario instead of silently matching. Mock::given(method("GET")) + .and(path_regex(r"^/keys/[^/]+/health$")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "agentId": "x", "oneTimePreKeyCount": one_time_count, @@ -119,11 +122,14 @@ async fn rotates_only_the_signed_key_when_the_relay_advertises_one_we_cannot_bac } #[tokio::test] -async fn tops_up_only_the_pool_when_depleted_leaving_a_good_signed_key_alone() { - // Every one-time key was consumed by peers, but the signed pre-key the relay - // advertises is still backed — refill without needless signed-key churn. +async fn tops_up_only_the_pool_when_low_leaving_a_good_signed_key_alone() { + // Peers have drawn the pool DOWN but not to zero: the relay still reports a + // positive count and only raises `lowOneTimePreKeys`. This pins the low-flag + // branch specifically (the empty-pool `count <= 0` branch is covered by the + // first-boot case), while the advertised signed pre-key is still backed — so + // refill must happen without any signed-key churn. let signer = test_signer(); - let server = relay(0, json!("spk_known")).await; + let server = relay(4, json!("spk_known")).await; let client = client_for(&server); let store = store_backing(&signer, "spk_known").await;