test fixtures: use pre generated keys (#216)

* test fixtures: use pre generated keys

* more pythonic

Author: tcoratger

packages/testing/src/consensus_testing/keys.py

@@ -28,13 +28,9 @@
 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.containers import Signature
 from lean_spec.subspecs.xmss.interface import TEST_SIGNATURE_SCHEME
+from lean_spec.subspecs.xmss.types import HashDigestList, HashTreeOpening, Randomness
 from lean_spec.types import Bytes32, Uint64, ValidatorIndex
 
 from ..keys import XmssKeyManager

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

@@ -6,13 +6,9 @@
 """
 
 from .constants import PROD_CONFIG, TEST_CONFIG
-from .containers import (
-    HashTreeOpening,
-    PublicKey,
-    SecretKey,
-    Signature,
-)
+from .containers import PublicKey, SecretKey, Signature
 from .interface import GeneralizedXmssScheme
+from .types import HashTreeOpening
 
 __all__ = [
     "GeneralizedXmssScheme",

packages/testing/src/consensus_testing/test_keys.json

@@ -1,177 +1,28 @@
-"""Defines the data containers for the Generalized XMSS signature scheme."""
+"""
+Data containers for the Generalized XMSS signature scheme.
+
+This module defines the high-level containers: PublicKey, Signature, and SecretKey.
+Base types (HashDigestVector, Parameter, etc.) are defined in types.py.
+"""
 
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from lean_spec.subspecs.koalabear import Fp
-
 from ...types import Uint64
-from ...types.byte_arrays import BaseBytes
-from ...types.collections import SSZList, SSZVector
 from ...types.container import Container
-from .constants import PRF_KEY_LENGTH, PROD_CONFIG
+from .subtree import HashSubTree
+from .types import (
+    HashDigestList,
+    HashDigestVector,
+    HashTreeOpening,
+    Parameter,
+    PRFKey,
+    Randomness,
+)
 
 if TYPE_CHECKING:
     from .interface import GeneralizedXmssScheme
-    from .subtree import HashSubTree
-
-
-class PRFKey(BaseBytes):
-    """
-    The PRF master secret key.
-
-    This is a high-entropy byte string that acts as the single root secret from
-    which all one-time signing keys are deterministically derived.
-    """
-
-    LENGTH = PRF_KEY_LENGTH
-
-
-HASH_DIGEST_LENGTH = PROD_CONFIG.HASH_LEN_FE
-"""
-The fixed length of a hash digest in field elements.
-
-Derived from `PROD_CONFIG.HASH_LEN_FE`. This corresponds to the output length
-of the Poseidon2 hash function used in the XMSS scheme.
-
-TODO: Make the configuration generic and don't hardcode `PROD_CONFIG`.
-"""
-
-# Calculate the maximum number of nodes in a sparse Merkle tree layer:
-# - A bottom tree has at most 2^(LOG_LIFETIME/2) leaves
-# - With padding, we may add up to 2 additional nodes
-# - To be generous and future-proof, we use 2^(LOG_LIFETIME/2 + 1)
-NODE_LIST_LIMIT = 1 << (PROD_CONFIG.LOG_LIFETIME // 2 + 1)
-"""
-The maximum number of nodes that can be stored in a sparse Merkle tree layer.
-
-Calculated as `2^(LOG_LIFETIME/2 + 1)` from PROD_CONFIG to accommodate:
-- Bottom trees with up to `2^(LOG_LIFETIME/2)` nodes
-- Padding overhead (up to 2 additional nodes)
-- Future-proofing with 2x margin
-
-TODO: Make the configuration generic and don't hardcode `PROD_CONFIG`.
-"""
-
-
-class HashDigestVector(SSZVector):
-    """
-    A single hash digest represented as a fixed-size vector of field elements.
-
-    This is the SSZ-compliant representation of a Poseidon2 hash output.
-    In SSZ notation: `Vector[Fp, HASH_DIGEST_LENGTH]`
-
-    The fixed size enables efficient serialization when used in collections,
-    as SSZ can pack these back-to-back without per-element offsets.
-    """
-
-    ELEMENT_TYPE = Fp
-    LENGTH = HASH_DIGEST_LENGTH
-
-
-class HashDigestList(SSZList):
-    """
-    Variable-length list of hash digests.
-
-    In SSZ notation: `List[Vector[Fp, HASH_DIGEST_LENGTH], NODE_LIST_LIMIT]`
-
-    This type is used to represent collections of hash digests in the XMSS scheme.
-    """
-
-    ELEMENT_TYPE = HashDigestVector
-    LIMIT = NODE_LIST_LIMIT
-
-
-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.
-
-    TODO: Make the configuration generic and don't hardcode `PROD_CONFIG`.
-    """
-
-    ELEMENT_TYPE = Fp
-    LENGTH = PROD_CONFIG.PARAMETER_LEN
-
-
-class Randomness(SSZVector):
-    """
-    The randomness `rho` (ρ) used during signing.
-
-    This value provides a variable input to the message hash, allowing the signer to
-    repeatedly try hashing until a valid "codeword" is found. It must be included in
-    the final signature for the verifier to reproduce the same hash.
-
-    SSZ notation: `Vector[Fp, RAND_LEN_FE]`
-
-    TODO: Make the configuration generic and don't hardcode `PROD_CONFIG`.
-    """
-
-    ELEMENT_TYPE = Fp
-    LENGTH = PROD_CONFIG.RAND_LEN_FE
-
-
-class HashTreeOpening(Container):
-    """
-    A Merkle authentication path.
-
-    This object contains the minimal proof required to connect a specific leaf
-    to the Merkle root. It consists of the list of all sibling nodes along the
-    path from the leaf to the top of the tree.
-
-    SSZ Container with fields:
-    - siblings: List[Vector[Fp, HASH_DIGEST_LENGTH], NODE_LIST_LIMIT]
-    """
-
-    siblings: HashDigestList
-    """SSZ-compliant list of sibling hashes, from bottom to top."""
-
-
-class HashTreeLayer(Container):
-    """
-    Represents a single horizontal "slice" of the sparse Merkle tree.
-
-    Because the tree is sparse, we only store the nodes that are actually computed
-    for the active range of leaves, not the entire conceptual layer.
-    """
-
-    start_index: Uint64
-    """The starting index of the first node in this layer."""
-    nodes: HashDigestList
-    """SSZ-compliant list of hash digests stored for this layer."""
-
-
-LAYERS_LIMIT = PROD_CONFIG.LOG_LIFETIME + 1
-"""
-The maximum number of layers in a subtree.
-
-This is `LOG_LIFETIME + 1` to accommodate all layers from 0 (leaves) to LOG_LIFETIME (root),
-inclusive. For PROD_CONFIG with LOG_LIFETIME=32, this allows up to 33 layers.
-
-TODO: Make the configuration generic and don't hardcode `PROD_CONFIG`.
-"""
-
-
-class HashTreeLayers(SSZList):
-    """
-    Variable-length list of Merkle tree layers.
-
-    In SSZ notation: `List[HashTreeLayer, LAYERS_LIMIT]`
-
-    This type represents the layers of a subtree, from the lowest layer up to the root.
-
-    The number of layers varies based on the subtree structure:
-    - Bottom trees: `LOG_LIFETIME/2` layers
-    - Top trees: `LOG_LIFETIME/2` layers
-    - Maximum: `LOG_LIFETIME + 1` layers
-    """
-
-    ELEMENT_TYPE = HashTreeLayer
-    LIMIT = LAYERS_LIMIT
 
 
 class PublicKey(Container):
@@ -277,7 +128,7 @@ class SecretKey(Container):
     `sqrt(LIFETIME)`, with a minimum of `2 * sqrt(LIFETIME)`.
     """
 
-    top_tree: "HashSubTree"
+    top_tree: HashSubTree
     """
     The top tree containing the root and top `LOG_LIFETIME/2` layers.
 
@@ -297,15 +148,15 @@ class SecretKey(Container):
 
     """
 
-    left_bottom_tree: "HashSubTree"
+    left_bottom_tree: HashSubTree
     """
     The left bottom tree in the sliding window.
 
     This covers epochs:
     [left_bottom_tree_index * sqrt(LIFETIME), (left_bottom_tree_index + 1) * sqrt(LIFETIME))
     """
 
-    right_bottom_tree: "HashSubTree"
+    right_bottom_tree: HashSubTree
     """
     The right bottom tree in the sliding window.
 

src/lean_spec/subspecs/xmss/__init__.py

@@ -25,7 +25,7 @@
     TEST_CONFIG,
     XmssConfig,
 )
