|
| 1 | +"""Reaggregation test fixture format.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +from typing import ClassVar |
| 6 | + |
| 7 | +from consensus_testing.keys import XmssKeyManager |
| 8 | +from consensus_testing.test_fixtures.base import BaseConsensusFixture, BaseTestSpec |
| 9 | +from lean_spec.spec.crypto.merkleization import hash_tree_root |
| 10 | +from lean_spec.spec.forks import AggregationBits, Checkpoint, Slot, ValidatorIndex |
| 11 | +from lean_spec.spec.forks.lstar.containers import ( |
| 12 | + AggregatedAttestation, |
| 13 | + AggregatedAttestations, |
| 14 | + AttestationData, |
| 15 | + Block, |
| 16 | + BlockBody, |
| 17 | + MultiMessageAggregate, |
| 18 | + SingleMessageAggregate, |
| 19 | +) |
| 20 | +from lean_spec.spec.ssz import Bytes32 |
| 21 | + |
| 22 | +ATTESTATION_SLOT: Slot = Slot(1) |
| 23 | +"""Attestation slot, one before the block that carries it.""" |
| 24 | + |
| 25 | +BLOCK_SLOT: Slot = Slot(2) |
| 26 | +"""Block slot, one after the attestation it carries.""" |
| 27 | + |
| 28 | +PROPOSER_INDEX: ValidatorIndex = ValidatorIndex(0) |
| 29 | +"""Validator index that signs the block.""" |
| 30 | + |
| 31 | +CHAIN_ROOT: Bytes32 = Bytes32(b"\x11" * 32) |
| 32 | +"""Head and target root the attestation votes for.""" |
| 33 | + |
| 34 | +GENESIS_ROOT: Bytes32 = Bytes32(b"\x33" * 32) |
| 35 | +"""Source root anchoring the attestation.""" |
| 36 | + |
| 37 | +PARENT_ROOT: Bytes32 = Bytes32(b"\xaa" * 32) |
| 38 | +"""Parent root of the synthetic block.""" |
| 39 | + |
| 40 | + |
| 41 | +class ReaggregationFixture(BaseConsensusFixture): |
| 42 | + """Emitted vector for proof re-aggregation.""" |
| 43 | + |
| 44 | + block_proof: str |
| 45 | + """The block's multi-message proof that gets split, as hex.""" |
| 46 | + |
| 47 | + public_keys_per_message: list[list[str]] |
| 48 | + """Per-message public key layout the original proof was built with, as hex.""" |
| 49 | + |
| 50 | + attestation_message: str |
| 51 | + """Hash tree root of the split attestation, as hex.""" |
| 52 | + |
| 53 | + attestation_slot: int |
| 54 | + """Slot of the split attestation.""" |
| 55 | + |
| 56 | + block_attesters: list[int] |
| 57 | + """Validator indices whose signed attestations came in the block.""" |
| 58 | + |
| 59 | + local_attesters: list[int] |
| 60 | + """Validator indices whose signed attestations are in the node's local pool.""" |
| 61 | + |
| 62 | + combined_attesters: list[int] |
| 63 | + """Validator indices covered by the re-aggregated proof, the block and local union.""" |
| 64 | + |
| 65 | + reaggregated_proof: str |
| 66 | + """The re-aggregated single-message proof bytes, as hex.""" |
| 67 | + |
| 68 | + |
| 69 | +class ReaggregationTest(BaseTestSpec): |
| 70 | + """ |
| 71 | + Split one attestation's proof out of a block, then merge it with the local partial. |
| 72 | +
|
| 73 | + The reference proof bytes are not deterministic. |
| 74 | + Each vector is checked by verifying against the expected attesters' keys, not by byte match. |
| 75 | + """ |
| 76 | + |
| 77 | + format_name: ClassVar[str] = "reaggregation_test" |
| 78 | + description: ClassVar[str] = "Tests aggregate proof split and merge clients must reproduce" |
| 79 | + |
| 80 | + block_attesters: list[ValidatorIndex] |
| 81 | + """Validator indices whose signed attestations were aggregated into the block.""" |
| 82 | + |
| 83 | + local_attesters: list[ValidatorIndex] = [] |
| 84 | + """Validator indices whose signed attestations are in the node's local pool.""" |
| 85 | + |
| 86 | + def generate(self) -> ReaggregationFixture: |
| 87 | + """Build the merged proof, split the attestation out, merge it, and verify.""" |
| 88 | + key_manager = XmssKeyManager.shared() |
| 89 | + |
| 90 | + attestation_data = AttestationData( |
| 91 | + slot=ATTESTATION_SLOT, |
| 92 | + head=Checkpoint(root=CHAIN_ROOT, slot=ATTESTATION_SLOT), |
| 93 | + target=Checkpoint(root=CHAIN_ROOT, slot=ATTESTATION_SLOT), |
| 94 | + source=Checkpoint(root=GENESIS_ROOT, slot=Slot(0)), |
| 95 | + ) |
| 96 | + attestation_message = hash_tree_root(attestation_data) |
| 97 | + |
| 98 | + signing_validators = list(dict.fromkeys([*self.block_attesters, *self.local_attesters])) |
| 99 | + shared_signatures = { |
| 100 | + validator_index: key_manager.sign_attestation_data(validator_index, attestation_data) |
| 101 | + for validator_index in signing_validators |
| 102 | + } |
| 103 | + |
| 104 | + # Phase 1: the block's attestation component, then the proposal component. |
| 105 | + attestation_component = key_manager.sign_and_aggregate( |
| 106 | + self.block_attesters, |
| 107 | + attestation_data, |
| 108 | + precomputed_signatures=shared_signatures, |
| 109 | + ) |
| 110 | + attestation_keys = [ |
| 111 | + key_manager.get_public_keys(validator_index)[0] |
| 112 | + for validator_index in self.block_attesters |
| 113 | + ] |
| 114 | + block = Block( |
| 115 | + slot=BLOCK_SLOT, |
| 116 | + proposer_index=PROPOSER_INDEX, |
| 117 | + parent_root=PARENT_ROOT, |
| 118 | + state_root=Bytes32.zero(), |
| 119 | + body=BlockBody( |
| 120 | + attestations=AggregatedAttestations( |
| 121 | + data=[ |
| 122 | + AggregatedAttestation( |
| 123 | + aggregation_bits=AggregationBits.from_indices(self.block_attesters), |
| 124 | + data=attestation_data, |
| 125 | + ) |
| 126 | + ] |
| 127 | + ) |
| 128 | + ), |
| 129 | + ) |
| 130 | + block_root = hash_tree_root(block) |
| 131 | + proposal_key = key_manager.get_public_keys(PROPOSER_INDEX)[1] |
| 132 | + proposal_signature = key_manager.sign_block_root(PROPOSER_INDEX, BLOCK_SLOT, block_root) |
| 133 | + proposal_component = SingleMessageAggregate.aggregate( |
| 134 | + children=[], |
| 135 | + raw_xmss=[(PROPOSER_INDEX, proposal_key, proposal_signature)], |
| 136 | + message=block_root, |
| 137 | + slot=BLOCK_SLOT, |
| 138 | + ) |
| 139 | + public_keys_per_message = [attestation_keys, [proposal_key]] |
| 140 | + |
| 141 | + # Phase 2: merge the two components into a single, multi-message block proof. |
| 142 | + block_proof = MultiMessageAggregate.aggregate( |
| 143 | + [attestation_component, proposal_component], |
| 144 | + public_keys_per_aggregate=public_keys_per_message, |
| 145 | + ) |
| 146 | + |
| 147 | + # Phase 3: split the attestation's component back out by its message. |
| 148 | + block_bits = AggregationBits.from_indices(self.block_attesters) |
| 149 | + recovered_proof = block_proof.split_by_message( |
| 150 | + message=attestation_message, |
| 151 | + public_keys_per_message=public_keys_per_message, |
| 152 | + participants=block_bits, |
| 153 | + ) |
| 154 | + |
| 155 | + # The split output must verify on its own against the block attesters' keys. |
| 156 | + # A later merge could otherwise mask a malformed recovered component. |
| 157 | + recovered_proof.verify(attestation_keys, attestation_message, attestation_data.slot) |
| 158 | + |
| 159 | + # Phase 4: merge the recovered proof with the local partial. |
| 160 | + if self.local_attesters: |
| 161 | + local_partial = key_manager.sign_and_aggregate( |
| 162 | + self.local_attesters, |
| 163 | + attestation_data, |
| 164 | + precomputed_signatures=shared_signatures, |
| 165 | + ) |
| 166 | + local_keys = [ |
| 167 | + key_manager.get_public_keys(validator_index)[0] |
| 168 | + for validator_index in self.local_attesters |
| 169 | + ] |
| 170 | + reaggregated_proof = SingleMessageAggregate.aggregate( |
| 171 | + children=[ |
| 172 | + (recovered_proof, attestation_keys), |
| 173 | + (local_partial, local_keys), |
| 174 | + ], |
| 175 | + raw_xmss=[], |
| 176 | + message=attestation_message, |
| 177 | + slot=attestation_data.slot, |
| 178 | + ) |
| 179 | + else: |
| 180 | + reaggregated_proof = recovered_proof |
| 181 | + |
| 182 | + # Phase 5: the merged proof must cover exactly the union of block and local attesters. |
| 183 | + combined_attester_indices = list(reaggregated_proof.participants.to_validator_indices()) |
| 184 | + expected_attester_indices = sorted({*self.block_attesters, *self.local_attesters}) |
| 185 | + assert combined_attester_indices == expected_attester_indices, ( |
| 186 | + f"re-aggregated proof covers {combined_attester_indices}, " |
| 187 | + f"expected the union {expected_attester_indices}" |
| 188 | + ) |
| 189 | + |
| 190 | + # The merged proof must verify against the union attesters' keys. |
| 191 | + reaggregated_proof.verify( |
| 192 | + [ |
| 193 | + key_manager.get_public_keys(validator_index)[0] |
| 194 | + for validator_index in combined_attester_indices |
| 195 | + ], |
| 196 | + attestation_message, |
| 197 | + attestation_data.slot, |
| 198 | + ) |
| 199 | + |
| 200 | + return ReaggregationFixture( |
| 201 | + block_proof="0x" + bytes(block_proof.proof.data).hex(), |
| 202 | + public_keys_per_message=[ |
| 203 | + ["0x" + public_key.encode_bytes().hex() for public_key in component_keys] |
| 204 | + for component_keys in public_keys_per_message |
| 205 | + ], |
| 206 | + attestation_message="0x" + bytes(attestation_message).hex(), |
| 207 | + attestation_slot=int(attestation_data.slot), |
| 208 | + block_attesters=[int(validator_index) for validator_index in self.block_attesters], |
| 209 | + local_attesters=[int(validator_index) for validator_index in self.local_attesters], |
| 210 | + combined_attesters=[ |
| 211 | + int(validator_index) for validator_index in combined_attester_indices |
| 212 | + ], |
| 213 | + reaggregated_proof="0x" + bytes(reaggregated_proof.proof.data).hex(), |
| 214 | + ) |
0 commit comments