From 58a34e00a11839a814f364ec4942abe6780a29e9 Mon Sep 17 00:00:00 2001
From: tersec <tersec@users.noreply.github.com>
Date: Fri, 25 Oct 2024 14:04:35 +0000
Subject: [PATCH] fix inconsistent aggregation bits len in Electra (#6679)

---
 beacon_chain/spec/validator.nim | 36 +++++++++++++++++++++++----------
 tests/test_attestation_pool.nim |  6 ++++++
 2 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/beacon_chain/spec/validator.nim b/beacon_chain/spec/validator.nim
index ed99f571c9..f9f5554a93 100644
--- a/beacon_chain/spec/validator.nim
+++ b/beacon_chain/spec/validator.nim
@@ -565,21 +565,32 @@ func get_committee_index_one*(bits: AttestationCommitteeBits): Opt[CommitteeInde
 
 proc compute_on_chain_aggregate*(
     network_aggregates: openArray[electra.Attestation]): Opt[electra.Attestation] =
-  # aggregates = sorted(network_aggregates, key=lambda a: get_committee_indices(a.committee_bits)[0])
-  let aggregates = network_aggregates.sortedByIt(it.committee_bits.get_committee_index_one().expect("just one"))
-
-  let data = aggregates[0].data
-
-  var agg: AggregateSignature
-  var committee_bits: AttestationCommitteeBits
+  let
+    aggregates = network_aggregates.sortedByIt(
+      it.committee_bits.get_committee_index_one().expect("just one"))
+    data = aggregates[0].data
 
-  var totalLen = 0
+  var
+    agg: AggregateSignature
+    committee_bits: AttestationCommitteeBits
+    prev_committee_index: Opt[CommitteeIndex]
+    totalLen = 0
   for i, a in aggregates:
+    let committee_index = ? get_committee_index_one(a.committee_bits)
+    if prev_committee_index.isNone:
+      prev_committee_index = Opt.some committee_index
+    elif committee_index.distinctBase <= prev_committee_index.get.distinctBase:
+      continue
+    prev_committee_index = Opt.some committee_index
+
     totalLen += a.aggregation_bits.len
 
-  var aggregation_bits = ElectraCommitteeValidatorsBits.init(totalLen)
-  var pos = 0
-  var prev_committee_index: Opt[CommitteeIndex]
+  prev_committee_index.reset()
+
+  var
+    aggregation_bits = ElectraCommitteeValidatorsBits.init(totalLen)
+    pos = 0
+    filledLen = 0
   for i, a in aggregates:
     let
       committee_index = ? get_committee_index_one(a.committee_bits)
@@ -594,6 +605,7 @@ proc compute_on_chain_aggregate*(
     for b in a.aggregation_bits:
       aggregation_bits[pos] = b
       pos += 1
+    filledLen += a.aggregation_bits.len
 
     let sig = ? a.signature.load() # Expensive
     if first:
@@ -603,6 +615,8 @@ proc compute_on_chain_aggregate*(
 
     committee_bits[int(committee_index)] = true
 
+  doAssert totalLen == filledLen
+
   let signature = agg.finish()
 
   ok electra.Attestation(
diff --git a/tests/test_attestation_pool.nim b/tests/test_attestation_pool.nim
index aac347f0a6..4fccbf0873 100644
--- a/tests/test_attestation_pool.nim
+++ b/tests/test_attestation_pool.nim
@@ -961,6 +961,8 @@ suite "Attestation pool electra processing" & preset():
 
     check:
       verifyAttestationSignature(attestations[0])
+      check_attestation(
+        state[].electraData.data, attestations[0], {}, cache, true).isOk
 
       # A single final chain aggregated attestation should be created
       # with same data, 2 committee bits and 3 aggregation bits
@@ -1048,6 +1050,8 @@ suite "Attestation pool electra processing" & preset():
       check:
         attestations.len() == 1
         attestations[0].aggregation_bits.countOnes() == 3
+        check_attestation(
+          state[].electraData.data, attestations[0], {}, cache, true).isOk
         verifyAttestationSignature(attestations[0])
         # Can get either aggregate here, random!
         verifyAttestationSignature(pool[].getElectraAggregatedAttestation(
@@ -1063,6 +1067,8 @@ suite "Attestation pool electra processing" & preset():
       check:
         attestations.len() == 1
         attestations[0].aggregation_bits.countOnes() == 4
+        check_attestation(
+          state[].electraData.data, attestations[0], {}, cache, true).isOk
         verifyAttestationSignature(attestations[0])
         verifyAttestationSignature(pool[].getElectraAggregatedAttestation(
           1.Slot, hash_tree_root(attestations[0].data), 0.CommitteeIndex).get)