Clean up Electra observed aggregates (sigp#5929)

* Use consistent key in observed_attestations

* Remove unwraps from observed aggregates

beacon_node/beacon_chain/src/attestation_verification.rs

@@ -60,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,
 };
@@ -311,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.
 ///

beacon_node/beacon_chain/src/observed_aggregates.rs

@@ -11,7 +11,9 @@ 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<E> = ObservedAggregates<
     SyncCommitteeContribution<E>,
@@ -21,6 +23,17 @@ pub type ObservedSyncContributions<E> = ObservedAggregates<
 pub type ObservedAggregateAttestations<E> =
     ObservedAggregates<Attestation<E>, E, BitList<<E as types::EthSpec>::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`.
@@ -127,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<Self::Item, Error> {
         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<Hash256, Error> {
+        Ok(ObservedAttestationKey {
+            committee_index: self.committee_index().ok_or(Error::RootError)?,
+            attestation_data: self.data().clone(),
+        }
+        .tree_hash_root())
     }
 }
 
@@ -490,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"
                         );
@@ -503,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"
                         );

beacon_node/beacon_chain/tests/attestation_verification.rs

@@ -2,7 +2,6 @@
 
 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};
@@ -853,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()
                 ))