From 3ac3ddb2b7d4e4a824593db0a57fc6f9e39ec710 Mon Sep 17 00:00:00 2001 From: Michael Sproul Date: Tue, 18 Jun 2024 00:23:02 +1000 Subject: [PATCH] Clean up Electra observed aggregates (#5929) * Use consistent key in observed_attestations * Remove unwraps from observed aggregates --- .../src/attestation_verification.rs | 15 ++-- beacon_node/beacon_chain/src/lib.rs | 2 +- .../beacon_chain/src/observed_aggregates.rs | 74 +++++++++++++------ .../tests/attestation_verification.rs | 6 +- 4 files changed, 61 insertions(+), 36 deletions(-) diff --git a/beacon_node/beacon_chain/src/attestation_verification.rs b/beacon_node/beacon_chain/src/attestation_verification.rs index d82685c48c0..f2e3565ae66 100644 --- a/beacon_node/beacon_chain/src/attestation_verification.rs +++ b/beacon_node/beacon_chain/src/attestation_verification.rs @@ -35,8 +35,10 @@ mod batch; use crate::{ - beacon_chain::VALIDATOR_PUBKEY_CACHE_LOCK_TIMEOUT, metrics, - observed_aggregates::ObserveOutcome, observed_attesters::Error as ObservedAttestersError, + beacon_chain::VALIDATOR_PUBKEY_CACHE_LOCK_TIMEOUT, + metrics, + observed_aggregates::{ObserveOutcome, ObservedAttestationKey}, + observed_attesters::Error as ObservedAttestersError, BeaconChain, BeaconChainError, BeaconChainTypes, }; use bls::verify_signature_sets; @@ -58,9 +60,8 @@ use state_processing::{ use std::borrow::Cow; use strum::AsRefStr; use tree_hash::TreeHash; -use tree_hash_derive::TreeHash; use types::{ - Attestation, AttestationData, AttestationRef, BeaconCommittee, BeaconStateError, + Attestation, AttestationRef, BeaconCommittee, BeaconStateError, BeaconStateError::NoCommitteeFound, ChainSpec, CommitteeIndex, Epoch, EthSpec, ForkName, Hash256, IndexedAttestation, SelectionProof, SignedAggregateAndProof, Slot, SubnetId, }; @@ -309,12 +310,6 @@ struct IndexedAggregatedAttestation<'a, T: BeaconChainTypes> { observed_attestation_key_root: Hash256, } -#[derive(TreeHash)] -pub struct ObservedAttestationKey { - pub committee_index: u64, - pub attestation_data: AttestationData, -} - /// Wraps a `Attestation` that has been verified up until the point that an `IndexedAttestation` can /// be derived. /// diff --git a/beacon_node/beacon_chain/src/lib.rs b/beacon_node/beacon_chain/src/lib.rs index f419429e090..466ab0b67e7 100644 --- a/beacon_node/beacon_chain/src/lib.rs +++ b/beacon_node/beacon_chain/src/lib.rs @@ -39,7 +39,7 @@ mod light_client_server_cache; pub mod metrics; pub mod migrate; mod naive_aggregation_pool; -mod observed_aggregates; +pub mod observed_aggregates; mod observed_attesters; mod observed_blob_sidecars; pub mod observed_block_producers; diff --git a/beacon_node/beacon_chain/src/observed_aggregates.rs b/beacon_node/beacon_chain/src/observed_aggregates.rs index 5e161a998b5..b02bff96a2d 100644 --- a/beacon_node/beacon_chain/src/observed_aggregates.rs +++ b/beacon_node/beacon_chain/src/observed_aggregates.rs @@ -6,11 +6,14 @@ use ssz_types::{BitList, BitVector}; use std::collections::HashMap; use std::marker::PhantomData; use tree_hash::TreeHash; +use tree_hash_derive::TreeHash; use types::consts::altair::{ SYNC_COMMITTEE_SUBNET_COUNT, TARGET_AGGREGATORS_PER_SYNC_SUBCOMMITTEE, }; use types::slot_data::SlotData; -use types::{Attestation, AttestationRef, EthSpec, Hash256, Slot, SyncCommitteeContribution}; +use types::{ + Attestation, AttestationData, AttestationRef, EthSpec, Hash256, Slot, SyncCommitteeContribution, +}; pub type ObservedSyncContributions = ObservedAggregates< SyncCommitteeContribution, @@ -20,6 +23,17 @@ pub type ObservedSyncContributions = ObservedAggregates< pub type ObservedAggregateAttestations = ObservedAggregates, E, BitList<::MaxValidatorsPerSlot>>; +/// Attestation data augmented with committee index +/// +/// This is hashed and used to key the map of observed aggregate attestations. This is important +/// post-Electra where the attestation data committee index is 0 and we want to avoid accidentally +/// comparing aggregation bits for *different* committees. +#[derive(TreeHash)] +pub struct ObservedAttestationKey { + pub committee_index: u64, + pub attestation_data: AttestationData, +} + /// A trait use to associate capacity constants with the type being stored in `ObservedAggregates`. pub trait Consts { /// The default capacity of items stored per slot, in a single `SlotHashSet`. @@ -92,11 +106,11 @@ pub trait SubsetItem { /// Returns the item that gets stored in `ObservedAggregates` for later subset /// comparison with incoming aggregates. - fn get_item(&self) -> Self::Item; + fn get_item(&self) -> Result; /// Returns a unique value that keys the object to the item that is being stored /// in `ObservedAggregates`. - fn root(&self) -> Hash256; + fn root(&self) -> Result; } impl<'a, E: EthSpec> SubsetItem for AttestationRef<'a, E> { @@ -126,19 +140,22 @@ impl<'a, E: EthSpec> SubsetItem for AttestationRef<'a, E> { } /// Returns the sync contribution aggregation bits. - fn get_item(&self) -> Self::Item { + fn get_item(&self) -> Result { match self { - Self::Base(att) => { - // TODO(electra) fix unwrap - att.extend_aggregation_bits().unwrap() - } - Self::Electra(att) => att.aggregation_bits.clone(), + Self::Base(att) => att + .extend_aggregation_bits() + .map_err(|_| Error::GetItemError), + Self::Electra(att) => Ok(att.aggregation_bits.clone()), } } - /// Returns the hash tree root of the attestation data. - fn root(&self) -> Hash256 { - self.data().tree_hash_root() + /// Returns the hash tree root of the attestation data augmented with the committee index. + fn root(&self) -> Result { + Ok(ObservedAttestationKey { + committee_index: self.committee_index().ok_or(Error::RootError)?, + attestation_data: self.data().clone(), + } + .tree_hash_root()) } } @@ -153,19 +170,19 @@ impl<'a, E: EthSpec> SubsetItem for &'a SyncCommitteeContribution { } /// Returns the sync contribution aggregation bits. - fn get_item(&self) -> Self::Item { - self.aggregation_bits.clone() + fn get_item(&self) -> Result { + Ok(self.aggregation_bits.clone()) } /// Returns the hash tree root of the root, slot and subcommittee index /// of the sync contribution. - fn root(&self) -> Hash256 { - SyncCommitteeData { + fn root(&self) -> Result { + Ok(SyncCommitteeData { root: self.beacon_block_root, slot: self.slot, subcommittee_index: self.subcommittee_index, } - .tree_hash_root() + .tree_hash_root()) } } @@ -192,6 +209,8 @@ pub enum Error { expected: Slot, attestation: Slot, }, + GetItemError, + RootError, } /// A `HashMap` that contains entries related to some `Slot`. @@ -234,7 +253,7 @@ impl SlotHashSet { // If true, we replace the new item with its existing subset. This allows us // to hold fewer items in the list. } else if item.is_superset(existing) { - *existing = item.get_item(); + *existing = item.get_item()?; return Ok(ObserveOutcome::New); } } @@ -252,7 +271,7 @@ impl SlotHashSet { return Err(Error::ReachedMaxObservationsPerSlot(self.max_capacity)); } - let item = item.get_item(); + let item = item.get_item()?; self.map.entry(root).or_default().push(item); Ok(ObserveOutcome::New) } @@ -345,7 +364,7 @@ where root_opt: Option, ) -> Result { let index = self.get_set_index(item.get_slot())?; - let root = root_opt.unwrap_or_else(|| item.root()); + let root = root_opt.map_or_else(|| item.root(), Ok)?; self.sets .get_mut(index) @@ -487,7 +506,10 @@ mod tests { for a in &items { assert_eq!( - store.is_known_subset(a.as_reference(), a.as_reference().root()), + store.is_known_subset( + a.as_reference(), + a.as_reference().root().unwrap() + ), Ok(false), "should indicate an unknown attestation is unknown" ); @@ -500,12 +522,18 @@ mod tests { for a in &items { assert_eq!( - store.is_known_subset(a.as_reference(), a.as_reference().root()), + store.is_known_subset( + a.as_reference(), + a.as_reference().root().unwrap() + ), Ok(true), "should indicate a known attestation is known" ); assert_eq!( - store.observe_item(a.as_reference(), Some(a.as_reference().root())), + store.observe_item( + a.as_reference(), + Some(a.as_reference().root().unwrap()) + ), Ok(ObserveOutcome::Subset), "should acknowledge an existing attestation" ); diff --git a/beacon_node/beacon_chain/tests/attestation_verification.rs b/beacon_node/beacon_chain/tests/attestation_verification.rs index 63740c4736c..3e35e58296a 100644 --- a/beacon_node/beacon_chain/tests/attestation_verification.rs +++ b/beacon_node/beacon_chain/tests/attestation_verification.rs @@ -2,8 +2,8 @@ use beacon_chain::attestation_verification::{ batch_verify_aggregated_attestations, batch_verify_unaggregated_attestations, Error, - ObservedAttestationKey, }; +use beacon_chain::observed_aggregates::ObservedAttestationKey; use beacon_chain::test_utils::{MakeAttestationOptions, HARNESS_GENESIS_TIME}; use beacon_chain::{ attestation_verification::Error as AttnError, @@ -852,7 +852,9 @@ async fn aggregated_gossip_verification() { err, AttnError::AttestationSupersetKnown(hash) if hash == ObservedAttestationKey { - committee_index: tester.valid_aggregate.message().aggregate().expect("should get committee index"), + committee_index: tester.valid_aggregate.message().aggregate() + .committee_index() + .expect("should get committee index"), attestation_data: tester.valid_aggregate.message().aggregate().data().clone(), }.tree_hash_root() ))