-from .containers import HashDigestVector, PublicKey, SecretKey, Signature
+from .containers import PublicKey, SecretKey, Signature
 from .merkle_tree import (
     PROD_MERKLE_TREE,
     TEST_MERKLE_TREE,
@@ -38,6 +38,7 @@
     TEST_TWEAK_HASHER,
     TweakHasher,
 )
+from .types import HashDigestVector
 from .utils import bottom_tree_from_prf_key, expand_activation_time
 
 
@@ -389,7 +390,7 @@ def sign(self, sk: SecretKey, epoch: Uint64, message: bytes) -> Signature:
         # - The Merkle path,
         # - The randomness `rho` needed for verification.
         # Wrap ots_hashes in SSZ types
-        from .containers import HashDigestList, HashDigestVector
+        from .types import HashDigestList, HashDigestVector
 
         ssz_hashes = [HashDigestVector(data=hash_digest) for hash_digest in ots_hashes]
         return Signature(path=path, rho=rho, hashes=HashDigestList(data=ssz_hashes))

src/lean_spec/subspecs/xmss/containers.py

@@ -42,14 +42,6 @@
     TEST_CONFIG,
     XmssConfig,
 )
-from .containers import (
-    HashDigestList,
-    HashDigestVector,
-    HashTreeLayer,
-    HashTreeLayers,
-    HashTreeOpening,
-    Parameter,
-)
 from .rand import PROD_RAND, TEST_RAND, Rand
 from .subtree import HashSubTree
 from .tweak_hash import (
@@ -58,6 +50,14 @@
     TreeTweak,
     TweakHasher,
 )
