feat: add signature aggregation using python bindings

packages/testing/src/consensus_testing/keys.py

@@ -30,14 +30,15 @@
 from typing import TYPE_CHECKING, Iterator, Self
 
 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 ByteListMib
 from lean_spec.types import Uint64
 
 if TYPE_CHECKING:
@@ -216,23 +217,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[ByteListMib] = []
+        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(ByteListMib(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(ByteListMib(data=payload))
+
+        return AttestationSignatures(data=proof_blobs)
 
 
 def _generate_single_keypair(num_epochs: int) -> dict[str, str]:

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

@@ -106,11 +106,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
@@ -168,28 +167,33 @@ def verify_signatures(self, parent_state: "State") -> bool:
 
         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,
-                ), "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 ByteListMib
 
 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[ByteListMib]):
+    """
+    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 = ByteListMib
     LIMIT = int(VALIDATOR_REGISTRY_LIMIT)

src/lean_spec/subspecs/forkchoice/store.py

@@ -458,27 +458,20 @@ def on_block(self, signed_block_with_attestation: SignedBlockWithAttestation) ->
 
         # Process block body attestations.
         aggregated_attestations = signed_block_with_attestation.message.block.body.attestations
-        attestation_signatures = signed_block_with_attestation.signature.attestation_signatures
 
         assert len(aggregated_attestations) == len(attestation_signatures), (
             "Attestation signature groups must match aggregated attestations"
         )
 
-        for aggregated_attestation, aggregated_signature in zip(
-            aggregated_attestations, attestation_signatures, strict=True
-        ):
+        for aggregated_attestation in aggregated_attestations:
             validator_ids = aggregated_attestation.aggregation_bits.to_validator_indices()
-
-            assert len(validator_ids) == len(aggregated_signature), (
-                "Aggregated attestation signature count mismatch"
-            )
-
-            for validator_id, signature in zip(validator_ids, aggregated_signature, strict=True):
+            for validator_id in validator_ids:
+                # Signature bytes are not needed for forkchoice once the block-level
+                # aggregated proof has been verified.
                 store = store.on_attestation(
                     signed_attestation=SignedAttestation(
                         validator_id=validator_id,
                         message=aggregated_attestation.data,
-                        signature=signature,
                     ),
                     is_from_block=True,
                 )

src/lean_spec/subspecs/xmss/aggregation.py

@@ -0,0 +1,132 @@
+"""
+lean-multisig aggregation helpers bridging leanSpec containers to native bindings.
+
+This module wraps the Python bindings exposed by the `lean-multisig` project to provide
+XMSS signature aggregation + verification.
+
+The aggregated signatures are stored as raw payload bytes produced by
+`lean_multisig.aggregate_signatures`.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Sequence
+
+from lean_spec.subspecs.xmss.containers import PublicKey, Signature
+from lean_spec.types import Uint64
+
+
+class LeanMultisigError(RuntimeError):
+    """Base exception for lean-multisig aggregation helpers."""
+
+
+class LeanMultisigUnavailableError(LeanMultisigError):
+    """Raised when the lean-multisig Python bindings cannot be imported."""
+
+
+class LeanMultisigAggregationError(LeanMultisigError):
+    """Raised when lean-multisig fails to aggregate or verify signatures."""
+
+
+@lru_cache(maxsize=1)
+def _import_lean_multisig():
+    try:
+        import lean_multisig  # type: ignore
+    except ModuleNotFoundError as exc:  # pragma: no cover - import is environment-specific
+        raise LeanMultisigUnavailableError(
+            "lean-multisig bindings are required. Install them with `uv pip install lean-multisig` "
+            "(or your local editable install) from the leanSpec repository."
+        ) from exc
+    return lean_multisig
+
+
+@lru_cache(maxsize=1)
+def _ensure_prover_setup() -> None:
+    """Run the (expensive) prover setup routine exactly once."""
+    _import_lean_multisig().setup_prover()
+
+
+@lru_cache(maxsize=1)
+def _ensure_verifier_setup() -> None:
+    """Run the verifier setup routine exactly once."""
+    _import_lean_multisig().setup_verifier()
+
+
+def _coerce_epoch(epoch: int | Uint64) -> int:
+    value = int(epoch)
+    if value < 0 or value >= 2**32:
+        raise ValueError("epoch must fit in uint32 for lean-multisig aggregation")
+    return value
+
+
+def aggregate_signatures(
+    public_keys: Sequence[PublicKey],
+    signatures: Sequence[Signature],
+    message: bytes,
+    epoch: int | Uint64,
+) -> bytes:
+    """
+    Aggregate XMSS signatures using lean-multisig.
+
+    Args:
+        public_keys: Public keys of the signers, one per signature.
+        signatures: Individual XMSS signatures to aggregate.
+        message: The 32-byte message that was signed.
+        epoch: The epoch in which the signatures were created.
+
+    Returns:
+        Raw bytes of the aggregated signature payload.
+
+    Raises:
+        LeanMultisigError: If lean-multisig is unavailable or aggregation fails.
+    """
+    lean_multisig = _import_lean_multisig()
+    _ensure_prover_setup()
+    try:
+        # `lean_multisig` expects serialized keys/signatures as raw bytes.
+        # We use leanSpec's SSZ encoding for these containers.
+        pub_keys_bytes = [pk.encode_bytes() for pk in public_keys]
+        sig_bytes = [sig.encode_bytes() for sig in signatures]
+
+        aggregated_bytes = lean_multisig.aggregate_signatures(
+            pub_keys_bytes,
+            sig_bytes,
+            message,
+            _coerce_epoch(epoch),
+        )
+        return aggregated_bytes
+    except Exception as exc:
+        raise LeanMultisigAggregationError(f"lean-multisig aggregation failed: {exc}") from exc
+
+
+def verify_aggregated_payload(
+    public_keys: Sequence[PublicKey],
+    payload: bytes,
+    message: bytes,
+    epoch: int | Uint64,
+) -> None:
+    """
+    Verify a lean-multisig aggregated signature payload.
+
+    Args:
+        public_keys: Public keys of the signers, one per original signature.
+        payload: Raw bytes of the aggregated signature payload.
+        message: The 32-byte message that was signed.
+        epoch: The epoch in which the signatures were created.
+
+    Raises:
+        LeanMultisigError: If lean-multisig is unavailable or verification fails.
+    """
+    lean_multisig = _import_lean_multisig()
+    _ensure_verifier_setup()
+    try:
+        pub_keys_bytes = [pk.encode_bytes() for pk in public_keys]
+        lean_multisig.verify_aggregated_signatures(
+            pub_keys_bytes,
+            message,
+            payload,
+            _coerce_epoch(epoch),
+        )
+    except Exception as exc:
+        raise LeanMultisigAggregationError(f"lean-multisig verification failed: {exc}") from exc

src/lean_spec/types/byte_arrays.py

@@ -383,3 +383,9 @@ class ByteList2048(BaseByteList):
     """Variable-length byte list with a limit of 2048 bytes."""
 
     LIMIT = 2048
+
+
+class ByteListMib(BaseByteList):
+    """Variable-length byte list with a limit of 1048576 bytes."""
+
+    LIMIT = 1024 * 1024