feat(rust/signal): idempotent maintain_keys for the pre-key lifecycle

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<KeyMaintenance> {
+    // 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(),
+    })
+}

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;

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);
+}