+from .types import (
+    HashDigestList,
+    HashDigestVector,
+    HashTreeLayer,
+    HashTreeLayers,
+    HashTreeOpening,
+    Parameter,
+)
 from .utils import get_padded_layer
 
 

src/lean_spec/subspecs/xmss/interface.py

@@ -47,12 +47,12 @@
     TWEAK_PREFIX_MESSAGE,
     XmssConfig,
 )
-from .containers import Parameter, Randomness
 from .hypercube import (
     hypercube_find_layer,
     hypercube_part_size,
     map_to_vertex,
 )
+from .types import Parameter, Randomness
 from .utils import int_to_base_p
 
 

src/lean_spec/subspecs/xmss/merkle_tree.py

@@ -24,7 +24,7 @@
     TEST_CONFIG,
     XmssConfig,
 )
-from .containers import PRFKey, Randomness
+from .types import PRFKey, Randomness
 
 PRF_DOMAIN_SEP: bytes = bytes(
     [

src/lean_spec/subspecs/xmss/message_hash.py

@@ -9,7 +9,7 @@
 
 from ..koalabear import Fp, P
 from .constants import PROD_CONFIG, TEST_CONFIG, XmssConfig
-from .containers import Parameter, Randomness
+from .types import Parameter, Randomness
 
 
 class Rand(StrictBaseModel):