fix: prevent attestation cache poisoning on key rotation (#5)

src/attestation/key_registry.rs

@@ -0,0 +1,202 @@
+//! Validator key registry and rotation-safe verification cache (#5).
+//!
+//! ## The bug
+//!
+//! `verify_attestation()` cached each validator's public key keyed only by
+//! validator id, refreshed on a timer. `rotate_key()` cleared the cache entry
+//! and updated the on-chain registry, but the on-chain write is not visible for
+//! 3–5s (Soroban finality). A verify landing in that window missed the cache,
+//! reloaded the *old* key from the registry, re-cached it, and rejected every
+//! attestation now signed with the new key — a multi-second outage per
+//! rotation.
+//!
+//! ## The fix
+//!
+//! Three reinforcing measures from the resolution blueprint:
+//!
+//! 1. **Two-phase rotation (step 1):** a rotation keeps the old key valid as a
+//!    `previous_key` for a bounded window of ledgers, so attestations still
+//!    in flight under the old key continue to verify.
+//! 2. **Key-generation counter + versioned cache (steps 2 & 4):** each
+//!    validator carries a `key_gen` bumped on every rotation. The cache stores
+//!    the generation it was populated at; a verify whose registry generation is
+//!    newer than the cached one forces a reload, so a stale entry is never
+//!    reused after a rotation.
+//! 3. **Multi-key cache (step 3):** the reloaded entry holds *both* the new and
+//!    (within the window) the old key; verification accepts if any cached key
+//!    matches. The previous key is dropped once the window elapses.
+//!
+//! Together these guarantee `verify(attestation, keys) == true` for every
+//! attestation signed after a rotation — both late old-key attestations (within
+//! the window) and new-key attestations.
+
+extern crate alloc;
+use alloc::collections::BTreeMap;
+use alloc::vec::Vec;
+
+use crate::attestation::verifier::{verify_attestation_signature, AttestationData, SecretKey, Signature};
+use crate::crypto::domain::Domain;
+
+/// Validator identifier.
+pub type ValidatorId = u32;
+
+/// A validator verification key. (In this crate's signing model the key type is
+/// the same 32-byte value used by [`verify_attestation_signature`].)
+pub type VerifyKey = SecretKey;
+
+/// Number of ledgers the previous key stays valid after a rotation (~10 per the
+/// blueprint, covering the on-chain finality gap with margin).
+pub const ROTATION_WINDOW_LEDGERS: u64 = 10;
+
+/// An immutable snapshot of a validator's key state.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct KeySnapshot {
+    pub key_gen: u64,
+    pub active_key: VerifyKey,
+    pub previous_key: Option<VerifyKey>,
+    pub rotated_at_ledger: u64,
+}
+
+/// On-chain validator key registry.
+#[derive(Clone, Debug, Default)]
+pub struct KeyRegistry {
+    records: BTreeMap<ValidatorId, KeySnapshot>,
+}
+
+impl KeyRegistry {
+    /// Create an empty registry.
+    pub fn new() -> Self {
+        Self {
+            records: BTreeMap::new(),
+        }
+    }
+
+    /// Register a validator's initial key at generation 0.
+    pub fn register(&mut self, validator: ValidatorId, key: VerifyKey) {
+        self.records.insert(
+            validator,
+            KeySnapshot {
+                key_gen: 0,
+                active_key: key,
+                previous_key: None,
+                rotated_at_ledger: 0,
+            },
+        );
+    }
+
+    /// Rotate a validator's key: bump `key_gen`, promote `new_key` to active,
+    /// and retain the prior key as `previous_key` (valid during the rotation
+    /// window). Returns the new generation, or `None` if the validator is
+    /// unknown.
+    pub fn rotate_key(
+        &mut self,
+        validator: ValidatorId,
+        new_key: VerifyKey,
+        current_ledger: u64,
+    ) -> Option<u64> {
+        let record = self.records.get_mut(&validator)?;
+        record.previous_key = Some(record.active_key);
+        record.active_key = new_key;
+        record.key_gen += 1;
+        record.rotated_at_ledger = current_ledger;
+        Some(record.key_gen)
+    }
+
+    /// Current key generation for `validator`.
+    pub fn key_gen(&self, validator: ValidatorId) -> Option<u64> {
+        self.records.get(&validator).map(|r| r.key_gen)
+    }
+
+    /// Immutable snapshot of `validator`'s key state.
+    pub fn snapshot(&self, validator: ValidatorId) -> Option<KeySnapshot> {
+        self.records.get(&validator).copied()
+    }
+}
+
+/// In-memory verification cache, versioned by `key_gen`.
+#[derive(Clone, Debug)]
+pub struct VerificationCache {
+    window_ledgers: u64,
+    entries: BTreeMap<ValidatorId, KeySnapshot>,
+}
+
+impl Default for VerificationCache {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl VerificationCache {
+    /// Cache with the default rotation window.
+    pub fn new() -> Self {
+        Self::with_window(ROTATION_WINDOW_LEDGERS)
+    }
+
+    /// Cache with a custom rotation window (ledgers).
+    pub fn with_window(window_ledgers: u64) -> Self {
+        Self {
+            window_ledgers,
+            entries: BTreeMap::new(),
+        }
+    }
+
+    /// Generation the cache currently holds for `validator`, if any.
+    pub fn cached_key_gen(&self, validator: ValidatorId) -> Option<u64> {
+        self.entries.get(&validator).map(|c| c.key_gen)
+    }
+
+    /// Resolve the set of keys currently valid for `validator`, reloading from
+    /// the registry whenever the cached generation is behind. The previous key
+    /// is included only while inside the rotation window.
+    fn resolve_keys(
+        &mut self,
+        registry: &KeyRegistry,
+        validator: ValidatorId,
+        current_ledger: u64,
+    ) -> Option<Vec<VerifyKey>> {
+        let snapshot = registry.snapshot(validator)?;
+
+        let needs_reload = match self.entries.get(&validator) {
+            Some(cached) => cached.key_gen < snapshot.key_gen,
+            None => true,
+        };
+        if needs_reload {
+            self.entries.insert(validator, snapshot);
+        }
+
+        // Post-reload the cached entry equals the registry snapshot; otherwise
+        // it carries the same generation, so its keys are still current.
+        let cached = self.entries.get(&validator).copied()?;
+
+        let mut keys = Vec::new();
+        keys.push(cached.active_key);
+        if let Some(previous) = cached.previous_key {
+            if current_ledger.saturating_sub(cached.rotated_at_ledger) < self.window_ledgers {
+                keys.push(previous);
+            }
+        }
+        Some(keys)
+    }
+}
+
+/// Verify an attestation signature for `validator`, tolerating an in-flight key
+/// rotation. Accepts the signature if it validates under any key currently
+/// valid for the validator (the active key, plus the previous key while inside
+/// the rotation window). Refreshes the cache when the registry generation has
+/// advanced, so a rotation never leaves the cache serving a stale key.
+pub fn verify_with_rotation(
+    cache: &mut VerificationCache,
+    registry: &KeyRegistry,
+    validator: ValidatorId,
+    domain: &Domain,
+    data: &AttestationData,
+    signature: &Signature,
+    current_ledger: u64,
+) -> bool {
+    let keys = match cache.resolve_keys(registry, validator, current_ledger) {
+        Some(keys) => keys,
+        None => return false,
+    };
+    keys.iter()
+        .any(|key| verify_attestation_signature(key, domain, data, signature))
+}

src/attestation/mod.rs

@@ -2,4 +2,5 @@
 
 pub mod bitfield;
 pub mod bls_aggregator;
+pub mod key_registry;
 pub mod verifier;

tests/attestation_key_rotation_test.rs

@@ -0,0 +1,143 @@
+//! Key-rotation cache-poisoning tests (#5).
+//!
+//! Cargo only auto-discovers test targets at the top level of `tests/`, so this
+//! lives here rather than at `tests/attestation/`.
+
+use sorosusu_contracts::attestation::key_registry::{
+    verify_with_rotation, KeyRegistry, VerificationCache, ROTATION_WINDOW_LEDGERS,
+};
+use sorosusu_contracts::attestation::verifier::{
+    sign_attestation, verify_attestation_signature, AttestationData,
+};
+use sorosusu_contracts::crypto::domain::{compute_domain, DOMAIN_BEACON_ATTESTER, GENESIS_FORK_VERSION};
+
+use proptest::prelude::*;
+
+const OLD_KEY: [u8; 32] = [0x11; 32];
+const NEW_KEY: [u8; 32] = [0x22; 32];
+const VALIDATOR: u32 = 1;
+
+fn sample_data(slot: u64) -> AttestationData {
+    AttestationData {
+        slot,
+        index: 0,
+        beacon_block_root: [0xAA; 32],
+        source_epoch: 1,
+        source_root: [0xBB; 32],
+        target_epoch: 2,
+        target_root: [0xCC; 32],
+    }
+}
+
+/// Core regression for #5: after a rotation, a NEW-key attestation that would
+/// be rejected against the stale cached key is accepted because the versioned
+/// cache reloads; and an in-flight OLD-key attestation still verifies within
+/// the rotation window.
+#[test]
+fn rotation_window_accepts_old_and_new_keys() {
+    let domain = compute_domain(DOMAIN_BEACON_ATTESTER, GENESIS_FORK_VERSION);
+    let data = sample_data(42);
+
+    let mut registry = KeyRegistry::new();
+    let mut cache = VerificationCache::new();
+    registry.register(VALIDATOR, OLD_KEY);
+
+    let sig_old = sign_attestation(&OLD_KEY, &domain, &data);
+    let sig_new = sign_attestation(&NEW_KEY, &domain, &data);
+
+    // Before rotation: old key verifies and is cached at gen 0.
+    assert!(verify_with_rotation(&mut cache, &registry, VALIDATOR, &domain, &data, &sig_old, 100));
+    assert_eq!(cache.cached_key_gen(VALIDATOR), Some(0));
+
+    // Rotate to the new key at ledger 100.
+    assert_eq!(registry.rotate_key(VALIDATOR, NEW_KEY, 100), Some(1));
+
+    // The poisoning failure mode: verifying the new attestation against the old
+    // (stale) key fails.
+    assert!(!verify_attestation_signature(&OLD_KEY, &domain, &data, &sig_new));
+
+    // The fix: the versioned cache reloads (gen 0 -> 1) and the new attestation
+    // verifies.
+    assert!(verify_with_rotation(&mut cache, &registry, VALIDATOR, &domain, &data, &sig_new, 101));
+    assert_eq!(cache.cached_key_gen(VALIDATOR), Some(1));
+
+    // An in-flight old-key attestation still verifies within the window.
+    assert!(verify_with_rotation(&mut cache, &registry, VALIDATOR, &domain, &data, &sig_old, 102));
+}
+
+/// Once the rotation window elapses, the old key is no longer accepted but the
+/// new key still is.
+#[test]
+fn old_key_expires_after_window() {
+    let domain = compute_domain(DOMAIN_BEACON_ATTESTER, GENESIS_FORK_VERSION);
+    let data = sample_data(7);
+
+    let mut registry = KeyRegistry::new();
+    let mut cache = VerificationCache::new();
+    registry.register(VALIDATOR, OLD_KEY);
+    registry.rotate_key(VALIDATOR, NEW_KEY, 1_000);
+
+    let sig_old = sign_attestation(&OLD_KEY, &domain, &data);
+    let sig_new = sign_attestation(&NEW_KEY, &domain, &data);
+
+    let after_window = 1_000 + ROTATION_WINDOW_LEDGERS;
+    assert!(!verify_with_rotation(&mut cache, &registry, VALIDATOR, &domain, &data, &sig_old, after_window));
+    assert!(verify_with_rotation(&mut cache, &registry, VALIDATOR, &domain, &data, &sig_new, after_window));
+}
+
+/// Blueprint step 5: a rotation lands while a batch of verifications is in
+/// flight; old-key and new-key attestations alike all pass within the window.
+#[test]
+fn concurrent_rotation_and_verification_all_pass() {
+    let domain = compute_domain(DOMAIN_BEACON_ATTESTER, GENESIS_FORK_VERSION);
+
+    let mut registry = KeyRegistry::new();
+    let mut cache = VerificationCache::new();
+    registry.register(VALIDATOR, OLD_KEY);
+    registry.rotate_key(VALIDATOR, NEW_KEY, 200);
+
+    for i in 0..10u64 {
+        let ledger = 201 + i; // 201..210, all inside the [200, 210) window
+        let data = sample_data(i);
+        // Even slots are in-flight old-key attestations, odd are new-key.
+        let signer = if i % 2 == 0 { OLD_KEY } else { NEW_KEY };
+        let sig = sign_attestation(&signer, &domain, &data);
+        assert!(
+            verify_with_rotation(&mut cache, &registry, VALIDATOR, &domain, &data, &sig, ledger),
+            "verification {i} (ledger {ledger}) should pass"
+        );
+    }
+}
+
+/// Unknown validators never verify.
+#[test]
+fn unknown_validator_rejected() {
+    let domain = compute_domain(DOMAIN_BEACON_ATTESTER, GENESIS_FORK_VERSION);
+    let data = sample_data(1);
+    let mut cache = VerificationCache::new();
+    let registry = KeyRegistry::new();
+    let sig = sign_attestation(&OLD_KEY, &domain, &data);
+    assert!(!verify_with_rotation(&mut cache, &registry, 99, &domain, &data, &sig, 0));
+}
+
+proptest! {
+    /// For any ledger offset after a rotation: the new key always verifies; the
+    /// old key verifies iff still inside the rotation window.
+    #[test]
+    fn prop_window_acceptance(offset in 0_u64..30, use_new_key in any::<bool>()) {
+        let domain = compute_domain(DOMAIN_BEACON_ATTESTER, GENESIS_FORK_VERSION);
+        let data = sample_data(3);
+
+        let mut registry = KeyRegistry::new();
+        let mut cache = VerificationCache::new();
+        registry.register(VALIDATOR, OLD_KEY);
+        registry.rotate_key(VALIDATOR, NEW_KEY, 1_000);
+
+        let signer = if use_new_key { NEW_KEY } else { OLD_KEY };
+        let sig = sign_attestation(&signer, &domain, &data);
+        let got = verify_with_rotation(&mut cache, &registry, VALIDATOR, &domain, &data, &sig, 1_000 + offset);
+
+        let expected = use_new_key || offset < ROTATION_WINDOW_LEDGERS;
+        prop_assert_eq!(got, expected);
+    }
+}