xmss: make Parameter ssz compatible (leanEthereum#208)

Author: tcoratger

src/lean_spec/subspecs/xmss/containers.py

@@ -77,14 +77,18 @@ class HashDigestList(SSZList):
     LIMIT = NODE_LIST_LIMIT
 
 
-Parameter = List[Fp]
-"""
-A type alias for the public parameter `P`.
+class Parameter(SSZVector):
+    """
+    The public parameter P.
+
+    This is a unique, randomly generated value associated with a single key pair. It
+    is mixed into every hash computation to "personalize" the hash function, preventing
+    certain cross-key attacks. It is public knowledge.
+    """
+
+    ELEMENT_TYPE = Fp
+    LENGTH = PROD_CONFIG.PARAMETER_LEN
 
-This is a unique, randomly generated value associated with a single key pair. It
-is mixed into every hash computation to "personalize" the hash function, preventing
-certain cross-key attacks. It is public knowledge.
-"""
 
 Randomness = List[Fp]
 """
@@ -199,7 +203,7 @@ def __bytes__(self) -> bytes:
             True
         """
         return Fp.serialize_list(cast(List[Fp], list(self.root.data))) + Fp.serialize_list(
-            self.parameter
+            cast(List[Fp], list(self.parameter.data))
         )
 
     def to_bytes(self, config: XmssConfig) -> bytes:
@@ -266,7 +270,7 @@ def from_bytes(cls, data: bytes, config: XmssConfig) -> PublicKey:
         root = Fp.deserialize_list(data[:root_len], config.HASH_LEN_FE)
         parameter = Fp.deserialize_list(data[root_len:], config.PARAMETER_LEN)
 
-        return cls(root=HashDigestVector(data=root), parameter=parameter)
+        return cls(root=HashDigestVector(data=root), parameter=Parameter(data=parameter))
 
 
 class Signature(StrictBaseModel):

src/lean_spec/subspecs/xmss/message_hash.py

@@ -29,7 +29,7 @@
 
 from __future__ import annotations
 
-from typing import List
+from typing import List, cast
 
 from pydantic import model_validator
 
@@ -190,7 +190,13 @@ def apply(
             iteration_separator = [Fp(value=i)]
 
             # The input is: rho || P || epoch || message || iteration.
-            combined_input = rho + parameter + epoch_fe + message_fe + iteration_separator
+            combined_input = (
+                rho
+                + cast(List[Fp], list(parameter.data))
+                + epoch_fe
+                + message_fe
+                + iteration_separator
+            )
 
             # Hash the combined input using Poseidon2 compression mode.
             iteration_output = self.poseidon.compress(

src/lean_spec/subspecs/xmss/rand.py

@@ -32,7 +32,7 @@ def field_elements(self, length: int) -> List[Fp]:
 
     def parameter(self) -> Parameter:
         """Generates a random public parameter."""
-        return self.field_elements(self.config.PARAMETER_LEN)
+        return Parameter(data=self.field_elements(self.config.PARAMETER_LEN))
 
     def domain(self) -> List[Fp]:
         """Generates a random hash digest."""

src/lean_spec/subspecs/xmss/tweak_hash.py

@@ -27,7 +27,7 @@
 from __future__ import annotations
 
 from itertools import chain
-from typing import List, Union
+from typing import List, Union, cast
 
 from pydantic import Field, model_validator
 
@@ -189,14 +189,19 @@ def apply(
             # Case 1: Hashing a single digest (used in hash chains).
             #
             # We use the efficient width-16 compression mode.
-            input_vec = parameter + encoded_tweak + message_parts[0]
+            input_vec = cast(List[Fp], list(parameter.data)) + encoded_tweak + message_parts[0]
             return self.poseidon.compress(input_vec, 16, config.HASH_LEN_FE)
 
         elif len(message_parts) == 2:
             # Case 2: Hashing two digests (used for Merkle tree nodes).
             #
             # We use the slightly larger width-24 compression mode.
-            input_vec = parameter + encoded_tweak + message_parts[0] + message_parts[1]
+            input_vec = (
+                cast(List[Fp], list(parameter.data))
+                + encoded_tweak
+                + message_parts[0]
+                + message_parts[1]
+            )
             return self.poseidon.compress(input_vec, 24, config.HASH_LEN_FE)
 
         else:
@@ -205,7 +210,7 @@ def apply(
             # We use the robust sponge mode.
             # First, flatten the list of message parts into a single vector.
             flattened_message = list(chain.from_iterable(message_parts))
-            input_vec = parameter + encoded_tweak + flattened_message
+            input_vec = cast(List[Fp], list(parameter.data)) + encoded_tweak + flattened_message
 
             # Create a domain separator for the sponge mode based on the input dimensions.
             #

src/lean_spec/subspecs/xmss/utils.py

@@ -5,7 +5,7 @@
 from ...types.uint import Uint64
 from ..koalabear import Fp, P
 from .constants import XmssConfig
-from .containers import HashDigestList, HashDigestVector, HashTreeLayer, PRFKey
+from .containers import HashDigestList, HashDigestVector, HashTreeLayer, Parameter, PRFKey
 from .rand import Rand
 
 if TYPE_CHECKING:
@@ -168,7 +168,7 @@ def bottom_tree_from_prf_key(
     config: XmssConfig,
     prf_key: PRFKey,
     bottom_tree_index: Uint64,
-    parameter: List[Fp],
+    parameter: Parameter,
 ) -> "HashSubTree":
     """
     Generates a single bottom tree on-demand from the PRF key.

tests/lean_spec/subspecs/xmss/test_containers.py

@@ -10,6 +10,7 @@
     HashDigestList,
     HashDigestVector,
     HashTreeOpening,
+    Parameter,
     PublicKey,
     Signature,
 )
@@ -21,7 +22,7 @@ class TestPublicKey:
     def test_bytes_protocol(self) -> None:
         """Test that PublicKey implements Python's bytes protocol."""
         root = HashDigestVector(data=[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)])
-        parameter = [Fp(value=i + 100) for i in range(TEST_CONFIG.PARAMETER_LEN)]
+        parameter = Parameter(data=[Fp(value=i + 100) for i in range(TEST_CONFIG.PARAMETER_LEN)])
         pk = PublicKey(root=root, parameter=parameter)
 
         # Test __bytes__()
@@ -32,7 +33,7 @@ def test_bytes_protocol(self) -> None:
     def test_to_bytes_with_validation(self) -> None:
         """Test that to_bytes validates field lengths."""
         root = HashDigestVector(data=[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)])
-        parameter = [Fp(value=i) for i in range(TEST_CONFIG.PARAMETER_LEN)]
+        parameter = Parameter(data=[Fp(value=i) for i in range(TEST_CONFIG.PARAMETER_LEN)])
         pk = PublicKey(root=root, parameter=parameter)
 
         # Valid serialization
@@ -43,15 +44,14 @@ def test_to_bytes_with_validation(self) -> None:
         with pytest.raises(ValueError, match="requires exactly"):
             HashDigestVector(data=[Fp(value=0)] * 5)
 
-        # Invalid parameter length
-        invalid_pk = PublicKey(root=root, parameter=[Fp(value=0)] * 3)
-        with pytest.raises(ValueError, match="Invalid parameter length"):
-            invalid_pk.to_bytes(TEST_CONFIG)
+        # Invalid parameter length - Parameter validates length at construction
+        with pytest.raises(ValueError, match="requires exactly"):
+            Parameter(data=[Fp(value=0)] * 3)
 
     def test_roundtrip_test_config(self) -> None:
         """Test serialization round-trip with TEST_CONFIG."""
         root = HashDigestVector(data=[Fp(value=i * 10) for i in range(TEST_CONFIG.HASH_LEN_FE)])
-        parameter = [Fp(value=i * 20) for i in range(TEST_CONFIG.PARAMETER_LEN)]
+        parameter = Parameter(data=[Fp(value=i * 20) for i in range(TEST_CONFIG.PARAMETER_LEN)])
         original = PublicKey(root=root, parameter=parameter)
 
         # Serialize and deserialize
@@ -65,7 +65,7 @@ def test_roundtrip_test_config(self) -> None:
     def test_roundtrip_prod_config(self) -> None:
         """Test serialization round-trip with PROD_CONFIG."""
         root = HashDigestVector(data=[Fp(value=i) for i in range(PROD_CONFIG.HASH_LEN_FE)])
-        parameter = [Fp(value=i + 1000) for i in range(PROD_CONFIG.PARAMETER_LEN)]
+        parameter = Parameter(data=[Fp(value=i + 1000) for i in range(PROD_CONFIG.PARAMETER_LEN)])
         original = PublicKey(root=root, parameter=parameter)
 
         data = original.to_bytes(PROD_CONFIG)
@@ -86,7 +86,7 @@ def test_from_bytes_invalid_length(self) -> None:
     def test_serialization_format(self) -> None:
         """Test that serialization follows the documented format: root || parameter."""
         root = HashDigestVector(data=[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)])
-        parameter = [Fp(value=i + 100) for i in range(TEST_CONFIG.PARAMETER_LEN)]
+        parameter = Parameter(data=[Fp(value=i + 100) for i in range(TEST_CONFIG.PARAMETER_LEN)])
         pk = PublicKey(root=root, parameter=parameter)
 
         data = bytes(pk)
@@ -95,7 +95,7 @@ def test_serialization_format(self) -> None:
         from typing import List, cast
 
         root_data = Fp.serialize_list(cast(List[Fp], list(root.data)))
-        parameter_data = Fp.serialize_list(parameter)
+        parameter_data = Fp.serialize_list(cast(List[Fp], list(parameter.data)))
 
         assert data == root_data + parameter_data
         assert data[: len(root_data)] == root_data
@@ -255,7 +255,7 @@ class TestSerializationProperties:
     def test_public_key_deterministic(self) -> None:
         """Test that serialization is deterministic."""
         root = HashDigestVector(data=[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)])
-        parameter = [Fp(value=i) for i in range(TEST_CONFIG.PARAMETER_LEN)]
+        parameter = Parameter(data=[Fp(value=i) for i in range(TEST_CONFIG.PARAMETER_LEN)])
         pk = PublicKey(root=root, parameter=parameter)
 
         # Serialize multiple times
@@ -291,7 +291,7 @@ def test_different_values_produce_different_bytes(self) -> None:
         """Test that different values produce different serializations."""
         root1 = HashDigestVector(data=[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)])
         root2 = HashDigestVector(data=[Fp(value=i + 1) for i in range(TEST_CONFIG.HASH_LEN_FE)])
