From 0b8e46616ca55f465802e107f77f722084e8ddb5 Mon Sep 17 00:00:00 2001 From: Thomas Coratger Date: Mon, 29 Dec 2025 23:01:05 +0100 Subject: [PATCH 1/5] aggregation: cleanup and reorg --- .../testing/src/consensus_testing/keys.py | 26 ++- .../test_fixtures/fork_choice.py | 11 +- .../test_fixtures/state_transition.py | 5 +- .../test_fixtures/verify_signatures.py | 3 +- src/lean_spec/subspecs/chain/config.py | 3 +- src/lean_spec/subspecs/containers/__init__.py | 4 - .../containers/attestation/__init__.py | 3 +- .../{types.py => aggregation_bits.py} | 45 ++--- .../containers/attestation/attestation.py | 2 +- .../subspecs/containers/block/block.py | 8 +- .../subspecs/containers/block/types.py | 15 +- .../subspecs/containers/state/state.py | 181 +++++++----------- src/lean_spec/subspecs/forkchoice/store.py | 57 +++--- src/lean_spec/subspecs/xmss/aggregation.py | 135 ++++++------- src/lean_spec/types/__init__.py | 2 + .../containers/test_state_aggregation.py | 165 ++++++++-------- .../forkchoice/test_store_attestations.py | 9 +- .../subspecs/forkchoice/test_validator.py | 16 +- 18 files changed, 338 insertions(+), 352 deletions(-) rename src/lean_spec/subspecs/containers/attestation/{types.py => aggregation_bits.py} (56%) diff --git a/packages/testing/src/consensus_testing/keys.py b/packages/testing/src/consensus_testing/keys.py index 4819408e..fe35fbdc 100755 --- a/packages/testing/src/consensus_testing/keys.py +++ b/packages/testing/src/consensus_testing/keys.py @@ -39,14 +39,15 @@ from lean_spec.config import LEAN_ENV from lean_spec.subspecs.containers import AttestationData +from lean_spec.subspecs.containers.attestation import AggregationBits 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 ( - AttestationSignatureKey, - MultisigAggregatedSignature, + AggregatedSignatureProof, + SignatureKey, ) from lean_spec.subspecs.xmss.containers import KeyPair, PublicKey, Signature from lean_spec.subspecs.xmss.interface import ( @@ -54,7 +55,7 @@ TEST_SIGNATURE_SCHEME, GeneralizedXmssScheme, ) -from lean_spec.types import Uint64 +from lean_spec.types import Bytes32, Uint64 if TYPE_CHECKING: from collections.abc import Mapping @@ -276,7 +277,7 @@ def sign_attestation_data( def build_attestation_signatures( self, aggregated_attestations: AggregatedAttestations, - signature_lookup: Mapping[AttestationSignatureKey, Signature] | None = None, + signature_lookup: Mapping[SignatureKey, Signature] | None = None, ) -> AttestationSignatures: """ Build `AttestationSignatures` for already-aggregated attestations. @@ -286,7 +287,7 @@ def build_attestation_signatures( """ lookup = signature_lookup or {} - proof_blobs: list[MultisigAggregatedSignature] = [] + proofs: list[AggregatedSignatureProof] = [] for agg in aggregated_attestations: validator_ids = agg.aggregation_bits.to_validator_indices() message = agg.data.data_root_bytes() @@ -294,21 +295,26 @@ def build_attestation_signatures( 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)) + ( + lookup.get(SignatureKey(vid, Bytes32(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( + participants = AggregationBits.from_validator_indices(validator_ids) + proof = AggregatedSignatureProof.aggregate( + participants=participants, public_keys=public_keys, signatures=signatures, message=message, - epoch=epoch, + epoch=Uint64(epoch), ) - proof_blobs.append(aggregated_signature) + proofs.append(proof) - return AttestationSignatures(data=proof_blobs) + return AttestationSignatures(data=proofs) def _generate_single_keypair( diff --git a/packages/testing/src/consensus_testing/test_fixtures/fork_choice.py b/packages/testing/src/consensus_testing/test_fixtures/fork_choice.py index 3f62bec1..97ace502 100644 --- a/packages/testing/src/consensus_testing/test_fixtures/fork_choice.py +++ b/packages/testing/src/consensus_testing/test_fixtures/fork_choice.py @@ -29,7 +29,7 @@ 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.aggregation import AttestationSignatureKey +from lean_spec.subspecs.xmss.aggregation import SignatureKey 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 @@ -404,14 +404,14 @@ def _build_attestations_from_spec( block_registry: dict[str, Block], parent_root: Bytes32, key_manager: XmssKeyManager, - ) -> tuple[list[Attestation], dict[AttestationSignatureKey, Signature]]: + ) -> tuple[list[Attestation], dict[SignatureKey, Signature]]: """Build attestations list from BlockSpec and their signatures.""" if spec.attestations is None: return [], {} parent_state = store.states[parent_root] attestations = [] - signature_lookup: dict[AttestationSignatureKey, Signature] = {} + signature_lookup: dict[SignatureKey, Signature] = {} for att_spec in spec.attestations: if isinstance(att_spec, SignedAttestationSpec): @@ -423,9 +423,10 @@ def _build_attestations_from_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 + sig_key = SignatureKey( + attestation.validator_id, Bytes32(attestation.data.data_root_bytes()) ) + signature_lookup[sig_key] = signed_att.signature return attestations, signature_lookup diff --git a/packages/testing/src/consensus_testing/test_fixtures/state_transition.py b/packages/testing/src/consensus_testing/test_fixtures/state_transition.py index 7cb55c62..c77389d4 100644 --- a/packages/testing/src/consensus_testing/test_fixtures/state_transition.py +++ b/packages/testing/src/consensus_testing/test_fixtures/state_transition.py @@ -9,6 +9,7 @@ from lean_spec.subspecs.containers.block.types import AggregatedAttestations from lean_spec.subspecs.containers.state.state import State from lean_spec.subspecs.ssz.hash import hash_tree_root +from lean_spec.subspecs.xmss.aggregation import SignatureKey from lean_spec.types import Bytes32, Uint64 from ..keys import get_shared_key_manager @@ -262,7 +263,9 @@ def _build_block_from_spec(self, spec: BlockSpec, state: State) -> tuple[Block, 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( + SignatureKey( + att.validator_id, Bytes32(att.data.data_root_bytes()) + ): key_manager.sign_attestation_data( att.validator_id, att.data, ) diff --git a/packages/testing/src/consensus_testing/test_fixtures/verify_signatures.py b/packages/testing/src/consensus_testing/test_fixtures/verify_signatures.py index 6e5e4df0..4ca7ef89 100644 --- a/packages/testing/src/consensus_testing/test_fixtures/verify_signatures.py +++ b/packages/testing/src/consensus_testing/test_fixtures/verify_signatures.py @@ -21,6 +21,7 @@ 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.aggregation import SignatureKey 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 @@ -187,7 +188,7 @@ def _build_block_from_spec( # 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 + SignatureKey(att.validator_id, Bytes32(att.data.data_root_bytes())): sig for att, sig in zip(attestations, attestation_signature_inputs, strict=True) } diff --git a/src/lean_spec/subspecs/chain/config.py b/src/lean_spec/subspecs/chain/config.py index 955440e9..aa00fee7 100644 --- a/src/lean_spec/subspecs/chain/config.py +++ b/src/lean_spec/subspecs/chain/config.py @@ -7,7 +7,8 @@ from typing_extensions import Final -from lean_spec.types import StrictBaseModel, Uint64 +from lean_spec.types.base import StrictBaseModel +from lean_spec.types.uint import Uint64 # --- Time Parameters --- diff --git a/src/lean_spec/subspecs/containers/__init__.py b/src/lean_spec/subspecs/containers/__init__.py index 33ce84af..3f42e8f0 100644 --- a/src/lean_spec/subspecs/containers/__init__.py +++ b/src/lean_spec/subspecs/containers/__init__.py @@ -10,10 +10,8 @@ from .attestation import ( AggregatedAttestation, - AggregationBits, Attestation, AttestationData, - AttestationsByValidator, SignedAttestation, ) from .block import ( @@ -30,10 +28,8 @@ __all__ = [ "AggregatedAttestation", - "AggregationBits", "Attestation", "AttestationData", - "AttestationsByValidator", "Block", "BlockBody", "BlockHeader", diff --git a/src/lean_spec/subspecs/containers/attestation/__init__.py b/src/lean_spec/subspecs/containers/attestation/__init__.py index 615d2f4d..febbf61e 100644 --- a/src/lean_spec/subspecs/containers/attestation/__init__.py +++ b/src/lean_spec/subspecs/containers/attestation/__init__.py @@ -1,18 +1,17 @@ """Attestation containers and related types for the Lean spec.""" +from .aggregation_bits import AggregationBits from .attestation import ( AggregatedAttestation, Attestation, AttestationData, SignedAttestation, ) -from .types import AggregationBits, AttestationsByValidator __all__ = [ "AggregatedAttestation", "AggregationBits", "Attestation", "AttestationData", - "AttestationsByValidator", "SignedAttestation", ] diff --git a/src/lean_spec/subspecs/containers/attestation/types.py b/src/lean_spec/subspecs/containers/attestation/aggregation_bits.py similarity index 56% rename from src/lean_spec/subspecs/containers/attestation/types.py rename to src/lean_spec/subspecs/containers/attestation/aggregation_bits.py index 0e34c6b6..83841512 100644 --- a/src/lean_spec/subspecs/containers/attestation/types.py +++ b/src/lean_spec/subspecs/containers/attestation/aggregation_bits.py @@ -1,23 +1,19 @@ -"""Attestation-related SSZ types for the Lean consensus specification.""" +"""Aggregation bits for tracking validator participation.""" from __future__ import annotations -from typing import TYPE_CHECKING - -from lean_spec.types import Uint64 +from lean_spec.subspecs.chain.config import VALIDATOR_REGISTRY_LIMIT +from lean_spec.types import Boolean, Uint64 from lean_spec.types.bitfields import BaseBitlist -from ...chain.config import VALIDATOR_REGISTRY_LIMIT - -if TYPE_CHECKING: - from .attestation import AttestationData - -AttestationsByValidator = dict[Uint64, "AttestationData"] -"""Mapping from validator index to attestation data.""" - class AggregationBits(BaseBitlist): - """Bitlist representing validator participation in an attestation.""" + """ + Bitlist representing validator participation in an attestation or signature. + + A general-purpose bitfield for tracking which validators have participated + in some collective action (attestation, signature aggregation, etc.). + """ LIMIT = int(VALIDATOR_REGISTRY_LIMIT) @@ -36,19 +32,23 @@ def from_validator_indices(cls, indices: list[Uint64]) -> AggregationBits: AssertionError: If no indices are provided. AssertionError: If any index is outside the supported LIMIT. """ - ids = [int(i) for i in indices] - if not ids: + # Require at least one validator for a valid aggregation. + if not indices: raise AssertionError("Aggregated attestation must reference at least one validator") - max_id = max(ids) - if max_id >= cls.LIMIT: - raise AssertionError("Validator index out of range for aggregation bits") + # Convert to a set of native ints. + # + # This combines int conversion and deduplication in a single O(N) pass. + ids = {int(i) for i in indices} - bits = [False] * (max_id + 1) - for i in ids: - bits[i] = True + # Validate bounds: max index must be within registry limit. + if (max_id := max(ids)) >= cls.LIMIT: + raise AssertionError("Validator index out of range for aggregation bits") - return cls(data=bits) + # Build bit list: + # - True at positions present in indices, + # - False elsewhere. + return cls(data=[Boolean(i in ids) for i in range(max_id + 1)]) def to_validator_indices(self) -> list[Uint64]: """ @@ -60,6 +60,7 @@ def to_validator_indices(self) -> list[Uint64]: Raises: AssertionError: If no bits are set. """ + # Extract indices where bit is set; fail if none found. if not (indices := [Uint64(i) for i, bit in enumerate(self.data) if bool(bit)]): raise AssertionError("Aggregated attestation must reference at least one validator") diff --git a/src/lean_spec/subspecs/containers/attestation/attestation.py b/src/lean_spec/subspecs/containers/attestation/attestation.py index be6acb34..18080a5f 100644 --- a/src/lean_spec/subspecs/containers/attestation/attestation.py +++ b/src/lean_spec/subspecs/containers/attestation/attestation.py @@ -22,7 +22,7 @@ from ...xmss.containers import Signature from ..checkpoint import Checkpoint -from .types import AggregationBits +from .aggregation_bits import AggregationBits class AttestationData(Container): diff --git a/src/lean_spec/subspecs/containers/block/block.py b/src/lean_spec/subspecs/containers/block/block.py index f58cd211..a93d29b8 100644 --- a/src/lean_spec/subspecs/containers/block/block.py +++ b/src/lean_spec/subspecs/containers/block/block.py @@ -12,9 +12,7 @@ from typing import TYPE_CHECKING from lean_spec.subspecs.containers.slot import Slot -from lean_spec.subspecs.xmss.aggregation import ( - MultisigError, -) +from lean_spec.subspecs.xmss.aggregation import AggregationError from lean_spec.subspecs.xmss.interface import TARGET_SIGNATURE_SCHEME, GeneralizedXmssScheme from lean_spec.types import Bytes32, Uint64 from lean_spec.types.container import Container @@ -189,12 +187,12 @@ def verify_signatures( public_keys = [validators[vid].get_pubkey() for vid in validator_ids] try: - aggregated_signature.verify_aggregated_payload( + aggregated_signature.verify( public_keys=public_keys, message=attestation_data_root, epoch=aggregated_attestation.data.slot, ) - except MultisigError as exc: + except AggregationError as exc: raise AssertionError( f"Attestation aggregated signature verification failed: {exc}" ) from exc diff --git a/src/lean_spec/subspecs/containers/block/types.py b/src/lean_spec/subspecs/containers/block/types.py index e16a93d9..b9aaa672 100644 --- a/src/lean_spec/subspecs/containers/block/types.py +++ b/src/lean_spec/subspecs/containers/block/types.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING -from lean_spec.subspecs.xmss.aggregation import MultisigAggregatedSignature +from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof from lean_spec.types import Bytes32, SSZList from ...chain.config import VALIDATOR_REGISTRY_LIMIT @@ -24,13 +24,16 @@ class AggregatedAttestations(SSZList[AggregatedAttestation]): LIMIT = int(VALIDATOR_REGISTRY_LIMIT) -class AttestationSignatures(SSZList[MultisigAggregatedSignature]): +class AttestationSignatures(SSZList[AggregatedSignatureProof]): """ - List of per-attestation aggregated signature proof blobs. + List of per-attestation aggregated signature proofs. - Each entry corresponds to an aggregated attestation from the block body and contains - the raw bytes of the leanVM signature aggregation proof. + Each entry corresponds to an aggregated attestation from the block body. + + It contains: + - the participants bitfield, + - proof bytes from leanVM signature aggregation. """ - ELEMENT_TYPE = MultisigAggregatedSignature + ELEMENT_TYPE = AggregatedSignatureProof LIMIT = int(VALIDATOR_REGISTRY_LIMIT) diff --git a/src/lean_spec/subspecs/containers/state/state.py b/src/lean_spec/subspecs/containers/state/state.py index 7eb33fb4..01815d08 100644 --- a/src/lean_spec/subspecs/containers/state/state.py +++ b/src/lean_spec/subspecs/containers/state/state.py @@ -4,9 +4,8 @@ from lean_spec.subspecs.ssz.hash import hash_tree_root from lean_spec.subspecs.xmss.aggregation import ( - AggregatedSignaturePayloads, - AttestationSignatureKey, - MultisigAggregatedSignature, + AggregatedSignatureProof, + SignatureKey, ) from lean_spec.subspecs.xmss.containers import PublicKey, Signature from lean_spec.types import ( @@ -616,19 +615,17 @@ def _aggregate_signatures_from_gossip( validator_ids: list[Uint64], data_root: bytes, epoch: Slot, - gossip_signatures: dict[AttestationSignatureKey, "Signature"] | None = None, - ) -> tuple[MultisigAggregatedSignature, AggregationBits, set[Uint64]] | None: + gossip_signatures: dict[SignatureKey, "Signature"] | None = None, + ) -> tuple[AggregatedSignatureProof, set[Uint64]] | None: """ - Aggregate per-validator XMSS signatures into a single payload. + Aggregate per-validator XMSS signatures into a single proof. Returns: - ------- - - MultisigAggregatedSignature: The aggregated signature. - - AggregationBits: The aggregation bits. - - set[Uint64]: The validator ids whose id's were missing from the gossip signatures. + A tuple of (proof, missing_validator_ids) or None if no signatures found. + The proof contains the participants bitfield. Raises: - AssertionError: If the aggregation fails. + AggregationError: If the aggregation fails. """ if not gossip_signatures or not validator_ids: return None @@ -640,8 +637,8 @@ def _aggregate_signatures_from_gossip( missing_validator_ids: set[Uint64] = set() for validator_index in validator_ids: - # Attempt to retrieve the signature; fail fast if any are missing. - if (sig := gossip_signatures.get((validator_index, data_root))) is None: + key = SignatureKey(validator_index, Bytes32(data_root)) + if (sig := gossip_signatures.get(key)) is None: missing_validator_ids.add(validator_index) continue @@ -652,16 +649,15 @@ def _aggregate_signatures_from_gossip( if not included_validator_ids: return None - return ( - MultisigAggregatedSignature.aggregate_signatures( - public_keys=public_keys, - signatures=signatures, - message=data_root, - epoch=epoch, - ), - AggregationBits.from_validator_indices(list(included_validator_ids)), - missing_validator_ids, + participants = AggregationBits.from_validator_indices(list(included_validator_ids)) + proof = AggregatedSignatureProof.aggregate( + participants=participants, + public_keys=public_keys, + signatures=signatures, + message=data_root, + epoch=Uint64(epoch), ) + return proof, missing_validator_ids def build_block( self, @@ -671,10 +667,9 @@ def build_block( attestations: list[Attestation] | None = None, available_attestations: Iterable[Attestation] | None = None, known_block_roots: AbstractSet[Bytes32] | None = None, - gossip_signatures: dict[AttestationSignatureKey, "Signature"] | None = None, - aggregated_payloads: dict[AttestationSignatureKey, AggregatedSignaturePayloads] - | None = None, - ) -> tuple[Block, "State", list[AggregatedAttestation], list[MultisigAggregatedSignature]]: + gossip_signatures: dict[SignatureKey, "Signature"] | None = None, + aggregated_payloads: dict[SignatureKey, list[AggregatedSignatureProof]] | None = None, + ) -> tuple[Block, "State", list[AggregatedAttestation], list[AggregatedSignatureProof]]: """ Build a valid block on top of this state. @@ -737,7 +732,7 @@ def build_block( data = attestation.data validator_id = attestation.validator_id data_root = data.data_root_bytes() - attestation_key = (validator_id, data_root) + sig_key = SignatureKey(validator_id, Bytes32(data_root)) # Skip if target block is unknown if data.head.root not in known_block_roots: @@ -752,17 +747,14 @@ def build_block( continue # We can only include an attestation if we have some way to later provide - # an aggregated payload for its group: + # an aggregated proof for its group: # - either a per validator XMSS signature from gossip, or - # - at least one aggregated payload learned from a block that references + # - at least one aggregated proof learned from a block that references # this validator+data. - has_gossip_sig = bool(gossip_signatures and attestation_key in gossip_signatures) - - has_block_payload = bool( - aggregated_payloads and attestation_key in aggregated_payloads - ) + has_gossip_sig = bool(gossip_signatures and sig_key in gossip_signatures) + has_block_proof = bool(aggregated_payloads and sig_key in aggregated_payloads) - if has_gossip_sig or has_block_payload: + if has_gossip_sig or has_block_proof: new_attestations.append(attestation) # Fixed point reached: no new attestations found @@ -798,10 +790,9 @@ def build_block( def compute_aggregated_signatures( self, attestations: list[Attestation], - gossip_signatures: dict[AttestationSignatureKey, "Signature"] | None = None, - aggregated_payloads: dict[AttestationSignatureKey, AggregatedSignaturePayloads] - | None = None, - ) -> tuple[list[AggregatedAttestation], list[MultisigAggregatedSignature]]: + gossip_signatures: dict[SignatureKey, "Signature"] | None = None, + aggregated_payloads: dict[SignatureKey, list[AggregatedSignatureProof]] | None = None, + ) -> tuple[list[AggregatedAttestation], list[AggregatedSignatureProof]]: """ Compute aggregated signatures for a set of attestations. @@ -817,7 +808,7 @@ def compute_aggregated_signatures( A tuple of `(aggregated_attestations, aggregated_signatures)`. """ final_aggregated_attestations: list[AggregatedAttestation] = [] - final_aggregated_signatures: list[MultisigAggregatedSignature] = [] + final_aggregated_proofs: list[AggregatedSignatureProof] = [] # Aggregate all the attestations into a single aggregated attestation. completely_aggregated_attestations = AggregatedAttestation.aggregate_by_data(attestations) @@ -839,8 +830,7 @@ def compute_aggregated_signatures( data_root = completely_aggregated_attestation.data.data_root_bytes() slot = completely_aggregated_attestation.data.slot - aggregated_signatures: list[MultisigAggregatedSignature] = [] - aggregated_signatures_bitlists: list[AggregationBits] = [] + proofs: list[AggregatedSignatureProof] = [] # Try to find per validator XMSS signatures from gossip. gossip_result = self._aggregate_signatures_from_gossip( @@ -851,101 +841,76 @@ def compute_aggregated_signatures( ) if gossip_result is not None: - gossip_aggregated_signature, gossip_aggregated_bitlist, remaining_validator_ids = ( - gossip_result - ) - aggregated_signatures.append(gossip_aggregated_signature) - aggregated_signatures_bitlists.append(gossip_aggregated_bitlist) + gossip_proof, remaining_validator_ids = gossip_result + proofs.append(gossip_proof) else: remaining_validator_ids = set(validator_ids) - # Try to pick a single aggregated signature from the aggregated signatures - # that covers the most validators. + # Pick existing aggregated proofs to cover remaining validators. while remaining_validator_ids: - aggregated_signature, aggregated_signature_bitlist, remaining_validator_ids = ( - self._pick_from_aggregated_signatures( - remaining_validator_ids, - data_root, - aggregated_payloads, - ) + proof, remaining_validator_ids = self._pick_from_aggregated_proofs( + remaining_validator_ids, + data_root, + aggregated_payloads, ) - aggregated_signatures.append(aggregated_signature) - aggregated_signatures_bitlists.append(aggregated_signature_bitlist) - - # TODO: Recursively aggregate the signatures in aggregated_signatures and append it to - # final_aggregated_signatures. Since we currently don't support recursive aggregation, - # we will just append the aggregated signatures to final_aggregated_signatures. - # There might be a case where we have multiple aggregated signatures that cover the same - # validators. This is fine for now, eventually we will recursively aggregate into one. - for aggregated_signature, aggregated_signature_bitlist in zip( - aggregated_signatures, aggregated_signatures_bitlists, strict=True - ): + proofs.append(proof) + + # TODO: Recursively aggregate the proofs. Since we currently don't support + # recursive aggregation, we just append each proof separately. This is fine + # for now, eventually we will recursively aggregate into one. + for proof in proofs: final_aggregated_attestations.append( AggregatedAttestation( - aggregation_bits=aggregated_signature_bitlist, + aggregation_bits=proof.participants, data=completely_aggregated_attestation.data, ) ) - final_aggregated_signatures.append(aggregated_signature) + final_aggregated_proofs.append(proof) - return final_aggregated_attestations, final_aggregated_signatures + return final_aggregated_attestations, final_aggregated_proofs - def _pick_from_aggregated_signatures( + def _pick_from_aggregated_proofs( self, remaining_validator_ids: set[Uint64], data_root: bytes, - aggregated_payloads: dict[AttestationSignatureKey, AggregatedSignaturePayloads] - | None = None, - ) -> tuple[MultisigAggregatedSignature, AggregationBits, set[Uint64]]: + aggregated_payloads: dict[SignatureKey, list[AggregatedSignatureProof]] | None = None, + ) -> tuple[AggregatedSignatureProof, set[Uint64]]: """ - Pick a single aggregated signature from aggregated signatures that covers the most - remaining validator ids. + Pick an aggregated proof that covers the most remaining validators. Args: - remaining_validator_ids: The validator ids to pick from. - data_root: The data root to pick from. - aggregated_payloads: The aggregated payloads to pick from. + remaining_validator_ids: The validator ids still needing coverage. + data_root: The attestation data root. + aggregated_payloads: Previously learned proofs keyed by (validator_id, data_root). Returns: - - MultisigAggregatedSignature: The aggregated signature. - - AggregationBits: The aggregation bits. - - set[Uint64]: The remaining validator ids. + A tuple of (proof, remaining_validator_ids after this proof is applied). Raises: - ValueError: If remaining validator ids is empty or aggregated payloads is not provided. - ValueError: If the best aggregated signature is not found. + ValueError: If no suitable proof is found. """ if not remaining_validator_ids: - raise ValueError( - "remaining validator ids cannot be empty when picking aggregated signatures" - ) + raise ValueError("remaining validator ids cannot be empty") if aggregated_payloads is None: - raise ValueError("aggregated payloads is required when gossip coverage is incomplete") + raise ValueError("aggregated payloads required when gossip coverage incomplete") - best_aggregated_signature: MultisigAggregatedSignature | None = None - best_aggregated_signature_participants: set[Uint64] = set() - best_remaining_validator_ids: set[Uint64] = set() + best_proof: AggregatedSignatureProof | None = None + best_overlap: set[Uint64] = set() + best_remaining: set[Uint64] = set() representative_validator_id = next(iter(remaining_validator_ids)) + key = SignatureKey(representative_validator_id, Bytes32(data_root)) - for aggregated_signature_bitlist, aggregated_signature in aggregated_payloads.get( - (representative_validator_id, data_root), [] - ): - participants = set(aggregated_signature_bitlist.to_validator_indices()) + for proof in aggregated_payloads.get(key, []): + participants = set(proof.participants.to_validator_indices()) overlap = participants.intersection(remaining_validator_ids) - if len(overlap) > len(best_aggregated_signature_participants): - best_aggregated_signature = aggregated_signature - best_aggregated_signature_participants = overlap - best_remaining_validator_ids = set(remaining_validator_ids).difference(overlap) - - if best_aggregated_signature is None: - raise ValueError( - "Failed to locate an aggregated signature payload for the remaining validators" - ) + if len(overlap) > len(best_overlap): + best_proof = proof + best_overlap = overlap + best_remaining = remaining_validator_ids - overlap - return ( - best_aggregated_signature, - AggregationBits.from_validator_indices(list(best_aggregated_signature_participants)), - best_remaining_validator_ids, - ) + if best_proof is None: + raise ValueError("Failed to locate aggregated proof for remaining validators") + + return best_proof, best_remaining diff --git a/src/lean_spec/subspecs/forkchoice/store.py b/src/lean_spec/subspecs/forkchoice/store.py index bf0e07ac..ac987401 100644 --- a/src/lean_spec/subspecs/forkchoice/store.py +++ b/src/lean_spec/subspecs/forkchoice/store.py @@ -22,10 +22,8 @@ SECONDS_PER_SLOT, ) from lean_spec.subspecs.containers import ( - AggregationBits, Attestation, AttestationData, - AttestationsByValidator, Block, Checkpoint, Config, @@ -38,9 +36,8 @@ from lean_spec.subspecs.containers.state import StateLookup from lean_spec.subspecs.ssz.hash import hash_tree_root from lean_spec.subspecs.xmss.aggregation import ( - AggregatedSignaturePayloads, - AttestationSignatureKey, - MultisigAggregatedSignature, + AggregatedSignatureProof, + SignatureKey, ) from lean_spec.subspecs.xmss.containers import Signature from lean_spec.subspecs.xmss.interface import TARGET_SIGNATURE_SCHEME, GeneralizedXmssScheme @@ -129,7 +126,7 @@ class Store(Container): `Store`'s latest justified and latest finalized checkpoints. """ - latest_known_attestations: AttestationsByValidator = {} + latest_known_attestations: dict[Uint64, AttestationData] = {} """ Latest attestation data by validator that have been processed. @@ -138,7 +135,7 @@ class Store(Container): - Only stores the attestation data, not signatures. """ - latest_new_attestations: AttestationsByValidator = {} + latest_new_attestations: dict[Uint64, AttestationData] = {} """ Latest attestation data by validator that are pending processing. @@ -148,20 +145,20 @@ class Store(Container): - Only stores the attestation data, not signatures. """ - gossip_signatures: Dict[AttestationSignatureKey, Signature] = {} + gossip_signatures: Dict[SignatureKey, Signature] = {} """ Per-validator XMSS signatures learned from gossip. - Keyed by (validator_id, attestation_data_root). + Keyed by SignatureKey(validator_id, attestation_data_root). """ - aggregated_payloads: Dict[AttestationSignatureKey, AggregatedSignaturePayloads] = {} + aggregated_payloads: Dict[SignatureKey, list[AggregatedSignatureProof]] = {} """ - Aggregated signature payloads learned from blocks. + Aggregated signature proofs learned from blocks. - - Keyed by (validator_id, attestation_data_root). - - Values are lists of (aggregation bits, payload) tuples so we know exactly which - validators signed. + - Keyed by SignatureKey(validator_id, attestation_data_root). + - Values are lists of AggregatedSignatureProof, each containing the participants + bitfield indicating which validators signed. - Used for recursive signature aggregation when building blocks. - Populated by on_block. """ @@ -320,7 +317,8 @@ def on_gossip_attestation( # Store signature for later lookup during block building new_gossip_sigs = dict(self.gossip_signatures) - new_gossip_sigs[(validator_id, attestation_data.data_root_bytes())] = signature + sig_key = SignatureKey(validator_id, Bytes32(attestation_data.data_root_bytes())) + new_gossip_sigs[sig_key] = signature # Process the attestation data store = self.on_attestation(attestation=attestation, is_from_block=False) @@ -561,24 +559,22 @@ def on_block( "Attestation signature groups must match aggregated attestations" ) - # Copy the aggregated signature payloads map for updates + # Copy the aggregated proof map for updates # Must deep copy the lists to maintain immutability of previous store snapshots - new_block_sigs: Dict[AttestationSignatureKey, AggregatedSignaturePayloads] = copy.deepcopy( + new_block_proofs: Dict[SignatureKey, list[AggregatedSignatureProof]] = copy.deepcopy( store.aggregated_payloads ) - for att, sig in zip(aggregated_attestations, attestation_signatures, strict=True): + for att, proof in zip(aggregated_attestations, attestation_signatures, strict=True): validator_ids = att.aggregation_bits.to_validator_indices() data_root = att.data.data_root_bytes() - # Pre-calculate the payload record once for this entire group - record = (AggregationBits.from_validator_indices(validator_ids), sig) - for vid in validator_ids: - # Update Signature Map + # Update Proof Map # - # Store the payload so future block builders can reuse this aggregation - new_block_sigs.setdefault((vid, data_root), []).append(record) + # Store the proof so future block builders can reuse this aggregation + key = SignatureKey(vid, Bytes32(data_root)) + new_block_proofs.setdefault(key, []).append(proof) # Update Fork Choice # @@ -588,8 +584,8 @@ def on_block( is_from_block=True, ) - # Update store with new aggregated signature payloads - store = store.model_copy(update={"aggregated_payloads": new_block_sigs}) + # Update store with new aggregated proofs + store = store.model_copy(update={"aggregated_payloads": new_block_proofs}) # Update forkchoice head based on new block and attestations # @@ -606,9 +602,12 @@ def on_block( # 3. Influence fork choice only after interval 3 (end of slot) # # We also store the proposer's signature for potential future block building. - proposer_data_root = proposer_attestation.data.data_root_bytes() + proposer_sig_key = SignatureKey( + proposer_attestation.validator_id, + Bytes32(proposer_attestation.data.data_root_bytes()), + ) new_gossip_sigs = dict(store.gossip_signatures) - new_gossip_sigs[(proposer_attestation.validator_id, proposer_data_root)] = ( + new_gossip_sigs[proposer_sig_key] = ( signed_block_with_attestation.signature.proposer_signature ) @@ -1037,7 +1036,7 @@ def produce_block_with_signatures( self, slot: Slot, validator_index: Uint64, - ) -> tuple["Store", Block, list[MultisigAggregatedSignature]]: + ) -> tuple["Store", Block, list[AggregatedSignatureProof]]: """ Produce a block and per-aggregated-attestation signature payloads for the target slot. diff --git a/src/lean_spec/subspecs/xmss/aggregation.py b/src/lean_spec/subspecs/xmss/aggregation.py index 9cbe813c..5f0d34b3 100644 --- a/src/lean_spec/subspecs/xmss/aggregation.py +++ b/src/lean_spec/subspecs/xmss/aggregation.py @@ -1,123 +1,128 @@ -""" -Multisig aggregation helpers bridging leanSpec containers to bindings. - -This module wraps the Python bindings exposed by the `leanMultisig-py` project to provide -Multisig signature aggregation + verification. -""" +"""Signature aggregation for the Lean Ethereum consensus specification.""" from __future__ import annotations -from typing import TYPE_CHECKING, Self, Sequence +from typing import NamedTuple, Self, Sequence -from lean_multisig_py import aggregate_signatures as aggregate_signatures_py -from lean_multisig_py import setup_prover, setup_verifier -from lean_multisig_py import verify_aggregated_signatures as verify_aggregated_signatures_py +from lean_multisig_py import aggregate_signatures as _aggregate_signatures_py +from lean_multisig_py import setup_prover as _setup_prover +from lean_multisig_py import setup_verifier as _setup_verifier +from lean_multisig_py import verify_aggregated_signatures as _verify_aggregated_signatures_py -from lean_spec.subspecs.xmss.containers import PublicKey -from lean_spec.subspecs.xmss.containers import Signature as XmssSignature -from lean_spec.types import Uint64 +from lean_spec.subspecs.containers.attestation import AggregationBits +from lean_spec.types import Bytes32, Uint64 from lean_spec.types.byte_arrays import ByteListMiB +from lean_spec.types.container import Container + +from .containers import PublicKey, Signature + + +class SignatureKey(NamedTuple): + """ + Key for looking up individual validator signatures. + + Used to index signature caches by (validator, message) pairs. + """ + + validator_id: Uint64 + """The validator who produced the signature.""" -if TYPE_CHECKING: - from lean_spec.subspecs.containers.attestation import AggregationBits + data_root: Bytes32 + """The hash of the signed data (e.g., attestation data root).""" -AttestationSignatureKey = tuple[Uint64, bytes] -"""Key type for looking up signatures: (validator id, attestation data root).""" +class AggregationError(Exception): + """Raised when signature aggregation or verification fails.""" -class MultisigError(RuntimeError): - """Base exception for multisig aggregation helpers.""" +class AggregatedSignatureProof(Container): + """ + Cryptographic proof that a set of validators signed a message. -class MultisigAggregationError(MultisigError): - """Raised when multisig fails to aggregate or verify signatures.""" + This container encapsulates the output of the leanVM signature aggregation, + combining the participant set with the proof bytes. This design ensures + the proof is self-describing: it carries information about which validators + it covers. + The proof can verify that all participants signed the same message in the + same epoch, using a single verification operation instead of checking + each signature individually. + """ -class MultisigAggregatedSignature(ByteListMiB): - """Variable-length byte list with a limit of 1048576 bytes.""" + participants: AggregationBits + """Bitfield indicating which validators' signatures are included.""" + + proof_data: ByteListMiB + """The raw aggregated proof bytes from leanVM.""" - # This function will change for recursive aggregation - # which might additionally require hints. @classmethod - def aggregate_signatures( + def aggregate( cls, + participants: AggregationBits, public_keys: Sequence[PublicKey], - signatures: Sequence[XmssSignature], + signatures: Sequence[Signature], message: bytes, epoch: Uint64, ) -> Self: """ - Aggregate XMSS signatures. + Aggregate individual XMSS signatures into a single proof. Args: - public_keys: Public keys of the signers, one per signature. + participants: Bitfield of validators whose signatures are included. + public_keys: Public keys of the signers (must match signatures order). signatures: Individual XMSS signatures to aggregate. message: The 32-byte message that was signed. epoch: The epoch in which the signatures were created. Returns: - The aggregated signature payload. + An aggregated signature proof covering all participants. Raises: - MultisigError: If lean-multisig-py is unavailable or aggregation fails. + AggregationError: If aggregation fails. """ - setup_prover() + _setup_prover() try: - pub_keys_bytes = [pk.encode_bytes() for pk in public_keys] - sig_bytes = [sig.encode_bytes() for sig in signatures] - - # In test mode, we return a single zero byte payload. - # TODO: Remove test mode once leanVM is supports correct signature encoding. - aggregated_bytes = aggregate_signatures_py( - pub_keys_bytes, - sig_bytes, + # TODO: Remove test_mode once leanVM supports correct signature encoding. + proof_bytes = _aggregate_signatures_py( + [pk.encode_bytes() for pk in public_keys], + [sig.encode_bytes() for sig in signatures], message, epoch, test_mode=True, ) - return cls(data=aggregated_bytes) + return cls( + participants=participants, + proof_data=ByteListMiB(data=proof_bytes), + ) except Exception as exc: - raise MultisigAggregationError(f"Multisig aggregation failed: {exc}") from exc + raise AggregationError(f"Signature aggregation failed: {exc}") from exc - # This function will change for recursive aggregation verification - # which might additionally require hints. - def verify_aggregated_payload( + def verify( self, public_keys: Sequence[PublicKey], message: bytes, epoch: Uint64, ) -> None: """ - Verify a lean-multisig-py aggregated signature payload. + Verify this aggregated signature proof. Args: - public_keys: Public keys of the signers, one per original signature. - payload: MultisigAggregatedSignature of the aggregated signature payload. + public_keys: Public keys of the participants (order must match participants bitfield). message: The 32-byte message that was signed. epoch: The epoch in which the signatures were created. Raises: - MultisigError: If lean-multisig-py is unavailable or verification fails. + AggregationError: If verification fails. """ - setup_verifier() + _setup_verifier() try: - pub_keys_bytes = [pk.encode_bytes() for pk in public_keys] - - # In test mode, we allow verification of a single zero byte payload. - # TODO: Remove test mode once leanVM is supports correct signature encoding. - verify_aggregated_signatures_py( - pub_keys_bytes, + # TODO: Remove test_mode once leanVM supports correct signature encoding. + _verify_aggregated_signatures_py( + [pk.encode_bytes() for pk in public_keys], message, - self.encode_bytes(), + self.proof_data.encode_bytes(), int(epoch), test_mode=True, ) except Exception as exc: - raise MultisigAggregationError(f"Multisig verification failed: {exc}") from exc - - -AggregatedSignaturePayload = tuple["AggregationBits", "MultisigAggregatedSignature"] -"""Aggregated signature payload with its participant bitlist.""" - -AggregatedSignaturePayloads = list[AggregatedSignaturePayload] -"""List of aggregated signature payloads with their participant bitlists.""" + raise AggregationError(f"Signature verification failed: {exc}") from exc diff --git a/src/lean_spec/types/__init__.py b/src/lean_spec/types/__init__.py index 0182d312..6dcbb36f 100644 --- a/src/lean_spec/types/__init__.py +++ b/src/lean_spec/types/__init__.py @@ -2,6 +2,7 @@ from .base import CamelModel, StrictBaseModel from .basispt import BasisPoint +from .bitfields import BaseBitlist from .boolean import Boolean from .byte_arrays import ZERO_HASH, Bytes32, Bytes52, Bytes3116 from .collections import SSZList, SSZVector @@ -18,6 +19,7 @@ __all__ = [ # Core types + "BaseBitlist", "Uint64", "BasisPoint", "Bytes32", diff --git a/tests/lean_spec/subspecs/containers/test_state_aggregation.py b/tests/lean_spec/subspecs/containers/test_state_aggregation.py index 2b337e88..60f8c2da 100644 --- a/tests/lean_spec/subspecs/containers/test_state_aggregation.py +++ b/tests/lean_spec/subspecs/containers/test_state_aggregation.py @@ -17,7 +17,7 @@ from lean_spec.subspecs.containers.validator import Validator from lean_spec.subspecs.koalabear import Fp from lean_spec.subspecs.ssz.hash import hash_tree_root -from lean_spec.subspecs.xmss.aggregation import MultisigAggregatedSignature +from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof, SignatureKey from lean_spec.subspecs.xmss.containers import PublicKey, Signature from lean_spec.subspecs.xmss.types import ( HashDigestList, @@ -27,8 +27,19 @@ Randomness, ) from lean_spec.types import Bytes32, Bytes52, Uint64 +from lean_spec.types.byte_arrays import ByteListMiB -TEST_AGGREGATED_SIGNATURE = MultisigAggregatedSignature(data=b"\x00") + +def make_test_proof(validator_ids: list[Uint64], data: bytes = b"\x00") -> AggregatedSignatureProof: + """Create a test AggregatedSignatureProof with given participants.""" + return AggregatedSignatureProof( + participants=AggregationBits.from_validator_indices(validator_ids), + proof_data=ByteListMiB(data=data), + ) + + +# Default test proof with empty participants (will be replaced in most tests) +TEST_AGGREGATED_PROOF = make_test_proof([Uint64(0), Uint64(1)]) def make_bytes32(seed: int) -> Bytes32: @@ -100,8 +111,8 @@ def test_gossip_aggregation_succeeds_with_all_signatures() -> None: data_root = b"\x11" * 32 validator_ids = [Uint64(0), Uint64(1)] gossip_signatures = { - (Uint64(0), data_root): make_signature(0), - (Uint64(1), data_root): make_signature(1), + SignatureKey(Uint64(0), Bytes32(data_root)): make_signature(0), + SignatureKey(Uint64(1), Bytes32(data_root)): make_signature(1), } result = state._aggregate_signatures_from_gossip( @@ -112,16 +123,15 @@ def test_gossip_aggregation_succeeds_with_all_signatures() -> None: ) assert result is not None - aggregated_signature, aggregated_bitlist, remaining = result - assert aggregated_signature == TEST_AGGREGATED_SIGNATURE - assert set(aggregated_bitlist.to_validator_indices()) == set(validator_ids) + proof, remaining = result + assert set(proof.participants.to_validator_indices()) == set(validator_ids) assert remaining == set() def test_gossip_aggregation_returns_partial_result_when_some_missing() -> None: state = make_state(2) data_root = b"\x22" * 32 - gossip_signatures = {(Uint64(0), data_root): make_signature(0)} + gossip_signatures = {SignatureKey(Uint64(0), Bytes32(data_root)): make_signature(0)} result = state._aggregate_signatures_from_gossip( [Uint64(0), Uint64(1)], @@ -131,9 +141,8 @@ def test_gossip_aggregation_returns_partial_result_when_some_missing() -> None: ) assert result is not None - aggregated_signature, aggregated_bitlist, remaining = result - assert aggregated_signature == TEST_AGGREGATED_SIGNATURE - assert aggregated_bitlist.to_validator_indices() == [Uint64(0)] + proof, remaining = result + assert proof.participants.to_validator_indices() == [Uint64(0)] assert remaining == {Uint64(1)} @@ -141,7 +150,7 @@ def test_gossip_aggregation_returns_none_if_no_signature_matches() -> None: state = make_state(2) data_root = b"\x33" * 32 # Gossip data exists but for a different validator key, so no signatures match - gossip_signatures = {(Uint64(9), data_root): make_signature(0)} + gossip_signatures = {SignatureKey(Uint64(9), Bytes32(data_root)): make_signature(0)} result = state._aggregate_signatures_from_gossip( [Uint64(0), Uint64(1)], @@ -153,96 +162,84 @@ def test_gossip_aggregation_returns_none_if_no_signature_matches() -> None: assert result is None -def test_pick_from_aggregated_signatures_prefers_widest_overlap() -> None: +def test_pick_from_aggregated_proofs_prefers_widest_overlap() -> None: state = make_state(3) data_root = b"\x44" * 32 remaining_validator_ids = {Uint64(0), Uint64(1)} - narrow_bits = AggregationBits.from_validator_indices([Uint64(0)]) - best_bits = AggregationBits.from_validator_indices([Uint64(0), Uint64(1)]) - narrow_signature = MultisigAggregatedSignature(data=b"narrow") - best_signature = MultisigAggregatedSignature(data=b"best") + narrow_proof = make_test_proof([Uint64(0)], b"narrow") + best_proof = make_test_proof([Uint64(0), Uint64(1)], b"best") aggregated_payloads = { - (Uint64(0), data_root): [ - (narrow_bits, narrow_signature), - (best_bits, best_signature), - ], - (Uint64(1), data_root): [ - (best_bits, best_signature), - (narrow_bits, narrow_signature), - ], + SignatureKey(Uint64(0), Bytes32(data_root)): [narrow_proof, best_proof], + SignatureKey(Uint64(1), Bytes32(data_root)): [best_proof, narrow_proof], } - signature, bitlist, remaining = state._pick_from_aggregated_signatures( + proof, remaining = state._pick_from_aggregated_proofs( remaining_validator_ids=remaining_validator_ids, data_root=data_root, aggregated_payloads=aggregated_payloads, ) - assert signature == best_signature - assert set(bitlist.to_validator_indices()) == {Uint64(0), Uint64(1)} + assert set(proof.participants.to_validator_indices()) == {Uint64(0), Uint64(1)} assert remaining == set() -def test_pick_from_aggregated_signatures_returns_remaining_for_partial_payload() -> None: +def test_pick_from_aggregated_proofs_returns_remaining_for_partial_payload() -> None: state = make_state(2) data_root = b"\x45" * 32 remaining_validator_ids = {Uint64(0), Uint64(1)} - partial_bits_0 = AggregationBits.from_validator_indices([Uint64(0)]) - partial_bits_1 = AggregationBits.from_validator_indices([Uint64(1)]) - partial_signature_0 = MultisigAggregatedSignature(data=b"partial-0") - partial_signature_1 = MultisigAggregatedSignature(data=b"partial-1") + partial_proof_0 = make_test_proof([Uint64(0)], b"partial-0") + partial_proof_1 = make_test_proof([Uint64(1)], b"partial-1") aggregated_payloads = { - (Uint64(0), data_root): [(partial_bits_0, partial_signature_0)], - (Uint64(1), data_root): [(partial_bits_1, partial_signature_1)], + SignatureKey(Uint64(0), Bytes32(data_root)): [partial_proof_0], + SignatureKey(Uint64(1), Bytes32(data_root)): [partial_proof_1], } - signature, bitlist, remaining = state._pick_from_aggregated_signatures( + proof, remaining = state._pick_from_aggregated_proofs( remaining_validator_ids=remaining_validator_ids, data_root=data_root, aggregated_payloads=aggregated_payloads, ) - covered_validators = set(bitlist.to_validator_indices()) + covered_validators = set(proof.participants.to_validator_indices()) assert covered_validators <= {Uint64(0), Uint64(1)} - assert signature in {partial_signature_0, partial_signature_1} assert remaining == remaining_validator_ids - covered_validators -def test_pick_from_aggregated_signatures_requires_payloads() -> None: +def test_pick_from_aggregated_proofs_requires_payloads() -> None: state = make_state(1) - with pytest.raises(ValueError, match="aggregated payloads is required"): - state._pick_from_aggregated_signatures( + with pytest.raises(ValueError, match="aggregated payloads required"): + state._pick_from_aggregated_proofs( remaining_validator_ids={Uint64(0)}, data_root=b"\x55" * 32, aggregated_payloads=None, ) -def test_pick_from_aggregated_signatures_errors_on_empty_remaining() -> None: +def test_pick_from_aggregated_proofs_errors_on_empty_remaining() -> None: state = make_state(1) with pytest.raises(ValueError, match="remaining validator ids cannot be empty"): - state._pick_from_aggregated_signatures( + state._pick_from_aggregated_proofs( remaining_validator_ids=set(), data_root=b"\x66" * 32, aggregated_payloads={}, ) -def test_pick_from_aggregated_signatures_errors_when_no_candidates() -> None: +def test_pick_from_aggregated_proofs_errors_when_no_candidates() -> None: state = make_state(1) data_root = b"\x77" * 32 - with pytest.raises(ValueError, match="Failed to locate an aggregated signature payload"): - state._pick_from_aggregated_signatures( + with pytest.raises(ValueError, match="Failed to locate aggregated proof"): + state._pick_from_aggregated_proofs( remaining_validator_ids={Uint64(0)}, data_root=data_root, - aggregated_payloads={(Uint64(0), data_root): []}, + aggregated_payloads={SignatureKey(Uint64(0), Bytes32(data_root)): []}, ) @@ -252,15 +249,18 @@ def test_compute_aggregated_signatures_prefers_full_gossip_payload() -> None: att_data = make_attestation_data(2, make_bytes32(3), make_bytes32(4), source=source) attestations = [Attestation(validator_id=Uint64(i), data=att_data) for i in range(2)] data_root = att_data.data_root_bytes() - gossip_signatures = {(Uint64(i), data_root): make_signature(i) for i in range(2)} + gossip_signatures = { + SignatureKey(Uint64(i), Bytes32(data_root)): make_signature(i) for i in range(2) + } - aggregated_atts, aggregated_sigs = state.compute_aggregated_signatures( + aggregated_atts, aggregated_proofs = state.compute_aggregated_signatures( attestations, gossip_signatures=gossip_signatures, ) assert len(aggregated_atts) == 1 - assert aggregated_sigs == [TEST_AGGREGATED_SIGNATURE] + assert len(aggregated_proofs) == 1 + assert set(aggregated_proofs[0].participants.to_validator_indices()) == {Uint64(0), Uint64(1)} def test_compute_aggregated_signatures_splits_when_needed() -> None: @@ -269,16 +269,15 @@ def test_compute_aggregated_signatures_splits_when_needed() -> None: att_data = make_attestation_data(3, make_bytes32(5), make_bytes32(6), source=source) attestations = [Attestation(validator_id=Uint64(i), data=att_data) for i in range(3)] data_root = att_data.data_root_bytes() - gossip_signatures = {(Uint64(0), data_root): make_signature(0)} + gossip_signatures = {SignatureKey(Uint64(0), Bytes32(data_root)): make_signature(0)} - block_bits = AggregationBits.from_validator_indices([Uint64(1), Uint64(2)]) - block_signature = MultisigAggregatedSignature(data=b"block-12") + block_proof = make_test_proof([Uint64(1), Uint64(2)], b"block-12") aggregated_payloads = { - (Uint64(1), data_root): [(block_bits, block_signature)], - (Uint64(2), data_root): [(block_bits, block_signature)], + SignatureKey(Uint64(1), Bytes32(data_root)): [block_proof], + SignatureKey(Uint64(2), Bytes32(data_root)): [block_proof], } - aggregated_atts, aggregated_sigs = state.compute_aggregated_signatures( + aggregated_atts, aggregated_proofs = state.compute_aggregated_signatures( attestations, gossip_signatures=gossip_signatures, aggregated_payloads=aggregated_payloads, @@ -290,8 +289,12 @@ def test_compute_aggregated_signatures_splits_when_needed() -> None: ] assert (0,) in seen_participants assert (1, 2) in seen_participants - assert block_signature in aggregated_sigs - assert TEST_AGGREGATED_SIGNATURE in aggregated_sigs + # Check we have both proofs + proof_participants = [ + tuple(int(v) for v in p.participants.to_validator_indices()) for p in aggregated_proofs + ] + assert (0,) in proof_participants + assert (1, 2) in proof_participants def test_build_block_collects_valid_available_attestations() -> None: @@ -314,10 +317,10 @@ def test_build_block_collects_valid_available_attestations() -> None: attestation = Attestation(validator_id=Uint64(0), data=att_data) data_root = att_data.data_root_bytes() - gossip_signatures = {(Uint64(0), data_root): make_signature(0)} + gossip_signatures = {SignatureKey(Uint64(0), Bytes32(data_root)): make_signature(0)} # Proposer for slot 1 with 2 validators: slot % num_validators = 1 % 2 = 1 - block, post_state, aggregated_atts, aggregated_sigs = state.build_block( + block, post_state, aggregated_atts, aggregated_proofs = state.build_block( slot=Slot(1), proposer_index=Uint64(1), parent_root=parent_root, @@ -330,7 +333,8 @@ def test_build_block_collects_valid_available_attestations() -> None: assert post_state.latest_block_header.slot == Slot(1) assert list(block.body.attestations.data) == aggregated_atts - assert aggregated_sigs == [TEST_AGGREGATED_SIGNATURE] + assert len(aggregated_proofs) == 1 + assert aggregated_proofs[0].participants.to_validator_indices() == [Uint64(0)] assert block.body.attestations.data[0].aggregation_bits.to_validator_indices() == [Uint64(0)] @@ -354,7 +358,7 @@ def test_build_block_skips_attestations_without_signatures() -> None: attestation = Attestation(validator_id=Uint64(0), data=att_data) # Proposer for slot 1 with 1 validator: slot % num_validators = 1 % 1 = 0 - block, post_state, aggregated_atts, aggregated_sigs = state.build_block( + block, post_state, aggregated_atts, aggregated_proofs = state.build_block( slot=Slot(1), proposer_index=Uint64(0), parent_root=parent_root, @@ -367,7 +371,7 @@ def test_build_block_skips_attestations_without_signatures() -> None: assert post_state.latest_block_header.slot == Slot(1) assert aggregated_atts == [] - assert aggregated_sigs == [] + assert aggregated_proofs == [] assert list(block.body.attestations.data) == [] @@ -375,7 +379,7 @@ def test_gossip_aggregation_with_empty_validator_list() -> None: """Empty validator list should return None.""" state = make_state(2) data_root = b"\x99" * 32 - gossip_signatures = {(Uint64(0), data_root): make_signature(0)} + gossip_signatures = {SignatureKey(Uint64(0), Bytes32(data_root)): make_signature(0)} result = state._aggregate_signatures_from_gossip( [], # empty validator list @@ -449,20 +453,20 @@ def test_compute_aggregated_signatures_with_multiple_data_groups() -> None: data_root2 = att_data2.data_root_bytes() gossip_signatures = { - (Uint64(0), data_root1): make_signature(0), - (Uint64(1), data_root1): make_signature(1), - (Uint64(2), data_root2): make_signature(2), - (Uint64(3), data_root2): make_signature(3), + SignatureKey(Uint64(0), Bytes32(data_root1)): make_signature(0), + SignatureKey(Uint64(1), Bytes32(data_root1)): make_signature(1), + SignatureKey(Uint64(2), Bytes32(data_root2)): make_signature(2), + SignatureKey(Uint64(3), Bytes32(data_root2)): make_signature(3), } - aggregated_atts, aggregated_sigs = state.compute_aggregated_signatures( + aggregated_atts, aggregated_proofs = state.compute_aggregated_signatures( attestations, gossip_signatures=gossip_signatures, ) # Should have 2 aggregated attestations (one per data group) assert len(aggregated_atts) == 2 - assert len(aggregated_sigs) == 2 + assert len(aggregated_proofs) == 2 def test_compute_aggregated_signatures_falls_back_to_block_payload() -> None: @@ -474,24 +478,25 @@ def test_compute_aggregated_signatures_falls_back_to_block_payload() -> None: data_root = att_data.data_root_bytes() # Only gossip signature for validator 0 (incomplete) - gossip_signatures = {(Uint64(0), data_root): make_signature(0)} + gossip_signatures = {SignatureKey(Uint64(0), Bytes32(data_root)): make_signature(0)} # Block payload covers both validators - block_bits = AggregationBits.from_validator_indices([Uint64(0), Uint64(1)]) - block_signature = MultisigAggregatedSignature(data=b"block-fallback") + block_proof = make_test_proof([Uint64(0), Uint64(1)], b"block-fallback") aggregated_payloads = { - (Uint64(0), data_root): [(block_bits, block_signature)], - (Uint64(1), data_root): [(block_bits, block_signature)], + SignatureKey(Uint64(0), Bytes32(data_root)): [block_proof], + SignatureKey(Uint64(1), Bytes32(data_root)): [block_proof], } - aggregated_atts, aggregated_sigs = state.compute_aggregated_signatures( + aggregated_atts, aggregated_proofs = state.compute_aggregated_signatures( attestations, gossip_signatures=gossip_signatures, aggregated_payloads=aggregated_payloads, ) - # Should include both gossip-covered and fallback payload attestations/signatures + # Should include both gossip-covered and fallback payload attestations/proofs assert len(aggregated_atts) == 2 - assert len(aggregated_sigs) == 2 - assert block_signature in aggregated_sigs - assert TEST_AGGREGATED_SIGNATURE in aggregated_sigs + assert len(aggregated_proofs) == 2 + # Check we have one proof with validator 0 and one proof with both validators + proof_participants = [set(p.participants.to_validator_indices()) for p in aggregated_proofs] + assert {Uint64(0)} in proof_participants + assert {Uint64(0), Uint64(1)} in proof_participants diff --git a/tests/lean_spec/subspecs/forkchoice/test_store_attestations.py b/tests/lean_spec/subspecs/forkchoice/test_store_attestations.py index c96f87fb..57e03ab6 100644 --- a/tests/lean_spec/subspecs/forkchoice/test_store_attestations.py +++ b/tests/lean_spec/subspecs/forkchoice/test_store_attestations.py @@ -23,6 +23,7 @@ from lean_spec.subspecs.containers.validator import Validator from lean_spec.subspecs.forkchoice import Store from lean_spec.subspecs.ssz.hash import hash_tree_root +from lean_spec.subspecs.xmss.aggregation import SignatureKey from lean_spec.types import Bytes32, Bytes52, Uint64 @@ -59,7 +60,9 @@ def test_on_block_processes_multi_validator_aggregations() -> None: # Store signatures in gossip_signatures data_root = attestation_data.data_root_bytes() gossip_sigs = { - (validator_id, data_root): key_manager.sign_attestation_data(validator_id, attestation_data) + SignatureKey(validator_id, Bytes32(data_root)): key_manager.sign_attestation_data( + validator_id, attestation_data + ) for validator_id in (Uint64(1), Uint64(2)) } @@ -148,7 +151,7 @@ def test_on_block_preserves_immutability_of_aggregated_payloads() -> None: validator_id: attestation_data_1 for validator_id in (Uint64(1), Uint64(2)) } gossip_sigs_1 = { - (validator_id, data_root_1): key_manager.sign_attestation_data( + SignatureKey(validator_id, Bytes32(data_root_1)): key_manager.sign_attestation_data( validator_id, attestation_data_1 ) for validator_id in (Uint64(1), Uint64(2)) @@ -214,7 +217,7 @@ def test_on_block_preserves_immutability_of_aggregated_payloads() -> None: validator_id: attestation_data_2 for validator_id in (Uint64(1), Uint64(2)) } gossip_sigs_2 = { - (validator_id, data_root_2): key_manager.sign_attestation_data( + SignatureKey(validator_id, Bytes32(data_root_2)): key_manager.sign_attestation_data( validator_id, attestation_data_2 ) for validator_id in (Uint64(1), Uint64(2)) diff --git a/tests/lean_spec/subspecs/forkchoice/test_validator.py b/tests/lean_spec/subspecs/forkchoice/test_validator.py index afbe2a20..bd6ed349 100644 --- a/tests/lean_spec/subspecs/forkchoice/test_validator.py +++ b/tests/lean_spec/subspecs/forkchoice/test_validator.py @@ -26,6 +26,7 @@ from lean_spec.subspecs.forkchoice import Store from lean_spec.subspecs.koalabear import Fp from lean_spec.subspecs.ssz.hash import hash_tree_root +from lean_spec.subspecs.xmss.aggregation import SignatureKey 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 @@ -195,12 +196,10 @@ def test_produce_block_with_attestations(self, sample_store: Store) -> None: ) sample_store.latest_known_attestations[Uint64(5)] = signed_5.message sample_store.latest_known_attestations[Uint64(6)] = signed_6.message - sample_store.gossip_signatures[(Uint64(5), signed_5.message.data_root_bytes())] = ( - signed_5.signature - ) - sample_store.gossip_signatures[(Uint64(6), signed_6.message.data_root_bytes())] = ( - signed_6.signature - ) + sig_key_5 = SignatureKey(Uint64(5), Bytes32(signed_5.message.data_root_bytes())) + sig_key_6 = SignatureKey(Uint64(6), Bytes32(signed_6.message.data_root_bytes())) + sample_store.gossip_signatures[sig_key_5] = signed_5.signature + sample_store.gossip_signatures[sig_key_6] = signed_6.signature slot = Slot(2) validator_idx = Uint64(2) # Proposer for slot 2 @@ -291,9 +290,8 @@ def test_produce_block_state_consistency(self, sample_store: Store) -> None: target=sample_store.get_attestation_target(), ) sample_store.latest_known_attestations[Uint64(7)] = signed_7.message - sample_store.gossip_signatures[(Uint64(7), signed_7.message.data_root_bytes())] = ( - signed_7.signature - ) + sig_key_7 = SignatureKey(Uint64(7), Bytes32(signed_7.message.data_root_bytes())) + sample_store.gossip_signatures[sig_key_7] = signed_7.signature store, block, _signatures = sample_store.produce_block_with_signatures( slot, From 1a1f5cbd571910257ecf4c88137a1b136856286a Mon Sep 17 00:00:00 2001 From: Thomas Coratger Date: Mon, 29 Dec 2025 23:02:51 +0100 Subject: [PATCH 2/5] cleanup --- packages/testing/src/consensus_testing/keys.py | 2 +- src/lean_spec/subspecs/containers/state/state.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/testing/src/consensus_testing/keys.py b/packages/testing/src/consensus_testing/keys.py index fe35fbdc..b60e5794 100755 --- a/packages/testing/src/consensus_testing/keys.py +++ b/packages/testing/src/consensus_testing/keys.py @@ -310,7 +310,7 @@ def build_attestation_signatures( public_keys=public_keys, signatures=signatures, message=message, - epoch=Uint64(epoch), + epoch=epoch, ) proofs.append(proof) diff --git a/src/lean_spec/subspecs/containers/state/state.py b/src/lean_spec/subspecs/containers/state/state.py index 01815d08..297874b8 100644 --- a/src/lean_spec/subspecs/containers/state/state.py +++ b/src/lean_spec/subspecs/containers/state/state.py @@ -655,7 +655,7 @@ def _aggregate_signatures_from_gossip( public_keys=public_keys, signatures=signatures, message=data_root, - epoch=Uint64(epoch), + epoch=epoch, ) return proof, missing_validator_ids From 4a2ffbf15ac90653d21aceaca50d21e99823cc9e Mon Sep 17 00:00:00 2001 From: Thomas Coratger Date: Mon, 29 Dec 2025 23:07:44 +0100 Subject: [PATCH 3/5] cleanup --- src/lean_spec/subspecs/containers/state/state.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/lean_spec/subspecs/containers/state/state.py b/src/lean_spec/subspecs/containers/state/state.py index 297874b8..24812121 100644 --- a/src/lean_spec/subspecs/containers/state/state.py +++ b/src/lean_spec/subspecs/containers/state/state.py @@ -637,6 +637,7 @@ def _aggregate_signatures_from_gossip( missing_validator_ids: set[Uint64] = set() for validator_index in validator_ids: + # Attempt to retrieve the signature; fail fast if any are missing. key = SignatureKey(validator_index, Bytes32(data_root)) if (sig := gossip_signatures.get(key)) is None: missing_validator_ids.add(validator_index) From d2fe1727e265eff87f591b79d36962abb7e89bed Mon Sep 17 00:00:00 2001 From: Thomas Coratger Date: Tue, 30 Dec 2025 19:26:26 +0100 Subject: [PATCH 4/5] fix Bytes32 --- .../testing/src/consensus_testing/keys.py | 4 +-- .../test_fixtures/fork_choice.py | 4 +-- .../test_fixtures/state_transition.py | 2 +- .../test_fixtures/verify_signatures.py | 2 +- .../containers/attestation/attestation.py | 6 ++--- .../subspecs/containers/state/state.py | 10 +++---- src/lean_spec/subspecs/forkchoice/store.py | 6 ++--- .../containers/test_state_aggregation.py | 26 +++++++++---------- .../forkchoice/test_store_attestations.py | 6 ++--- .../subspecs/forkchoice/test_validator.py | 6 ++--- 10 files changed, 34 insertions(+), 38 deletions(-) diff --git a/packages/testing/src/consensus_testing/keys.py b/packages/testing/src/consensus_testing/keys.py index b60e5794..bbbed28c 100755 --- a/packages/testing/src/consensus_testing/keys.py +++ b/packages/testing/src/consensus_testing/keys.py @@ -55,7 +55,7 @@ TEST_SIGNATURE_SCHEME, GeneralizedXmssScheme, ) -from lean_spec.types import Bytes32, Uint64 +from lean_spec.types import Uint64 if TYPE_CHECKING: from collections.abc import Mapping @@ -296,7 +296,7 @@ def build_attestation_signatures( public_keys: list[PublicKey] = [self.get_public_key(vid) for vid in validator_ids] signatures: list[Signature] = [ ( - lookup.get(SignatureKey(vid, Bytes32(message))) + lookup.get(SignatureKey(vid, message)) or self.sign_attestation_data(vid, agg.data) ) for vid in validator_ids diff --git a/packages/testing/src/consensus_testing/test_fixtures/fork_choice.py b/packages/testing/src/consensus_testing/test_fixtures/fork_choice.py index 97ace502..fdf9eff3 100644 --- a/packages/testing/src/consensus_testing/test_fixtures/fork_choice.py +++ b/packages/testing/src/consensus_testing/test_fixtures/fork_choice.py @@ -423,9 +423,7 @@ def _build_attestations_from_spec( attestation = Attestation(validator_id=signed_att.validator_id, data=signed_att.message) attestations.append(attestation) - sig_key = SignatureKey( - attestation.validator_id, Bytes32(attestation.data.data_root_bytes()) - ) + sig_key = SignatureKey(attestation.validator_id, attestation.data.data_root_bytes()) signature_lookup[sig_key] = signed_att.signature return attestations, signature_lookup diff --git a/packages/testing/src/consensus_testing/test_fixtures/state_transition.py b/packages/testing/src/consensus_testing/test_fixtures/state_transition.py index c77389d4..7241f861 100644 --- a/packages/testing/src/consensus_testing/test_fixtures/state_transition.py +++ b/packages/testing/src/consensus_testing/test_fixtures/state_transition.py @@ -264,7 +264,7 @@ def _build_block_from_spec(self, spec: BlockSpec, state: State) -> tuple[Block, key_manager = get_shared_key_manager(max_slot=spec.slot) gossip_signatures = { SignatureKey( - att.validator_id, Bytes32(att.data.data_root_bytes()) + att.validator_id, att.data.data_root_bytes() ): key_manager.sign_attestation_data( att.validator_id, att.data, diff --git a/packages/testing/src/consensus_testing/test_fixtures/verify_signatures.py b/packages/testing/src/consensus_testing/test_fixtures/verify_signatures.py index 4ca7ef89..6a5c66b2 100644 --- a/packages/testing/src/consensus_testing/test_fixtures/verify_signatures.py +++ b/packages/testing/src/consensus_testing/test_fixtures/verify_signatures.py @@ -188,7 +188,7 @@ def _build_block_from_spec( # 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 = { - SignatureKey(att.validator_id, Bytes32(att.data.data_root_bytes())): sig + SignatureKey(att.validator_id, att.data.data_root_bytes()): sig for att, sig in zip(attestations, attestation_signature_inputs, strict=True) } diff --git a/src/lean_spec/subspecs/containers/attestation/attestation.py b/src/lean_spec/subspecs/containers/attestation/attestation.py index 18080a5f..1a0e7fb6 100644 --- a/src/lean_spec/subspecs/containers/attestation/attestation.py +++ b/src/lean_spec/subspecs/containers/attestation/attestation.py @@ -18,7 +18,7 @@ from lean_spec.subspecs.containers.slot import Slot from lean_spec.subspecs.ssz import hash_tree_root -from lean_spec.types import Container, Uint64 +from lean_spec.types import Bytes32, Container, Uint64 from ...xmss.containers import Signature from ..checkpoint import Checkpoint @@ -40,9 +40,9 @@ class AttestationData(Container): source: Checkpoint """The checkpoint representing the source block as observed by the validator.""" - def data_root_bytes(self) -> bytes: + def data_root_bytes(self) -> Bytes32: """The root of the attestation data.""" - return bytes(hash_tree_root(self)) + return hash_tree_root(self) class Attestation(Container): diff --git a/src/lean_spec/subspecs/containers/state/state.py b/src/lean_spec/subspecs/containers/state/state.py index 24812121..cd700638 100644 --- a/src/lean_spec/subspecs/containers/state/state.py +++ b/src/lean_spec/subspecs/containers/state/state.py @@ -613,7 +613,7 @@ def state_transition(self, block: Block, valid_signatures: bool = True) -> "Stat def _aggregate_signatures_from_gossip( self, validator_ids: list[Uint64], - data_root: bytes, + data_root: Bytes32, epoch: Slot, gossip_signatures: dict[SignatureKey, "Signature"] | None = None, ) -> tuple[AggregatedSignatureProof, set[Uint64]] | None: @@ -638,7 +638,7 @@ def _aggregate_signatures_from_gossip( for validator_index in validator_ids: # Attempt to retrieve the signature; fail fast if any are missing. - key = SignatureKey(validator_index, Bytes32(data_root)) + key = SignatureKey(validator_index, data_root) if (sig := gossip_signatures.get(key)) is None: missing_validator_ids.add(validator_index) continue @@ -733,7 +733,7 @@ def build_block( data = attestation.data validator_id = attestation.validator_id data_root = data.data_root_bytes() - sig_key = SignatureKey(validator_id, Bytes32(data_root)) + sig_key = SignatureKey(validator_id, data_root) # Skip if target block is unknown if data.head.root not in known_block_roots: @@ -873,7 +873,7 @@ def compute_aggregated_signatures( def _pick_from_aggregated_proofs( self, remaining_validator_ids: set[Uint64], - data_root: bytes, + data_root: Bytes32, aggregated_payloads: dict[SignatureKey, list[AggregatedSignatureProof]] | None = None, ) -> tuple[AggregatedSignatureProof, set[Uint64]]: """ @@ -901,7 +901,7 @@ def _pick_from_aggregated_proofs( best_remaining: set[Uint64] = set() representative_validator_id = next(iter(remaining_validator_ids)) - key = SignatureKey(representative_validator_id, Bytes32(data_root)) + key = SignatureKey(representative_validator_id, data_root) for proof in aggregated_payloads.get(key, []): participants = set(proof.participants.to_validator_indices()) diff --git a/src/lean_spec/subspecs/forkchoice/store.py b/src/lean_spec/subspecs/forkchoice/store.py index ac987401..c0a93e88 100644 --- a/src/lean_spec/subspecs/forkchoice/store.py +++ b/src/lean_spec/subspecs/forkchoice/store.py @@ -317,7 +317,7 @@ def on_gossip_attestation( # Store signature for later lookup during block building new_gossip_sigs = dict(self.gossip_signatures) - sig_key = SignatureKey(validator_id, Bytes32(attestation_data.data_root_bytes())) + sig_key = SignatureKey(validator_id, attestation_data.data_root_bytes()) new_gossip_sigs[sig_key] = signature # Process the attestation data @@ -573,7 +573,7 @@ def on_block( # Update Proof Map # # Store the proof so future block builders can reuse this aggregation - key = SignatureKey(vid, Bytes32(data_root)) + key = SignatureKey(vid, data_root) new_block_proofs.setdefault(key, []).append(proof) # Update Fork Choice @@ -604,7 +604,7 @@ def on_block( # We also store the proposer's signature for potential future block building. proposer_sig_key = SignatureKey( proposer_attestation.validator_id, - Bytes32(proposer_attestation.data.data_root_bytes()), + proposer_attestation.data.data_root_bytes(), ) new_gossip_sigs = dict(store.gossip_signatures) new_gossip_sigs[proposer_sig_key] = ( diff --git a/tests/lean_spec/subspecs/containers/test_state_aggregation.py b/tests/lean_spec/subspecs/containers/test_state_aggregation.py index 60f8c2da..988ce875 100644 --- a/tests/lean_spec/subspecs/containers/test_state_aggregation.py +++ b/tests/lean_spec/subspecs/containers/test_state_aggregation.py @@ -249,9 +249,7 @@ def test_compute_aggregated_signatures_prefers_full_gossip_payload() -> None: att_data = make_attestation_data(2, make_bytes32(3), make_bytes32(4), source=source) attestations = [Attestation(validator_id=Uint64(i), data=att_data) for i in range(2)] data_root = att_data.data_root_bytes() - gossip_signatures = { - SignatureKey(Uint64(i), Bytes32(data_root)): make_signature(i) for i in range(2) - } + gossip_signatures = {SignatureKey(Uint64(i), data_root): make_signature(i) for i in range(2)} aggregated_atts, aggregated_proofs = state.compute_aggregated_signatures( attestations, @@ -269,12 +267,12 @@ def test_compute_aggregated_signatures_splits_when_needed() -> None: att_data = make_attestation_data(3, make_bytes32(5), make_bytes32(6), source=source) attestations = [Attestation(validator_id=Uint64(i), data=att_data) for i in range(3)] data_root = att_data.data_root_bytes() - gossip_signatures = {SignatureKey(Uint64(0), Bytes32(data_root)): make_signature(0)} + gossip_signatures = {SignatureKey(Uint64(0), data_root): make_signature(0)} block_proof = make_test_proof([Uint64(1), Uint64(2)], b"block-12") aggregated_payloads = { - SignatureKey(Uint64(1), Bytes32(data_root)): [block_proof], - SignatureKey(Uint64(2), Bytes32(data_root)): [block_proof], + SignatureKey(Uint64(1), data_root): [block_proof], + SignatureKey(Uint64(2), data_root): [block_proof], } aggregated_atts, aggregated_proofs = state.compute_aggregated_signatures( @@ -317,7 +315,7 @@ def test_build_block_collects_valid_available_attestations() -> None: attestation = Attestation(validator_id=Uint64(0), data=att_data) data_root = att_data.data_root_bytes() - gossip_signatures = {SignatureKey(Uint64(0), Bytes32(data_root)): make_signature(0)} + gossip_signatures = {SignatureKey(Uint64(0), data_root): make_signature(0)} # Proposer for slot 1 with 2 validators: slot % num_validators = 1 % 2 = 1 block, post_state, aggregated_atts, aggregated_proofs = state.build_block( @@ -453,10 +451,10 @@ def test_compute_aggregated_signatures_with_multiple_data_groups() -> None: data_root2 = att_data2.data_root_bytes() gossip_signatures = { - SignatureKey(Uint64(0), Bytes32(data_root1)): make_signature(0), - SignatureKey(Uint64(1), Bytes32(data_root1)): make_signature(1), - SignatureKey(Uint64(2), Bytes32(data_root2)): make_signature(2), - SignatureKey(Uint64(3), Bytes32(data_root2)): make_signature(3), + SignatureKey(Uint64(0), data_root1): make_signature(0), + SignatureKey(Uint64(1), data_root1): make_signature(1), + SignatureKey(Uint64(2), data_root2): make_signature(2), + SignatureKey(Uint64(3), data_root2): make_signature(3), } aggregated_atts, aggregated_proofs = state.compute_aggregated_signatures( @@ -478,13 +476,13 @@ def test_compute_aggregated_signatures_falls_back_to_block_payload() -> None: data_root = att_data.data_root_bytes() # Only gossip signature for validator 0 (incomplete) - gossip_signatures = {SignatureKey(Uint64(0), Bytes32(data_root)): make_signature(0)} + gossip_signatures = {SignatureKey(Uint64(0), data_root): make_signature(0)} # Block payload covers both validators block_proof = make_test_proof([Uint64(0), Uint64(1)], b"block-fallback") aggregated_payloads = { - SignatureKey(Uint64(0), Bytes32(data_root)): [block_proof], - SignatureKey(Uint64(1), Bytes32(data_root)): [block_proof], + SignatureKey(Uint64(0), data_root): [block_proof], + SignatureKey(Uint64(1), data_root): [block_proof], } aggregated_atts, aggregated_proofs = state.compute_aggregated_signatures( diff --git a/tests/lean_spec/subspecs/forkchoice/test_store_attestations.py b/tests/lean_spec/subspecs/forkchoice/test_store_attestations.py index 57e03ab6..1352dc4b 100644 --- a/tests/lean_spec/subspecs/forkchoice/test_store_attestations.py +++ b/tests/lean_spec/subspecs/forkchoice/test_store_attestations.py @@ -60,7 +60,7 @@ def test_on_block_processes_multi_validator_aggregations() -> None: # Store signatures in gossip_signatures data_root = attestation_data.data_root_bytes() gossip_sigs = { - SignatureKey(validator_id, Bytes32(data_root)): key_manager.sign_attestation_data( + SignatureKey(validator_id, data_root): key_manager.sign_attestation_data( validator_id, attestation_data ) for validator_id in (Uint64(1), Uint64(2)) @@ -151,7 +151,7 @@ def test_on_block_preserves_immutability_of_aggregated_payloads() -> None: validator_id: attestation_data_1 for validator_id in (Uint64(1), Uint64(2)) } gossip_sigs_1 = { - SignatureKey(validator_id, Bytes32(data_root_1)): key_manager.sign_attestation_data( + SignatureKey(validator_id, data_root_1): key_manager.sign_attestation_data( validator_id, attestation_data_1 ) for validator_id in (Uint64(1), Uint64(2)) @@ -217,7 +217,7 @@ def test_on_block_preserves_immutability_of_aggregated_payloads() -> None: validator_id: attestation_data_2 for validator_id in (Uint64(1), Uint64(2)) } gossip_sigs_2 = { - SignatureKey(validator_id, Bytes32(data_root_2)): key_manager.sign_attestation_data( + SignatureKey(validator_id, data_root_2): key_manager.sign_attestation_data( validator_id, attestation_data_2 ) for validator_id in (Uint64(1), Uint64(2)) diff --git a/tests/lean_spec/subspecs/forkchoice/test_validator.py b/tests/lean_spec/subspecs/forkchoice/test_validator.py index bd6ed349..408dcc40 100644 --- a/tests/lean_spec/subspecs/forkchoice/test_validator.py +++ b/tests/lean_spec/subspecs/forkchoice/test_validator.py @@ -196,8 +196,8 @@ def test_produce_block_with_attestations(self, sample_store: Store) -> None: ) sample_store.latest_known_attestations[Uint64(5)] = signed_5.message sample_store.latest_known_attestations[Uint64(6)] = signed_6.message - sig_key_5 = SignatureKey(Uint64(5), Bytes32(signed_5.message.data_root_bytes())) - sig_key_6 = SignatureKey(Uint64(6), Bytes32(signed_6.message.data_root_bytes())) + sig_key_5 = SignatureKey(Uint64(5), signed_5.message.data_root_bytes()) + sig_key_6 = SignatureKey(Uint64(6), signed_6.message.data_root_bytes()) sample_store.gossip_signatures[sig_key_5] = signed_5.signature sample_store.gossip_signatures[sig_key_6] = signed_6.signature @@ -290,7 +290,7 @@ def test_produce_block_state_consistency(self, sample_store: Store) -> None: target=sample_store.get_attestation_target(), ) sample_store.latest_known_attestations[Uint64(7)] = signed_7.message - sig_key_7 = SignatureKey(Uint64(7), Bytes32(signed_7.message.data_root_bytes())) + sig_key_7 = SignatureKey(Uint64(7), signed_7.message.data_root_bytes()) sample_store.gossip_signatures[sig_key_7] = signed_7.signature store, block, _signatures = sample_store.produce_block_with_signatures( From 0ded890a0fe91829ae3a22bbd7743327ec66ee0b Mon Sep 17 00:00:00 2001 From: Thomas Coratger Date: Tue, 30 Dec 2025 19:33:43 +0100 Subject: [PATCH 5/5] fix linter --- .../containers/test_state_aggregation.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/lean_spec/subspecs/containers/test_state_aggregation.py b/tests/lean_spec/subspecs/containers/test_state_aggregation.py index 988ce875..bae36e83 100644 --- a/tests/lean_spec/subspecs/containers/test_state_aggregation.py +++ b/tests/lean_spec/subspecs/containers/test_state_aggregation.py @@ -117,7 +117,7 @@ def test_gossip_aggregation_succeeds_with_all_signatures() -> None: result = state._aggregate_signatures_from_gossip( validator_ids, - data_root, + Bytes32(data_root), Slot(3), gossip_signatures, ) @@ -135,7 +135,7 @@ def test_gossip_aggregation_returns_partial_result_when_some_missing() -> None: result = state._aggregate_signatures_from_gossip( [Uint64(0), Uint64(1)], - data_root, + Bytes32(data_root), Slot(2), gossip_signatures, ) @@ -154,7 +154,7 @@ def test_gossip_aggregation_returns_none_if_no_signature_matches() -> None: result = state._aggregate_signatures_from_gossip( [Uint64(0), Uint64(1)], - data_root, + Bytes32(data_root), Slot(2), gossip_signatures, ) @@ -177,7 +177,7 @@ def test_pick_from_aggregated_proofs_prefers_widest_overlap() -> None: proof, remaining = state._pick_from_aggregated_proofs( remaining_validator_ids=remaining_validator_ids, - data_root=data_root, + data_root=Bytes32(data_root), aggregated_payloads=aggregated_payloads, ) @@ -200,7 +200,7 @@ def test_pick_from_aggregated_proofs_returns_remaining_for_partial_payload() -> proof, remaining = state._pick_from_aggregated_proofs( remaining_validator_ids=remaining_validator_ids, - data_root=data_root, + data_root=Bytes32(data_root), aggregated_payloads=aggregated_payloads, ) @@ -215,7 +215,7 @@ def test_pick_from_aggregated_proofs_requires_payloads() -> None: with pytest.raises(ValueError, match="aggregated payloads required"): state._pick_from_aggregated_proofs( remaining_validator_ids={Uint64(0)}, - data_root=b"\x55" * 32, + data_root=Bytes32(b"\x55" * 32), aggregated_payloads=None, ) @@ -226,7 +226,7 @@ def test_pick_from_aggregated_proofs_errors_on_empty_remaining() -> None: with pytest.raises(ValueError, match="remaining validator ids cannot be empty"): state._pick_from_aggregated_proofs( remaining_validator_ids=set(), - data_root=b"\x66" * 32, + data_root=Bytes32(b"\x66" * 32), aggregated_payloads={}, ) @@ -238,7 +238,7 @@ def test_pick_from_aggregated_proofs_errors_when_no_candidates() -> None: with pytest.raises(ValueError, match="Failed to locate aggregated proof"): state._pick_from_aggregated_proofs( remaining_validator_ids={Uint64(0)}, - data_root=data_root, + data_root=Bytes32(data_root), aggregated_payloads={SignatureKey(Uint64(0), Bytes32(data_root)): []}, ) @@ -381,7 +381,7 @@ def test_gossip_aggregation_with_empty_validator_list() -> None: result = state._aggregate_signatures_from_gossip( [], # empty validator list - data_root, + Bytes32(data_root), Slot(1), gossip_signatures, ) @@ -396,7 +396,7 @@ def test_gossip_aggregation_with_none_gossip_signatures() -> None: result = state._aggregate_signatures_from_gossip( [Uint64(0), Uint64(1)], - data_root, + Bytes32(data_root), Slot(1), None, # None gossip_signatures ) @@ -411,7 +411,7 @@ def test_gossip_aggregation_with_empty_gossip_signatures() -> None: result = state._aggregate_signatures_from_gossip( [Uint64(0), Uint64(1)], - data_root, + Bytes32(data_root), Slot(1), {}, # empty dict )