feat: add signature aggregation using python bindings (#238)

* feat: add signature aggregation using python bindings

* add highlevel modified flow for attestation import and proposal packing

* feat: update containers to handle attestation instead of signedAttestations

* fix: ci, add: test aggregation and verification

* fix: update store mappings

* fix: address comments

* fix: address review comments from codex

* rename: Signature to XmssSignature

* fix: update aggregation API signature

* update: block building logic with greedy signature aggregation

* fix: make names intuitive and add type alias

* add: extensive testing

* refactor: aggregated sig container, address other comments

* fix: simplify block building

---------

Co-authored-by: harkamal <[email protected]>

Author: anshalshukla

packages/testing/src/consensus_testing/keys.py

@@ -39,12 +39,13 @@
 
 from lean_spec.config import LEAN_ENV
 from lean_spec.subspecs.containers import AttestationData
-from lean_spec.subspecs.containers.attestation.types import NaiveAggregatedSignature
 from lean_spec.subspecs.containers.block.types import (
     AggregatedAttestations,
     AttestationSignatures,
 )
 from lean_spec.subspecs.containers.slot import Slot
+from lean_spec.subspecs.containers.state.types import AttestationSignatureKey
+from lean_spec.subspecs.xmss.aggregation import MultisigAggregatedSignature
 from lean_spec.subspecs.xmss.containers import KeyPair, PublicKey, Signature
 from lean_spec.subspecs.xmss.interface import (
     PROD_SIGNATURE_SCHEME,
@@ -273,24 +274,39 @@ def sign_attestation_data(
     def build_attestation_signatures(
         self,
         aggregated_attestations: AggregatedAttestations,
-        signature_lookup: Mapping[tuple[Uint64, bytes], Signature] | None = None,
+        signature_lookup: Mapping[AttestationSignatureKey, Signature] | None = None,
     ) -> AttestationSignatures:
-        """Build `AttestationSignatures` for already-aggregated attestations."""
+        """
+        Build `AttestationSignatures` for already-aggregated attestations.
+
+        For each aggregated attestation, collect the participating validators' public keys and
+        signatures, then produce a single leanVM aggregated signature proof.
+        """
         lookup = signature_lookup or {}
-        return AttestationSignatures(
-            data=[
-                NaiveAggregatedSignature(
-                    data=[
-                        (
-                            lookup.get((vid, agg.data.data_root_bytes()))
-                            or self.sign_attestation_data(vid, agg.data)
-                        )
-                        for vid in agg.aggregation_bits.to_validator_indices()
-                    ]
-                )
-                for agg in aggregated_attestations
+
+        proof_blobs: list[MultisigAggregatedSignature] = []
+        for agg in aggregated_attestations:
+            validator_ids = agg.aggregation_bits.to_validator_indices()
+            message = agg.data.data_root_bytes()
+            epoch = agg.data.slot
+
+            public_keys: list[PublicKey] = [self.get_public_key(vid) for vid in validator_ids]
+            signatures: list[Signature] = [
+                (lookup.get((vid, message)) or self.sign_attestation_data(vid, agg.data))
+                for vid in validator_ids
             ]
-        )
+
+            # If the caller supplied raw signatures and any are invalid,
+            # aggregation should fail with exception.
+            aggregated_signature = MultisigAggregatedSignature.aggregate_signatures(
+                public_keys=public_keys,
+                signatures=signatures,
+                message=message,
+                epoch=epoch,
+            )
+            proof_blobs.append(aggregated_signature)
+
+        return AttestationSignatures(data=proof_blobs)
 
 
 def _generate_single_keypair(

packages/testing/src/consensus_testing/test_fixtures/fork_choice.py

@@ -26,10 +26,10 @@
 from lean_spec.subspecs.containers.slot import Slot
 from lean_spec.subspecs.containers.state import Validators
 from lean_spec.subspecs.containers.state.state import State
+from lean_spec.subspecs.containers.state.types import AttestationSignatureKey
 from lean_spec.subspecs.forkchoice import Store
 from lean_spec.subspecs.koalabear import Fp
 from lean_spec.subspecs.ssz import hash_tree_root
-from lean_spec.subspecs.xmss.constants import PROD_CONFIG
 from lean_spec.subspecs.xmss.containers import Signature
 from lean_spec.subspecs.xmss.types import HashDigestList, HashTreeOpening, Randomness
 from lean_spec.types import Bytes32, Uint64
@@ -231,7 +231,10 @@ def make_fixture(self) -> ForkChoiceTest:
 
                 elif isinstance(step, AttestationStep):
                     # Process attestation from gossip (immutable)
-                    store = store.on_attestation(step.attestation, is_from_block=False)
+                    store = store.on_gossip_attestation(
+                        step.attestation,
+                        scheme=LEAN_ENV_TO_SCHEMES[self.lean_env],
+                    )
 
                 else:
                     raise ValueError(f"Step {i}: unknown step type {type(step).__name__}")
@@ -301,7 +304,12 @@ def _build_block_from_spec(
         parent_root = self._resolve_parent_root(spec, store, block_registry)
 
         # Build attestations from spec
-        attestations = self._build_attestations_from_spec(spec, store, block_registry, parent_root)
+        attestations, attestation_signatures = self._build_attestations_from_spec(
+            spec, store, block_registry, parent_root, key_manager
+        )
+
+        gossip_signatures = dict(store.gossip_signatures)
+        gossip_signatures.update(attestation_signatures)
 
         # Use State.build_block for core block building (pure spec logic)
         parent_state = store.states[parent_root]
@@ -310,6 +318,8 @@ def _build_block_from_spec(
             proposer_index=proposer_index,
             parent_root=parent_root,
             attestations=attestations,
+            gossip_signatures=gossip_signatures,
+            aggregated_payloads=store.aggregated_payloads,
         )
 
         # Create proposer attestation for this block
@@ -325,8 +335,9 @@ def _build_block_from_spec(
         )
 
         # Sign all attestations and the proposer attestation
-        attestation_signatures = key_manager.build_attestation_signatures(
-            final_block.body.attestations
+        attestation_signatures_blob = key_manager.build_attestation_signatures(
+            final_block.body.attestations,
+            attestation_signatures,
         )
 
         proposer_signature = key_manager.sign_attestation_data(
@@ -340,7 +351,7 @@ def _build_block_from_spec(
                 proposer_attestation=proposer_attestation,
             ),
             signature=BlockSignatures(
-                attestation_signatures=attestation_signatures,
+                attestation_signatures=attestation_signatures_blob,
                 proposer_signature=proposer_signature,
             ),
         )
@@ -392,34 +403,38 @@ def _build_attestations_from_spec(
         store: Store,
         block_registry: dict[str, Block],
         parent_root: Bytes32,
-    ) -> list[Attestation]:
-        """Build attestations list from BlockSpec."""
+        key_manager: XmssKeyManager,
+    ) -> tuple[list[Attestation], dict[AttestationSignatureKey, Signature]]:
+        """Build attestations list from BlockSpec and their signatures."""
         if spec.attestations is None:
-            return []
+            return [], {}
 
         parent_state = store.states[parent_root]
         attestations = []
+        signature_lookup: dict[AttestationSignatureKey, Signature] = {}
 
         for att_spec in spec.attestations:
             if isinstance(att_spec, SignedAttestationSpec):
                 signed_att = self._build_signed_attestation_from_spec(
-                    att_spec, block_registry, parent_state
-                )
-                attestations.append(
-                    Attestation(validator_id=signed_att.validator_id, data=signed_att.message)
+                    att_spec, block_registry, parent_state, key_manager
                 )
             else:
-                attestations.append(
-                    Attestation(validator_id=att_spec.validator_id, data=att_spec.message)
-                )
+                signed_att = att_spec
+
+            attestation = Attestation(validator_id=signed_att.validator_id, data=signed_att.message)
+            attestations.append(attestation)
+            signature_lookup[(attestation.validator_id, attestation.data.data_root_bytes())] = (
+                signed_att.signature
+            )
 
-        return attestations
+        return attestations, signature_lookup
 
     def _build_signed_attestation_from_spec(
         self,
         spec: SignedAttestationSpec,
         block_registry: dict[str, Block],
         state: State,
+        key_manager: XmssKeyManager,
     ) -> SignedAttestation:
         """
         Build a SignedAttestation from a SignedAttestationSpec.
@@ -466,15 +481,21 @@ def _build_signed_attestation_from_spec(
         )
 
         # Create signed attestation
+        if spec.signature is not None:
+            signature = spec.signature
+        elif spec.valid_signature:
+            signature = key_manager.sign_attestation_data(
+                attestation.validator_id, attestation.data
+            )
+        else:
+            signature = Signature(
+                path=HashTreeOpening(siblings=HashDigestList(data=[])),
+                rho=Randomness(data=[Fp(0) for _ in range(Randomness.LENGTH)]),
+                hashes=HashDigestList(data=[]),
+            )
+
         return SignedAttestation(
             validator_id=attestation.validator_id,
             message=attestation.data,
-            signature=(
-                spec.signature
-                or Signature(
-                    path=HashTreeOpening(siblings=HashDigestList(data=[])),
-                    rho=Randomness(data=[Fp(0) for _ in range(PROD_CONFIG.RAND_LEN_FE)]),
-                    hashes=HashDigestList(data=[]),
-                )
-            ),
+            signature=signature,
         )

packages/testing/src/consensus_testing/test_fixtures/state_transition.py

@@ -11,6 +11,7 @@
 from lean_spec.subspecs.ssz.hash import hash_tree_root
 from lean_spec.types import Bytes32, Uint64
 
+from ..keys import get_shared_key_manager
 from ..test_types import BlockSpec, StateExpectation
 from .base import BaseConsensusFixture
 
@@ -258,10 +259,24 @@ def _build_block_from_spec(self, spec: BlockSpec, state: State) -> tuple[Block,
             for vid in agg.aggregation_bits.to_validator_indices()
         ]
 
+        if plain_attestations:
+            key_manager = get_shared_key_manager(max_slot=spec.slot)
+            gossip_signatures = {
+                (att.validator_id, att.data.data_root_bytes()): key_manager.sign_attestation_data(
+                    att.validator_id,
+                    att.data,
+                )
+                for att in plain_attestations
+            }
+        else:
+            gossip_signatures = {}
+
         block, post_state, _, _ = state.build_block(
             slot=spec.slot,
             proposer_index=proposer_index,
             parent_root=parent_root,
             attestations=plain_attestations,
+            gossip_signatures=gossip_signatures,
+            aggregated_payloads={},
         )
         return block, post_state

packages/testing/src/consensus_testing/test_fixtures/verify_signatures.py

@@ -16,10 +16,14 @@
     BlockWithAttestation,
     SignedBlockWithAttestation,
 )
+from lean_spec.subspecs.containers.block.types import AttestationSignatures
 from lean_spec.subspecs.containers.checkpoint import Checkpoint
 from lean_spec.subspecs.containers.state.state import State
 from lean_spec.subspecs.koalabear import Fp
 from lean_spec.subspecs.ssz import hash_tree_root
+from lean_spec.subspecs.xmss.constants import TARGET_CONFIG
+from lean_spec.subspecs.xmss.containers import Signature
+from lean_spec.subspecs.xmss.types import HashDigestList, HashTreeOpening, Randomness
 from lean_spec.types import Bytes32, Uint64
 
 from ..keys import XmssKeyManager, get_shared_key_manager
@@ -179,26 +183,26 @@ def _build_block_from_spec(
             spec, state, key_manager
         )
 
+        # Provide signatures to State.build_block so it can include attestations during
+        # fixed-point collection when available_attestations/known_block_roots are used.
+        # This might contain invalid signatures as we are not validating them here.
+        gossip_signatures = {
+            (att.validator_id, att.data.data_root_bytes()): sig
+            for att, sig in zip(attestations, attestation_signature_inputs, strict=True)
+        }
+
         # Use State.build_block for core block building (pure spec logic)
-        final_block, _, _, _ = state.build_block(
+        final_block, _, _, aggregated_signatures = state.build_block(
             slot=spec.slot,
             proposer_index=proposer_index,
             parent_root=parent_root,
             attestations=attestations,
+            gossip_signatures=gossip_signatures,
+            aggregated_payloads={},
         )
 
-        # Preserve per-attestation validity from the spec.
-        #
-        # For signature tests we must ensure that the signatures in the input spec are used
-        # for any intentionally-invalid signature from the input spec remains invalid
-        # in the produced `SignedBlockWithAttestation`.
-        signature_lookup: dict[tuple[Uint64, bytes], Signature] = {
-            (att.validator_id, att.data.data_root_bytes()): sig
-            for att, sig in zip(attestations, attestation_signature_inputs, strict=True)
-        }
-        attestation_signatures = key_manager.build_attestation_signatures(
-            final_block.body.attestations,
-            signature_lookup=signature_lookup,
+        attestation_signatures = AttestationSignatures(
+            data=aggregated_signatures,
         )
 
         # Create proposer attestation for this block
@@ -220,14 +224,10 @@ def _build_block_from_spec(
                 proposer_attestation.data,
             )
         else:
-            # Generate an invalid dummy signature (all zeros)
-            from lean_spec.subspecs.xmss.constants import TEST_CONFIG
-            from lean_spec.subspecs.xmss.containers import Signature
-            from lean_spec.subspecs.xmss.types import HashDigestList, HashTreeOpening, Randomness
-
+            # Generate a structurally valid but cryptographically invalid signature (all zeros).
             proposer_attestation_signature = Signature(
                 path=HashTreeOpening(siblings=HashDigestList(data=[])),
-                rho=Randomness(data=[Fp(0) for _ in range(TEST_CONFIG.RAND_LEN_FE)]),
+                rho=Randomness(data=[Fp(0) for _ in range(TARGET_CONFIG.RAND_LEN_FE)]),
                 hashes=HashDigestList(data=[]),
             )
 
@@ -333,14 +333,10 @@ def _build_signed_attestation_from_spec(
                 attestation.data,
             )
         else:
-            # Generate an invalid dummy signature (all zeros)
-            from lean_spec.subspecs.xmss.constants import TEST_CONFIG
-            from lean_spec.subspecs.xmss.containers import Signature
-            from lean_spec.subspecs.xmss.types import HashDigestList, HashTreeOpening, Randomness
-
+            # Generate a structurally valid but cryptographically invalid signature (all zeros).
             signature = Signature(
                 path=HashTreeOpening(siblings=HashDigestList(data=[])),
-                rho=Randomness(data=[Fp(0) for _ in range(TEST_CONFIG.RAND_LEN_FE)]),
+                rho=Randomness(data=[Fp(0) for _ in range(TARGET_CONFIG.RAND_LEN_FE)]),
                 hashes=HashDigestList(data=[]),
             )