feat: add signature aggregation using python bindings

packages/testing/src/consensus_testing/keys.py

@@ -40,13 +40,15 @@
 
 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.xmss.aggregation import aggregate_signatures
 from lean_spec.subspecs.xmss.containers import PublicKey, SecretKey, Signature
+from lean_spec.subspecs.xmss.interface import TEST_SIGNATURE_SCHEME, GeneralizedXmssScheme
+from lean_spec.types.byte_arrays import LeanAggregatedSignature
 from lean_spec.subspecs.xmss.interface import (
     PROD_SIGNATURE_SCHEME,
     TEST_SIGNATURE_SCHEME,
@@ -308,23 +310,43 @@ def build_attestation_signatures(
         self,
         aggregated_attestations: AggregatedAttestations,
         signature_lookup: Mapping[tuple[Uint64, bytes], Signature] | None = None,
+        payload_lookup: Mapping[bytes, bytes] | 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 blob using
+        `xmss_aggregate_signatures` (via `aggregate_signatures`).
+        """
         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[LeanAggregatedSignature] = []
+        for agg in aggregated_attestations:
+            validator_ids = agg.aggregation_bits.to_validator_indices()
+            message = agg.data.data_root_bytes()
+            epoch = agg.data.slot
+
+            if payload_lookup is not None and message in payload_lookup:
+                proof_blobs.append(LeanAggregatedSignature(data=payload_lookup[message]))
+                continue
+
+            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.
+            payload = aggregate_signatures(
+                public_keys=public_keys,
+                signatures=signatures,
+                message=message,
+                epoch=epoch,
+            )
+            proof_blobs.append(LeanAggregatedSignature(data=payload))
+
+        return AttestationSignatures(data=proof_blobs)
 
 
 def _generate_single_keypair(

pyproject.toml

@@ -26,7 +26,11 @@ classifiers = [
     "Programming Language :: Python :: 3.14",
 ]
 requires-python = ">=3.12"
-dependencies = ["pydantic>=2.12.0,<3", "typing-extensions>=4.4"]
+dependencies = [
+    "pydantic>=2.12.0,<3",
+    "typing-extensions>=4.4",
+    "lean-multisig>=0.1.0",
+]
 
 [project.license]
 file = "LICENSE"
@@ -118,6 +122,7 @@ test = [
     "pytest-xdist>=3.6.1,<4",
     "hypothesis>=6.138.14",
     "lean-ethereum-testing",
+    "lean-multisig>=0.1.0",
 ]
 lint = [
     "ty>=0.0.1a34",

src/lean_spec/subspecs/containers/block/block.py

@@ -107,11 +107,10 @@ class BlockSignatures(Container):
     """Attestation signatures for the aggregated attestations in the block body.
 
     Each entry corresponds to an aggregated attestation from the block body and
-    contains all XMSS signatures from the participating validators.
+    contains the leanVM aggregated signature proof bytes for the participating validators.
 
     TODO:
-    - Currently, this is list of lists of signatures.
-    - The list of signatures will be replaced by a BytesArray to include leanVM aggregated proof.
+    - Eventually this field will be replaced by a single SNARK aggregating *all* signatures.
     """
 
     proposer_signature: XmssSignature
@@ -172,29 +171,33 @@ def verify_signatures(
 
         validators = parent_state.validators
 
+        from lean_spec.subspecs.xmss.aggregation import (
+            LeanMultisigError,
+            verify_aggregated_payload,
+        )
+
         for aggregated_attestation, aggregated_signature in zip(
             aggregated_attestations, attestation_signatures, strict=True
         ):
             validator_ids = aggregated_attestation.aggregation_bits.to_validator_indices()
 
-            assert len(aggregated_signature) == len(validator_ids), (
-                "Aggregated attestation signature count mismatch"
-            )
-
             attestation_root = aggregated_attestation.data.data_root_bytes()
 
-            # Verify each validator's attestation signature
-            for validator_id, signature in zip(validator_ids, aggregated_signature, strict=True):
+            # Verify the leanVM aggregated proof for this attestation
+            for validator_id in validator_ids:
                 # Ensure validator exists in the active set
                 assert validator_id < Uint64(len(validators)), "Validator index out of range"
-                validator = validators[validator_id]
-
-                assert signature.verify(
-                    validator.get_pubkey(),
-                    aggregated_attestation.data.slot,
-                    attestation_root,
-                    scheme,
-                ), "Attestation signature verification failed"
+
+            public_keys = [validators[vid].get_pubkey() for vid in validator_ids]
+            try:
+                verify_aggregated_payload(
+                    public_keys=public_keys,
+                    payload=bytes(aggregated_signature),
+                    message=attestation_root,
+                    epoch=aggregated_attestation.data.slot,
+                )
+            except LeanMultisigError as exc:
+                raise AssertionError("Attestation aggregated signature verification failed") from exc
 
         # Verify proposer attestation signature
         proposer_attestation = self.message.proposer_attestation

src/lean_spec/subspecs/containers/block/types.py

@@ -1,9 +1,10 @@
 """Block-specific SSZ types for the Lean Ethereum consensus specification."""
 
 from lean_spec.types import SSZList
+from lean_spec.types.byte_arrays import LeanAggregatedSignature
 
 from ...chain.config import VALIDATOR_REGISTRY_LIMIT
-from ..attestation import AggregatedAttestation, AttestationData, NaiveAggregatedSignature
+from ..attestation import AggregatedAttestation, AttestationData
 
 
 class AggregatedAttestations(SSZList[AggregatedAttestation]):
@@ -22,8 +23,14 @@ def has_duplicate_data(self) -> bool:
         return False
 
 
-class AttestationSignatures(SSZList[NaiveAggregatedSignature]):
-    """List of per-attestation naive signature lists aligned with block body attestations."""
+class AttestationSignatures(SSZList[LeanAggregatedSignature]):
+    """
+    List of per-attestation aggregated signature proof blobs.
 
-    ELEMENT_TYPE = NaiveAggregatedSignature
+    Each entry corresponds to an aggregated attestation from the block body and contains
+    the raw bytes of the leanVM XMSSAggregatedSignature proof produced by
+    `xmss_aggregate_signatures`.
+    """
+
+    ELEMENT_TYPE = LeanAggregatedSignature
     LIMIT = int(VALIDATOR_REGISTRY_LIMIT)

src/lean_spec/subspecs/containers/state/state.py

@@ -12,7 +12,8 @@
     is_proposer,
 )
 
-from ..attestation import AggregatedAttestation, Attestation, SignedAttestation
+from ..attestation import AggregatedAttestation, Attestation
+from lean_spec.types.byte_arrays import LeanAggregatedSignature
 
 if TYPE_CHECKING:
     from lean_spec.subspecs.xmss.containers import Signature
@@ -626,27 +627,34 @@ def build_block(
         proposer_index: Uint64,
         parent_root: Bytes32,
         attestations: list[Attestation] | None = None,
-        available_signed_attestations: Iterable[SignedAttestation] | None = None,
+        available_attestations: Iterable[Attestation] | None = None,
         known_block_roots: AbstractSet[Bytes32] | None = None,
+        gossip_attestation_signatures: dict[tuple[Uint64, bytes], "Signature"] | None = None,
+        block_attestation_signatures: dict[tuple[Uint64, bytes], list[LeanAggregatedSignature]] | None = None,
     ) -> tuple[Block, "State", list[Attestation], list["Signature"]]:
         """
         Build a valid block on top of this state.
 
         Computes the post-state and creates a block with the correct state root.
 
-        If `available_signed_attestations` and `known_block_roots` are provided,
+        If `available_attestations` and `known_block_roots` are provided,
         performs fixed-point attestation collection: iteratively adds valid
         attestations until no more can be included. This is necessary because
         processing attestations may update the justified checkpoint, which may
         make additional attestations valid.
 
+        Signatures are looked up from the provided signature maps using
+        (validator_id, attestation_data_root) as the key.
+
         Args:
             slot: Target slot for the block.
             proposer_index: Validator index of the proposer.
             parent_root: Root of the parent block.
             attestations: Initial attestations to include.
-            available_signed_attestations: Pool of attestations to collect from.
+            available_attestations: Pool of attestations to collect from.
             known_block_roots: Set of known block roots for attestation validation.
+            gossip_attestation_signatures: Map of (validator_id, data_root) to XMSS signatures.
+            block_attestation_signatures: Map of (validator_id, data_root) to aggregated signature payloads.
 
         Returns:
             Tuple of (Block, post-State, collected attestations, signatures).
@@ -677,19 +685,37 @@ def build_block(
             post_state = self.process_slots(slot).process_block(candidate_block)
 
             # No attestation source provided: done after computing post_state
-            if available_signed_attestations is None or known_block_roots is None:
+            if available_attestations is None or known_block_roots is None:
                 break
 
             # Find new valid attestations matching post-state justification
             new_attestations: list[Attestation] = []
             new_signatures: list[Signature] = []
 
-            for signed_attestation in available_signed_attestations:
-                data = signed_attestation.message
-                attestation = Attestation(
-                    validator_id=signed_attestation.validator_id,
-                    data=data,
-                )
+            for attestation in available_attestations:
+                data = attestation.data
+                validator_id = attestation.validator_id
+                data_root = data.data_root_bytes()
+
+                # Check if we have a signature for this (validator_id, data_root)
+                # First check gossip signatures, then block signatures (first entry if available)
+                signature = None
+                if gossip_attestation_signatures:
+                    signature = gossip_attestation_signatures.get((validator_id, data_root))
+
+                # If not found in gossip and block signatures available, try those
+                # Note: For now, we only use block signatures if gossip signature not found
+                # TODO: Implement recursive aggregation to combine multiple aggregated signatures
+                if signature is None and block_attestation_signatures:
+                    aggregated_sigs = block_attestation_signatures.get((validator_id, data_root))
+                    if aggregated_sigs:
+                        # For now, skip if only aggregated signatures available
+                        # because we don't have recursive aggregation yet
+                        continue
+
+                # Skip if no signature found
+                if signature is None:
+                    continue
 
                 # Skip if target block is unknown
                 if data.head.root not in known_block_roots:
@@ -702,7 +728,7 @@ def build_block(
                 # Add attestation if not already included
                 if attestation not in attestations:
                     new_attestations.append(attestation)
-                    new_signatures.append(signed_attestation.signature)
+                    new_signatures.append(signature)
 
             # Fixed point reached: no new attestations found
             if not new_attestations: