fork choice suite fixes

lean_client/fork_choice/src/handlers.rs

@@ -180,6 +180,11 @@ fn find_unknown_attestation_block(
     .find(|root| !store.blocks.contains_key(root))
 }
 
+pub enum AttestationOutcome {
+    Applied,
+    Queued(H256),
+}
+
 /// Process a signed attestation received via gossip network
 ///
 /// 1. Validates the attestation data
@@ -190,7 +195,7 @@ fn find_unknown_attestation_block(
 pub fn on_gossip_attestation(
     store: &mut Store,
     signed_attestation: SignedAttestation,
-) -> Result<()> {
+) -> Result<AttestationOutcome> {
     let _timer = METRICS.get().map(|metrics| {
         metrics
             .lean_attestation_validation_time_seconds
@@ -213,7 +218,7 @@ pub fn on_gossip_attestation(
             m.grandine_pending_fetch_roots
                 .set(store.pending_fetch_roots.len() as i64)
         });
-        return Ok(());
+        return Ok(AttestationOutcome::Queued(missing_root));
     }
 
     // Validate the attestation data first
@@ -226,12 +231,19 @@ pub fn on_gossip_attestation(
         });
     })?;
 
-    // Non-aggregators validate attestation data but do not store or verify individual
-    // signatures. Per leanSpec: only aggregators import gossip attestations for aggregation.
-    // Subnet filtering is already enforced at the p2p subscription layer.
-    if !store.is_aggregator {
-        return Ok(());
-    }
+    // Reject validators outside the registry (independent of signature checks).
+    let num_validators = store
+        .states
+        .get(&attestation_data.target.root)
+        .ok_or_else(|| anyhow!("no state for target block {}", attestation_data.target.root))?
+        .validators
+        .len_u64();
+    ensure!(
+        validator_id < num_validators,
+        "validator {} out of range (max {})",
+        validator_id,
+        num_validators
+    );
 
     let data_root = attestation_data.hash_tree_root();
     let sig_key = SignatureKey::new(signed_attestation.validator_id, data_root);
@@ -240,7 +252,13 @@ pub fn on_gossip_attestation(
     // Duplicate attestations arrive when the IDontWant buffer fills under CPU load,
     // causing peers to rebroadcast. Each verify costs ~100ms; the early exit breaks
     // the saturation loop without dropping vote data.
-    if !store.gossip_signatures.contains_key(&sig_key) {
+    let already_verified = if store.is_aggregator {
+        store.gossip_signatures.contains_key(&sig_key)
+    } else {
+        store.verified_gossip_signatures.contains(&sig_key)
+    };
+
+    if !already_verified {
         // Verify individual XMSS signature against the validator's public key.
         // State is available: the pending-block check above confirmed target.root is in
         // the store, and states are stored 1:1 with blocks in process_block_internal.
@@ -267,22 +285,30 @@ pub fn on_gossip_attestation(
             .verify(&pubkey, attestation_data.slot.0 as u32, data_root)
             .context("individual attestation signature verification failed")?;
 
-        store
-            .gossip_signatures
-            .insert(sig_key, signed_attestation.signature);
+        if store.is_aggregator {
+            store
+                .gossip_signatures
+                .insert(sig_key, signed_attestation.signature);
 
-        // Update gossip signatures gauge
-        METRICS.get().map(|metrics| {
-            metrics
-                .lean_gossip_signatures
-                .set(store.gossip_signatures.len() as i64);
-        });
+            // Update gossip signatures gauge
+            METRICS.get().map(|metrics| {
+                metrics
+                    .lean_gossip_signatures
+                    .set(store.gossip_signatures.len() as i64);
+            });
+        } else {
+            store.verified_gossip_signatures.insert(sig_key);
+        }
     } else {
         METRICS
             .get()
             .map(|m| m.grandine_xmss_verify_skipped_total.inc());
     }
 
+    if !store.is_aggregator {
+        return Ok(AttestationOutcome::Applied);
+    }
+
     store
         .attestation_data_by_root
         .insert(data_root, attestation_data.clone());
@@ -309,6 +335,7 @@ pub fn on_gossip_attestation(
                     .inc()
             });
         })
+        .map(|()| AttestationOutcome::Applied)
 }
 
 /// Process an attestation and place it into the correct attestation stage
@@ -413,7 +440,7 @@ pub fn on_attestation(
 pub fn on_aggregated_attestation(
     store: &mut Store,
     signed_aggregated_attestation: SignedAggregatedAttestation,
-) -> Result<()> {
+) -> Result<AttestationOutcome> {
     // Structure: { data: AttestationData, proof: AggregatedSignatureProof }
     let attestation_data = signed_aggregated_attestation.data.clone();
     let proof = signed_aggregated_attestation.proof.clone();
@@ -430,7 +457,7 @@ pub fn on_aggregated_attestation(
             m.grandine_pending_fetch_roots
                 .set(store.pending_fetch_roots.len() as i64)
         });
-        return Ok(());
+        return Ok(AttestationOutcome::Queued(missing_root));
     }
 
     // Validate attestation data (slot bounds, target validity, etc.)
@@ -507,7 +534,7 @@ pub fn on_aggregated_attestation(
             .set(store.latest_new_aggregated_payloads.len() as i64);
     });
 
-    Ok(())
+    Ok(AttestationOutcome::Applied)
 }
 
 /// Three-phase variant of `on_aggregated_attestation` that releases the store
@@ -892,6 +919,10 @@ pub fn apply_verified_block(
             adr.get(&key.data_root)
                 .map_or(true, |data| data.target.slot.0 > finalized_slot)
         });
