xmss: full ssz serialization (#210)

* xmss: full ssz serialization

* some fixes

* some fixes

* add todos

* cleanup

* fix test

Author: tcoratger

packages/testing/src/consensus_testing/keys.py

@@ -2,10 +2,10 @@
 
 from typing import NamedTuple, Optional
 
-from lean_spec.subspecs.containers import Attestation, Signature
+from lean_spec.subspecs.containers import Attestation
 from lean_spec.subspecs.containers.slot import Slot
 from lean_spec.subspecs.ssz.hash import hash_tree_root
-from lean_spec.subspecs.xmss.containers import PublicKey, SecretKey
+from lean_spec.subspecs.xmss.containers import PublicKey, SecretKey, Signature
 from lean_spec.subspecs.xmss.interface import (
     TEST_SIGNATURE_SCHEME,
     GeneralizedXmssScheme,
@@ -174,19 +174,8 @@ def sign_attestation(self, attestation: Attestation) -> Signature:
         # This produces a cryptographic hash of the entire attestation structure.
         message = bytes(hash_tree_root(attestation))
 
-        # Generate the XMSS signature using the validator's (now prepared) secret key.
-        xmss_sig = self.scheme.sign(sk, epoch, message)
-
-        # Convert the signature to the wire format (byte array).
-        signature_bytes = xmss_sig.to_bytes(self.scheme.config)
-
-        # Ensure the signature meets the consensus spec length (3116 bytes).
-        #
-        # This is necessary when using TEST_CONFIG (796 bytes) vs PROD_CONFIG.
-        # Padding with zeros on the right maintains compatibility.
-        padded_bytes = signature_bytes.ljust(Signature.LENGTH, b"\x00")
-
-        return Signature(padded_bytes)
+        # Generate and return the XMSS signature using the validator's (now prepared) secret key.
+        return self.scheme.sign(sk, epoch, message)
 
     def get_public_key(self, validator_index: ValidatorIndex) -> PublicKey:
         """

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

@@ -8,7 +8,6 @@
 from pydantic import model_validator
 
 from lean_spec.subspecs.chain.config import SECONDS_PER_SLOT
-from lean_spec.subspecs.containers import Signature
 from lean_spec.subspecs.containers.attestation import (
     Attestation,
     AttestationData,
@@ -26,7 +25,15 @@
 from lean_spec.subspecs.containers.state import Validators
 from lean_spec.subspecs.containers.state.state import State
 from lean_spec.subspecs.forkchoice import Store
+from lean_spec.subspecs.koalabear import Fp
 from lean_spec.subspecs.ssz import hash_tree_root
+from lean_spec.subspecs.xmss.constants import PROD_CONFIG
+from lean_spec.subspecs.xmss.containers import (
+    HashDigestList,
+    HashTreeOpening,
+    Randomness,
+    Signature,
+)
 from lean_spec.subspecs.xmss.interface import TEST_SIGNATURE_SCHEME
 from lean_spec.types import Bytes32, Uint64, ValidatorIndex
 
@@ -192,11 +199,7 @@ def make_fixture(self) -> ForkChoiceTest:
         # Update validator pubkeys to match key_manager's generated keys
         updated_validators = [
             validator.model_copy(
-                update={
-                    "pubkey": key_manager[ValidatorIndex(i)].public.to_bytes(
-                        key_manager.scheme.config
-                    )
-                }
+                update={"pubkey": key_manager[ValidatorIndex(i)].public.encode_bytes()}
             )
             for i, validator in enumerate(self.anchor_state.validators)
         ]
@@ -481,5 +484,13 @@ def _build_signed_attestation_from_spec(
         # Create signed attestation
         return SignedAttestation(
             message=attestation,
-            signature=spec.signature if spec.signature is not None else Signature.zero(),
+            signature=(
+                spec.signature
+                if spec.signature is not None
+                else Signature(
+                    path=HashTreeOpening(siblings=HashDigestList(data=[])),
+                    rho=Randomness(data=[Fp(0) for _ in range(PROD_CONFIG.RAND_LEN_FE)]),
+                    hashes=HashDigestList(data=[]),
+                )
+            ),
         )

packages/testing/src/consensus_testing/test_types/signed_attestation_spec.py

@@ -1,7 +1,7 @@
 """Lightweight attestation specification for test definitions."""
 
-from lean_spec.subspecs.containers.signature import Signature
 from lean_spec.subspecs.containers.slot import Slot
+from lean_spec.subspecs.xmss.containers import Signature
 from lean_spec.types import CamelModel, Uint64, ValidatorIndex
 
 
@@ -33,5 +33,5 @@ class SignedAttestationSpec(CamelModel):
     """
     The signature for the attestation.
 
-    If None, uses Signature.zero() for testing.
+    If None, uses an empty signature for testing.
     """

src/lean_spec/subspecs/containers/__init__.py

@@ -26,7 +26,6 @@
 )
 from .checkpoint import Checkpoint
 from .config import Config
-from .signature import Signature
 from .state import State
 from .validator import Validator
 
@@ -45,7 +44,6 @@
     "Checkpoint",
     "Config",
     "SignedBlockWithAttestation",
-    "Signature",
     "Validator",
     "State",
 ]

src/lean_spec/subspecs/containers/attestation/attestation.py

@@ -15,8 +15,8 @@
 from lean_spec.subspecs.containers.slot import Slot
 from lean_spec.types import Container, Uint64
 
+from ...xmss.containers import Signature
 from ..checkpoint import Checkpoint
-from ..signature import Signature
 from .types import AggregatedSignatures, AggregationBits
 
 

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

@@ -4,7 +4,7 @@
 from lean_spec.types.bitfields import BaseBitlist
 
 from ...chain.config import VALIDATOR_REGISTRY_LIMIT
-from ..signature import Signature
+from ...xmss.containers import Signature
 
 
 class AggregationBits(BaseBitlist):

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

@@ -3,8 +3,8 @@
 from lean_spec.types import SSZList
 
 from ...chain.config import VALIDATOR_REGISTRY_LIMIT
+from ...xmss.containers import Signature
 from ..attestation import Attestation
-from ..signature import Signature
 
 
 class Attestations(SSZList):

src/lean_spec/subspecs/containers/signature.py

@@ -20,7 +20,7 @@ class Validator(Container):
 
     def get_pubkey(self, scheme: GeneralizedXmssScheme = TEST_SIGNATURE_SCHEME) -> PublicKey:
         """Get the XMSS public key from this validator."""
-        return PublicKey.from_bytes(bytes(self.pubkey), scheme.config)
+        return PublicKey.decode_bytes(bytes(self.pubkey))
 
     def produce_attestation(self, data: AttestationData) -> Attestation:
         """

src/lean_spec/subspecs/containers/validator.py

@@ -28,14 +28,14 @@
     BlockBody,
     Checkpoint,
     Config,
-    Signature,
     SignedAttestation,
     SignedBlockWithAttestation,
     State,
 )
 from lean_spec.subspecs.containers.block import Attestations
 from lean_spec.subspecs.containers.slot import Slot
 from lean_spec.subspecs.ssz.hash import hash_tree_root
+from lean_spec.subspecs.xmss.containers import Signature
 from lean_spec.types import (
     ZERO_HASH,
     Bytes32,