xmss: make message Bytes32 (#335)

Author: tcoratger

src/lean_spec/subspecs/validator/service.py

@@ -332,11 +332,10 @@ def _sign_block(
         if entry is None:
             raise ValueError(f"No secret key for validator {validator_index}")
 
-        message_bytes = proposer_attestation_data.data_root_bytes()
         proposer_signature = TARGET_SIGNATURE_SCHEME.sign(
             entry.secret_key,
             block.slot,
-            bytes(message_bytes),
+            proposer_attestation_data.data_root_bytes(),
         )
 
         # Create the message wrapper.
@@ -385,11 +384,10 @@ def _sign_attestation(
         # Sign the attestation data root.
         #
         # Uses XMSS one-time signature for the current epoch (slot).
-        message_bytes = attestation_data.data_root_bytes()
         signature = TARGET_SIGNATURE_SCHEME.sign(
             entry.secret_key,
             attestation_data.slot,
-            bytes(message_bytes),
+            attestation_data.data_root_bytes(),
         )
 
         return SignedAttestation(

src/lean_spec/subspecs/xmss/aggregation.py

@@ -65,7 +65,7 @@ def aggregate(
         participants: AggregationBits,
         public_keys: Sequence[PublicKey],
         signatures: Sequence[Signature],
-        message: bytes,
+        message: Bytes32,
         epoch: Uint64,
         mode: str | None = None,
     ) -> Self:
@@ -106,7 +106,7 @@ def aggregate(
     def verify(
         self,
         public_keys: Sequence[PublicKey],
-        message: bytes,
+        message: Bytes32,
         epoch: Uint64,
         mode: str | None = None,
     ) -> None:

src/lean_spec/subspecs/xmss/containers.py

@@ -9,7 +9,7 @@
 
 from typing import TYPE_CHECKING, Mapping, NamedTuple
 
-from ...types import Uint64
+from ...types import Bytes32, Uint64
 from ...types.container import Container
 from .subtree import HashSubTree
 from .types import (
@@ -71,7 +71,7 @@ def verify(
         self,
         public_key: PublicKey,
         epoch: "Uint64",
-        message: bytes,
+        message: "Bytes32",
         scheme: "GeneralizedXmssScheme",
     ) -> bool:
         """

src/lean_spec/subspecs/xmss/interface.py

@@ -16,7 +16,7 @@
     TEST_TARGET_SUM_ENCODER,
     TargetSumEncoder,
 )
-from lean_spec.types import StrictBaseModel, Uint64
+from lean_spec.types import Bytes32, StrictBaseModel, Uint64
 
 from ._validation import enforce_strict_types
 from .constants import (
@@ -220,7 +220,7 @@ def key_gen(self, activation_epoch: Uint64, num_active_epochs: Uint64) -> KeyPai
         )
         return KeyPair(public=pk, secret=sk)
 
-    def sign(self, sk: SecretKey, epoch: Uint64, message: bytes) -> Signature:
+    def sign(self, sk: SecretKey, epoch: Uint64, message: Bytes32) -> Signature:
         """
         Produces a digital signature for a given message at a specific epoch.
 
@@ -362,7 +362,7 @@ def sign(self, sk: SecretKey, epoch: Uint64, message: bytes) -> Signature:
         # - The randomness `rho` needed for verification.
         return Signature(path=path, rho=rho, hashes=HashDigestList(data=ots_hashes))
 
-    def verify(self, pk: PublicKey, epoch: Uint64, message: bytes, sig: Signature) -> bool:
+    def verify(self, pk: PublicKey, epoch: Uint64, message: Bytes32, sig: Signature) -> bool:
         r"""
         Verifies a digital signature against a public key, message, and epoch.
 

src/lean_spec/subspecs/xmss/message_hash.py

@@ -36,7 +36,7 @@
     TEST_POSEIDON,
     PoseidonXmss,
 )
-from lean_spec.types import StrictBaseModel, Uint64
+from lean_spec.types import Bytes32, StrictBaseModel, Uint64
 
 from ..koalabear import Fp, P
 from ._validation import enforce_strict_types
@@ -70,7 +70,7 @@ def _validate_strict_types(self) -> "MessageHasher":
         enforce_strict_types(self, config=XmssConfig, poseidon=PoseidonXmss)
         return self
 
-    def encode_message(self, message: bytes) -> list[Fp]:
+    def encode_message(self, message: Bytes32) -> list[Fp]:
         """
         Encodes a 32-byte message into a list of field elements.
 
@@ -145,7 +145,7 @@ def apply(
         parameter: Parameter,
         epoch: Uint64,
         rho: Randomness,
-        message: bytes,
+        message: Bytes32,
     ) -> list[int]:
         """
         Applies the full "Top Level" message hash and mapping procedure.

src/lean_spec/subspecs/xmss/prf.py

@@ -15,7 +15,7 @@
 from pydantic import model_validator
 
 from lean_spec.subspecs.koalabear import Fp
-from lean_spec.types import StrictBaseModel, Uint64
+from lean_spec.types import Bytes32, StrictBaseModel, Uint64
 
 from ._validation import enforce_strict_types
 from .constants import (
@@ -178,7 +178,7 @@ def apply(self, key: PRFKey, epoch: Uint64, chain_index: Uint64) -> HashDigestVe
         return HashDigestVector(data=_bytes_to_field_elements(prf_output_bytes, config.HASH_LEN_FE))
 
     def get_randomness(
-        self, key: PRFKey, epoch: Uint64, message: bytes, counter: Uint64
+        self, key: PRFKey, epoch: Uint64, message: Bytes32, counter: Uint64
     ) -> Randomness:
         """
         Derives pseudorandom field elements for use in deterministic signing.

src/lean_spec/subspecs/xmss/target_sum.py

@@ -8,7 +8,7 @@
 
 from pydantic import model_validator
 
-from lean_spec.types import StrictBaseModel, Uint64
+from lean_spec.types import Bytes32, StrictBaseModel, Uint64
 
 from ._validation import enforce_strict_types
 from .constants import PROD_CONFIG, TEST_CONFIG, XmssConfig
@@ -41,7 +41,7 @@ def _validate_strict_types(self) -> "TargetSumEncoder":
         return self
 
     def encode(
-        self, parameter: Parameter, message: bytes, rho: Randomness, epoch: Uint64
+        self, parameter: Parameter, message: Bytes32, rho: Randomness, epoch: Uint64
     ) -> list[int] | None:
         """
         Encodes a message into a codeword if it meets the target sum criteria.

tests/lean_spec/subspecs/validator/test_service.py

@@ -601,7 +601,7 @@ async def produce_block() -> None:
         is_valid = TARGET_SIGNATURE_SCHEME.verify(
             pk=proposer_public_key,
             epoch=signed_block.message.block.slot,
-            message=bytes(message_bytes),
+            message=message_bytes,
             sig=signed_block.signature.proposer_signature,
         )
         assert is_valid, "Proposer signature failed verification"
@@ -648,7 +648,7 @@ async def produce_attestations() -> None:
             is_valid = TARGET_SIGNATURE_SCHEME.verify(
                 pk=public_key,
                 epoch=signed_att.message.slot,
-                message=bytes(message_bytes),
+                message=message_bytes,
                 sig=signed_att.signature,
             )
             assert is_valid, f"Attestation signature for validator {validator_id} failed"
@@ -750,7 +750,7 @@ async def produce_block() -> None:
         is_valid = TARGET_SIGNATURE_SCHEME.verify(
             pk=public_key,
             epoch=signed_block.message.block.slot,
-            message=bytes(message_bytes),
+            message=message_bytes,
             sig=signed_block.signature.proposer_signature,
         )
         assert is_valid
@@ -1005,7 +1005,7 @@ async def produce_attestations() -> None:
             is_valid = TARGET_SIGNATURE_SCHEME.verify(
                 pk=public_key,
                 epoch=test_slot,  # Must match the signing slot
-                message=bytes(message_bytes),
+                message=message_bytes,
                 sig=signed_att.signature,
             )
             assert is_valid, f"Signature for validator {validator_id} at slot {test_slot} failed"
@@ -1015,7 +1015,7 @@ async def produce_attestations() -> None:
             is_invalid = TARGET_SIGNATURE_SCHEME.verify(
                 pk=public_key,
                 epoch=wrong_epoch,
-                message=bytes(message_bytes),
+                message=message_bytes,
                 sig=signed_att.signature,
             )
             assert not is_invalid, "Signature should fail with wrong epoch"

tests/lean_spec/subspecs/xmss/test_interface.py

@@ -8,7 +8,7 @@
     TEST_SIGNATURE_SCHEME,
     GeneralizedXmssScheme,
 )
-from lean_spec.types import Uint64
+from lean_spec.types import Bytes32, Uint64
 
 
 def _test_correctness_roundtrip(
@@ -32,7 +32,7 @@ def _test_correctness_roundtrip(
     #
     # Pick a sample epoch within the active range to test signing.
     test_epoch = Uint64(activation_epoch + num_active_epochs // 2)
-    message = b"\x42" * scheme.config.MESSAGE_LENGTH
+    message = Bytes32(b"\x42" * 32)
 
     # Sign the message at the chosen epoch.
     #
@@ -46,7 +46,7 @@ def _test_correctness_roundtrip(
     # TEST INVALID CASES
     #
     # Verification must fail if the message is tampered with.
-    tampered_message = b"\x43" * scheme.config.MESSAGE_LENGTH
+    tampered_message = Bytes32(b"\x43" * 32)
 
     # With small test parameters (test configuration), there's a small chance that
     # the tampered message produces the same codeword as the original due to
@@ -176,7 +176,7 @@ def test_sign_requires_prepared_interval() -> None:
     assert int(outside_epoch) not in prepared_interval
 
     # Signing should fail
-    message = b"\x42" * scheme.config.MESSAGE_LENGTH
+    message = Bytes32(b"\x42" * 32)
     with pytest.raises(ValueError, match="outside the prepared interval"):
         scheme.sign(sk, outside_epoch, message)
 
@@ -189,7 +189,7 @@ def test_deterministic_signing() -> None:
 
     # Use epoch within prepared interval
     epoch = Uint64(4)
-    message = b"\x42" * scheme.config.MESSAGE_LENGTH
+    message = Bytes32(b"\x42" * 32)
 
     # Sign twice
     sig1 = scheme.sign(sk, epoch, message)
@@ -218,7 +218,7 @@ def test_rejects_epoch_beyond_lifetime(self) -> None:
 
         # Sign a valid message at a valid epoch.
         valid_epoch = Uint64(4)
-        message = b"\x42" * scheme.config.MESSAGE_LENGTH
+        message = Bytes32(b"\x42" * 32)
         signature = scheme.sign(sk, valid_epoch, message)
 
         # Verify with an epoch beyond LIFETIME.
@@ -234,7 +234,7 @@ def test_rejects_very_large_epoch(self) -> None:
         pk, sk = scheme.key_gen(Uint64(0), Uint64(scheme.config.LIFETIME))
 
         valid_epoch = Uint64(4)
-        message = b"\x42" * scheme.config.MESSAGE_LENGTH
+        message = Bytes32(b"\x42" * 32)
         signature = scheme.sign(sk, valid_epoch, message)
 
         # Try to verify with a huge epoch.

tests/lean_spec/subspecs/xmss/test_message_hash.py

@@ -12,7 +12,7 @@
 )
 from lean_spec.subspecs.xmss.rand import TEST_RAND
 from lean_spec.subspecs.xmss.utils import int_to_base_p
-from lean_spec.types import Uint64
+from lean_spec.types import Bytes32, Uint64
 
 
 def test_encode_message() -> None:
@@ -21,13 +21,13 @@ def test_encode_message() -> None:
     hasher = TEST_MESSAGE_HASHER
 
     # All-zero message
-    msg_zeros = b"\x00" * config.MESSAGE_LENGTH
+    msg_zeros = Bytes32(b"\x00" * 32)
     encoded_zeros = hasher.encode_message(msg_zeros)
     assert len(encoded_zeros) == config.MSG_LEN_FE
     assert all(fe.value == 0 for fe in encoded_zeros)
 
     # All-max message (0xff)
-    msg_max = b"\xff" * config.MESSAGE_LENGTH
+    msg_max = Bytes32(b"\xff" * 32)
     acc = int.from_bytes(msg_max, "little")
     expected_max = int_to_base_p(acc, config.MSG_LEN_FE)
     assert hasher.encode_message(msg_max) == expected_max
@@ -70,7 +70,7 @@ def test_apply_output_is_in_correct_hypercube_part() -> None:
     parameter = rand.parameter()
     epoch = Uint64(313)
     randomness = rand.rho()
-    message = b"\xaa" * config.MESSAGE_LENGTH
+    message = Bytes32(b"\xaa" * 32)
 
     # Call the message hash function.
     vertex = hasher.apply(parameter, epoch, randomness, message)