From d6e20ce367643b7b0c69647747863e5a9f1e8e4c Mon Sep 17 00:00:00 2001
From: Eitan Seri-Levi <eserilev@ucsc.edu>
Date: Wed, 15 May 2024 08:37:20 +0300
Subject: [PATCH] Subscribe to the correct subnets for electra attestations
 (#5782)

* subscribe to the correct att subnets for electra

* subscribe to the correct att subnets for electra
---
 beacon_node/beacon_chain/src/attestation_verification.rs   | 6 +++---
 beacon_node/beacon_chain/src/test_utils.rs                 | 4 ++--
 beacon_node/beacon_chain/tests/attestation_verification.rs | 4 ++--
 consensus/types/src/subnet_id.rs                           | 5 ++---
 4 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/beacon_node/beacon_chain/src/attestation_verification.rs b/beacon_node/beacon_chain/src/attestation_verification.rs
index f80d74aec44..f37f2ebfee4 100644
--- a/beacon_node/beacon_chain/src/attestation_verification.rs
+++ b/beacon_node/beacon_chain/src/attestation_verification.rs
@@ -824,9 +824,9 @@ impl<'a, T: BeaconChainTypes> IndexedUnaggregatedAttestation<'a, T> {
         committees_per_slot: u64,
         subnet_id: Option<SubnetId>,
         chain: &BeaconChain<T>,
-    ) -> Result<(u64, SubnetId), Error> {
-        let expected_subnet_id = SubnetId::compute_subnet_for_attestation_data::<T::EthSpec>(
-            indexed_attestation.data(),
+    ) -> Result<(u64, SubnetId), Error> { 
+        let expected_subnet_id = SubnetId::compute_subnet_for_attestation::<T::EthSpec>(
+            &attestation,
             committees_per_slot,
             &chain.spec,
         )
diff --git a/beacon_node/beacon_chain/src/test_utils.rs b/beacon_node/beacon_chain/src/test_utils.rs
index 7122d1837bf..db4c65f0caa 100644
--- a/beacon_node/beacon_chain/src/test_utils.rs
+++ b/beacon_node/beacon_chain/src/test_utils.rs
@@ -1169,8 +1169,8 @@ where
                             agg_sig
                         };
 
-                        let subnet_id = SubnetId::compute_subnet_for_attestation_data::<E>(
-                            attestation.data(),
+                        let subnet_id = SubnetId::compute_subnet_for_attestation::<E>(
+                            &attestation.to_ref(),
                             committee_count,
                             &self.chain.spec,
                         )
diff --git a/beacon_node/beacon_chain/tests/attestation_verification.rs b/beacon_node/beacon_chain/tests/attestation_verification.rs
index 3d5b57a4183..2c5ee0388bd 100644
--- a/beacon_node/beacon_chain/tests/attestation_verification.rs
+++ b/beacon_node/beacon_chain/tests/attestation_verification.rs
@@ -147,8 +147,8 @@ fn get_valid_unaggregated_attestation<T: BeaconChainTypes>(
         )
         .expect("should sign attestation");
 
-    let subnet_id = SubnetId::compute_subnet_for_attestation_data::<E>(
-        valid_attestation.data(),
+    let subnet_id = SubnetId::compute_subnet_for_attestation::<E>(
+        &valid_attestation.to_ref(),
         head.beacon_state
             .get_committee_count_at_slot(current_slot)
             .expect("should get committee count"),
diff --git a/consensus/types/src/subnet_id.rs b/consensus/types/src/subnet_id.rs
index 66786b51297..18f2b838f3b 100644
--- a/consensus/types/src/subnet_id.rs
+++ b/consensus/types/src/subnet_id.rs
@@ -1,6 +1,5 @@
 //! Identifies each shard by an integer identifier.
 use crate::{AttestationRef, ChainSpec, CommitteeIndex, Epoch, EthSpec, Slot};
-use lazy_static::lazy_static;
 use safe_arith::{ArithError, SafeArith};
 use serde::{Deserialize, Serialize};
 use std::ops::{Deref, DerefMut};
@@ -41,7 +40,7 @@ impl SubnetId {
     /// Compute the subnet for an attestation where each slot in the
     /// attestation epoch contains `committee_count_per_slot` committees.
     pub fn compute_subnet_for_attestation<E: EthSpec>(
-        attestation: AttestationRef<E>,
+        attestation: &AttestationRef<E>,
         committee_count_per_slot: u64,
         spec: &ChainSpec,
     ) -> Result<SubnetId, ArithError> {
@@ -49,7 +48,7 @@ impl SubnetId {
 
         Self::compute_subnet::<E>(
             attestation.data().slot,
-            committee_index,
+            attestation.committee_index(),
             committee_count_per_slot,
             spec,
         )