-        parameter = [Fp(value=0)] * TEST_CONFIG.PARAMETER_LEN
+        parameter = Parameter(data=[Fp(value=0)] * TEST_CONFIG.PARAMETER_LEN)
 
         pk1 = PublicKey(root=root1, parameter=parameter)
         pk2 = PublicKey(root=root2, parameter=parameter)

tests/lean_spec/subspecs/xmss/test_utils.py

@@ -7,6 +7,7 @@
 
 from lean_spec.subspecs.koalabear.field import Fp, P
 from lean_spec.subspecs.xmss.constants import TEST_CONFIG
+from lean_spec.subspecs.xmss.containers import Parameter
 from lean_spec.subspecs.xmss.merkle_tree import TEST_MERKLE_TREE
 from lean_spec.subspecs.xmss.prf import TEST_PRF
 from lean_spec.subspecs.xmss.tweak_hash import TEST_TWEAK_HASHER
@@ -116,7 +117,9 @@ def test_bottom_tree_from_prf_key() -> None:
     prf_key = TEST_PRF.key_gen()
 
     # Generate a random parameter
-    parameter = [Fp(value=secrets.randbelow(P)) for _ in range(config.PARAMETER_LEN)]
+    parameter = Parameter(
+        data=[Fp(value=secrets.randbelow(P)) for _ in range(config.PARAMETER_LEN)]
+    )
 
     # Generate bottom tree 0
     bottom_tree = bottom_tree_from_prf_key(
@@ -148,7 +151,9 @@ def test_bottom_tree_from_prf_key_deterministic() -> None:
     """Tests that bottom_tree_from_prf_key is deterministic."""
     config = TEST_CONFIG
     prf_key = TEST_PRF.key_gen()
-    parameter = [Fp(value=secrets.randbelow(P)) for _ in range(config.PARAMETER_LEN)]
+    parameter = Parameter(
+        data=[Fp(value=secrets.randbelow(P)) for _ in range(config.PARAMETER_LEN)]
+    )
 
     # Generate the same bottom tree twice
     tree1 = bottom_tree_from_prf_key(
@@ -179,7 +184,9 @@ def test_bottom_tree_from_prf_key_different_indices() -> None:
     """Tests that different bottom tree indices produce different trees."""
     config = TEST_CONFIG
     prf_key = TEST_PRF.key_gen()
-    parameter = [Fp(value=secrets.randbelow(P)) for _ in range(config.PARAMETER_LEN)]
+    parameter = Parameter(
+        data=[Fp(value=secrets.randbelow(P)) for _ in range(config.PARAMETER_LEN)]
+    )
 
     # Generate two different bottom trees
     tree0 = bottom_tree_from_prf_key(