From 8473e3a7772fb043b46aa7e13d15da6b2f0c8ee7 Mon Sep 17 00:00:00 2001
From: dapplion <35266934+dapplion@users.noreply.github.com>
Date: Mon, 17 Jun 2024 16:17:04 +0200
Subject: [PATCH] Remove unwraps in Attestation construction

---
 consensus/types/src/attestation.rs          | 31 ++++++++++++++++++
 validator_client/src/attestation_service.rs | 36 ++++++++++-----------
 2 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/consensus/types/src/attestation.rs b/consensus/types/src/attestation.rs
index 3df14feadef..f4d674fd78c 100644
--- a/consensus/types/src/attestation.rs
+++ b/consensus/types/src/attestation.rs
@@ -1,4 +1,5 @@
 use crate::slot_data::SlotData;
+use crate::ForkName;
 use crate::{test_utils::TestRandom, Hash256, Slot};
 use derivative::Derivative;
 use rand::RngCore;
@@ -22,6 +23,8 @@ pub enum Error {
     AlreadySigned(usize),
     SubnetCountIsZero(ArithError),
     IncorrectStateVariant,
+    InvalidCommitteeLength,
+    InvalidCommitteeIndex,
 }
 
 #[superstruct(
@@ -104,6 +107,34 @@ impl<E: EthSpec> Hash for Attestation<E> {
 }
 
 impl<E: EthSpec> Attestation<E> {
+    pub fn empty_for_signing(
+        committee_index: usize,
+        committee_length: usize,
+        data: AttestationData,
+        spec: &ChainSpec,
+    ) -> Result<Self, Error> {
+        if spec.fork_name_at_slot::<E>(data.slot) >= ForkName::Electra {
+            let mut committee_bits: BitVector<E::MaxCommitteesPerSlot> = BitVector::default();
+            committee_bits
+                .set(committee_index, true)
+                .map_err(|_| Error::InvalidCommitteeIndex)?;
+            Ok(Attestation::Electra(AttestationElectra {
+                aggregation_bits: BitList::with_capacity(committee_length)
+                    .map_err(|_| Error::InvalidCommitteeLength)?,
+                data,
+                committee_bits,
+                signature: AggregateSignature::infinity(),
+            }))
+        } else {
+            Ok(Attestation::Base(AttestationBase {
+                aggregation_bits: BitList::with_capacity(committee_length)
+                    .map_err(|_| Error::InvalidCommitteeLength)?,
+                data,
+                signature: AggregateSignature::infinity(),
+            }))
+        }
+    }
+
     /// Aggregate another Attestation into this one.
     ///
     /// The aggregation bitfields must be disjoint, and the data must be the same.
diff --git a/validator_client/src/attestation_service.rs b/validator_client/src/attestation_service.rs
index 404fe41249b..a3ab7a4c726 100644
--- a/validator_client/src/attestation_service.rs
+++ b/validator_client/src/attestation_service.rs
@@ -386,25 +386,23 @@ impl<T: SlotClock + 'static, E: EthSpec> AttestationService<T, E> {
                 return None;
             }
 
-            let mut attestation = if fork_name >= ForkName::Electra {
-                let mut committee_bits: BitVector<E::MaxCommitteesPerSlot> = BitVector::default();
-                committee_bits
-                    .set(duty.committee_index as usize, true)
-                    .unwrap();
-                Attestation::Electra(AttestationElectra {
-                    aggregation_bits: BitList::with_capacity(duty.committee_length as usize)
-                        .unwrap(),
-                    data: attestation_data.clone(),
-                    committee_bits,
-                    signature: AggregateSignature::infinity(),
-                })
-            } else {
-                Attestation::Base(AttestationBase {
-                    aggregation_bits: BitList::with_capacity(duty.committee_length as usize)
-                        .unwrap(),
-                    data: attestation_data.clone(),
-                    signature: AggregateSignature::infinity(),
-                })
+            let mut attestation = match Attestation::<E>::empty_for_signing(
+                duty.committee_index as usize,
+                duty.committee_length as usize,
+                attestation_data.clone(),
+                &self.context.eth2_config.spec,
+            ) {
+                Ok(attestation) => attestation,
+                Err(err) => {
+                    crit!(
+                        log,
+                        "Invalid validator duties during signing";
+                        "validator" => ?duty.pubkey,
+                        "duty" => ?duty,
+                        "err" => ?err,
+                    );
+                    return None;
+                }
             };
 
             match self