+        store.verified_gossip_signatures.retain(|key| {
+            adr.get(&key.data_root)
+                .map_or(true, |data| data.target.slot.0 > finalized_slot)
+        });
         store
             .latest_known_aggregated_payloads
             .retain(|data_root, _| {
@@ -978,6 +1009,10 @@ fn prune_with_retention_bounds(store: &mut Store) {
         adr.get(&key.data_root)
             .is_none_or(|data| data.target.slot.0 >= keep_min_slot)
     });
+    store.verified_gossip_signatures.retain(|key| {
+        adr.get(&key.data_root)
+            .is_none_or(|data| data.target.slot.0 >= keep_min_slot)
+    });
     store
         .latest_known_aggregated_payloads
         .retain(|data_root, _| {

lean_client/fork_choice/src/store.rs

@@ -79,6 +79,8 @@ pub struct Store {
 
     pub gossip_signatures: HashMap<SignatureKey, Signature>,
 
+    pub verified_gossip_signatures: HashSet<SignatureKey>,
+
     /// Aggregated signature proofs from block bodies (on-chain).
     /// These are attestations that have been included in blocks and are part of
     /// the "known" pool for safe target computation.
@@ -288,6 +290,7 @@ pub fn get_forkchoice_store(
         latest_known_attestations: HashMap::new(),
         latest_new_attestations: HashMap::new(),
         gossip_signatures: HashMap::new(),
+        verified_gossip_signatures: HashSet::new(),
         latest_known_aggregated_payloads: IndexMap::new(),
         latest_new_aggregated_payloads: IndexMap::new(),
         attestation_data_by_root: HashMap::new(),
@@ -395,12 +398,20 @@ pub fn get_latest_justified(states: &HashMap<H256, Arc<State>>) -> Option<&Check
 pub fn update_head(store: &mut Store) {
     let old_head = store.head;
 
-    let latest_votes = extract_attestations_from_aggregated_payloads(
+    let mut latest_votes = extract_attestations_from_aggregated_payloads(
         &store.latest_known_aggregated_payloads,
         &store.attestation_data_by_root,
         store.latest_finalized.slot,
     );
 
+    for (validator_id, data) in &store.latest_known_attestations {
+        if data.head.slot > store.latest_finalized.slot {
+            latest_votes
+                .entry(*validator_id)
+                .or_insert_with(|| data.clone());
+        }
+    }
+
     // Compute new head using LMD-GHOST from latest justified root
     let new_head = get_fork_choice_head(store, store.latest_justified.root, &latest_votes, 0);
     store.head = new_head;

lean_client/http_api/src/test_driver.rs

@@ -11,6 +11,7 @@
 //! The wire shapes here are dictated by the hive simulator at
 //! `simulators/lean/src/scenarios/spec_assets.rs`.
 
+use std::str::FromStr;
 use std::sync::Arc;
 
 use axum::{
@@ -22,7 +23,9 @@ use containers::{
 };
 use fork_choice::{
     block_cache::BlockCache,
-    handlers::{on_aggregated_attestation, on_block, on_gossip_attestation, on_tick},
+    handlers::{
+        AttestationOutcome, on_aggregated_attestation, on_block, on_gossip_attestation, on_tick,
+    },
     store::{MILLIS_PER_INTERVAL, SECONDS_PER_SLOT, Store, get_forkchoice_store},
 };
 use parking_lot::RwLock;
@@ -398,13 +401,23 @@ fn apply_step(
                 .map_err(|err| err.to_string())
         }
         ForkChoiceStep::Attestation { attestation, .. } => {
+            let signature = match attestation.signature.as_deref() {
+                Some(hex) => Signature::from_str(hex)
+                    .map_err(|err| format!("invalid fixture attestation signature: {err}"))?,
+                None => Signature::default(),
+            };
             let attestation: containers::Attestation = attestation.into();
             let signed = SignedAttestation {
                 validator_id: attestation.validator_id,
                 message: attestation.data,
-                signature: Signature::default(),
+                signature,
             };
-            on_gossip_attestation(store, signed).map_err(|err| err.to_string())
+            match on_gossip_attestation(store, signed).map_err(|err| err.to_string())? {
+                AttestationOutcome::Applied => Ok(()),
+                AttestationOutcome::Queued(root) => {
+                    Err(format!("attestation references unknown block {root:?}"))
+                }
+            }
         }
         ForkChoiceStep::GossipAggregatedAttestation { attestation, .. } => {
             let Some(step) = attestation else {
@@ -413,7 +426,12 @@ fn apply_step(
                 return Ok(());
             };
             let signed = build_signed_aggregated_attestation(step)?;
-            on_aggregated_attestation(store, signed).map_err(|err| err.to_string())
+            match on_aggregated_attestation(store, signed).map_err(|err| err.to_string())? {
+                AttestationOutcome::Applied => Ok(()),
+                AttestationOutcome::Queued(root) => {
+                    Err(format!("attestation references unknown block {root:?}"))
+                }
+            }
         }
         ForkChoiceStep::Checks { .. } => {
             // Pure-assertion step. The simulator validates against the

lean_client/spec_test_fixtures/src/common.rs

@@ -99,6 +99,8 @@ pub fn parse_root(hex_str: &str) -> H256 {
 pub struct TestAttestation {
     pub validator_index: u64,
     pub data: TestAttestationData,
+    #[serde(default)]
+    pub signature: Option<String>,
 }
 
 impl From<TestAttestation> for Attestation {

lean_client/src/main.rs

@@ -1813,6 +1813,7 @@ async fn main() -> Result<()> {
                                 on_attestation(&mut *store.write(), signed_attestation.clone(), false)
                             } else {
                                 on_gossip_attestation(&mut *store.write(), signed_attestation.clone())
+                                        .map(|_| ())
                             };
                             match result {
                                 Ok(()) => {