xmss: catchup with the latest updates (leanEthereum#152)

* xmss: catchup with the latest updates

* more updates

* more cleanup

* fix linter

* fix tests

* fix tests

* consensus testing: try a fix for the keys

* test fix

* revert

* wip

* some fix

* fix test

* touchups

Author: tcoratger

packages/testing/src/consensus_testing/keys.py

@@ -6,18 +6,28 @@
 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.interface import DEFAULT_SIGNATURE_SCHEME
+from lean_spec.subspecs.xmss.interface import (
+    TEST_SIGNATURE_SCHEME as DEFAULT_SIGNATURE_SCHEME,
+)
 from lean_spec.types import ValidatorIndex
 
 
 class KeyPair(NamedTuple):
-    """A validator’s XMSS key pair."""
+    """A validator's XMSS key pair."""
 
     public: PublicKey
-    """The validator’s public key (used for verification)."""
+    """The validator's public key (used for verification)."""
 
     secret: SecretKey
-    """The validator’s secret key (used for signing)."""
+    """The validator's secret key (used for signing)."""
+
+
+_KEY_CACHE: dict[tuple[int, int], KeyPair] = {}
+"""
+Cache keys across tests to avoid regenerating them for the same validator/lifetime combo.
+
+Key: (validator_index, num_active_epochs) -> KeyPair
+"""
 
 
 class XmssKeyManager:
@@ -76,6 +86,13 @@ def __getitem__(self, validator_index: ValidatorIndex) -> KeyPair:
         # - We include slot 0 (genesis) in the count
         num_active_epochs = self.max_slot.as_int() + 1
 
+        # Check global cache first (keys are reused across tests)
+        cache_key = (int(validator_index), num_active_epochs)
+        if cache_key in _KEY_CACHE:
+            key_pair = _KEY_CACHE[cache_key]
+            self._key_pairs[validator_index] = key_pair
+            return key_pair
+
         # Generate the key pair using the default XMSS scheme.
         #
         # The seed is set to 0 for deterministic test keys.
@@ -85,12 +102,13 @@ def __getitem__(self, validator_index: ValidatorIndex) -> KeyPair:
 
         # Store as a cohesive unit and return.
         key_pair = KeyPair(public=pk, secret=sk)
+        _KEY_CACHE[cache_key] = key_pair  # Cache globally for reuse across tests
         self._key_pairs[validator_index] = key_pair
         return key_pair
 
     def sign_attestation(self, attestation: Attestation) -> Signature:
         """
-        Sign an attestation with the validator’s XMSS key.
+        Sign an attestation with the validator's XMSS key.
 
         Parameters
         ----------
@@ -113,22 +131,50 @@ def sign_attestation(self, attestation: Attestation) -> Signature:
 
         # Lazy key retrieval: creates keys if first time seeing this validator.
         key_pair = self[validator_id]
-
-        # Compute the message digest from the attestation's SSZ tree root.
-        #
-        # This produces a cryptographic hash of the entire attestation structure.
-        message = bytes(hash_tree_root(attestation))
+        # Get the current secret key
+        sk = key_pair.secret
 
         # Map the attestation slot to an XMSS epoch.
         #
         # Each slot gets its own epoch to avoid key reuse.
         epoch = attestation.data.slot
 
-        # Generate the XMSS signature using the validator's secret key.
-        xmss_sig = DEFAULT_SIGNATURE_SCHEME.sign(key_pair.secret, epoch, message)
+        # Advance the key's prepared window until it covers the target epoch.
+        #
+        # We use the scheme that the key was generated with.
+        scheme = DEFAULT_SIGNATURE_SCHEME
+
+        # Loop until the epoch is inside the prepared interval
+        prepared_interval = scheme.get_prepared_interval(sk)
+        while int(epoch) not in prepared_interval:
+            # Check if we're advancing past the key's total lifetime
+            activation_interval = scheme.get_activation_interval(sk)
+            if prepared_interval.stop >= activation_interval.stop:
+                raise ValueError(
+                    f"Cannot sign for epoch {epoch}: "
+                    f"it is beyond the key's max lifetime {activation_interval.stop}"
+                )
+
+            # Advance the key and get the new key object
+            sk = scheme.advance_preparation(sk)
+
+            # Update the prepared interval for the next loop check
+            prepared_interval = scheme.get_prepared_interval(sk)
+
+        # Update the cached key pair with the new, advanced secret key.
+        # This ensures the *next* call to sign() uses the advanced state.
+        self._key_pairs[validator_id] = KeyPair(public=key_pair.public, secret=sk)
+
+        # Compute the message digest from the attestation's SSZ tree root.
+        #
+        # 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 = scheme.sign(sk, epoch, message)
 
         # Convert the signature to the wire format (byte array).
-        signature_bytes = xmss_sig.to_bytes(DEFAULT_SIGNATURE_SCHEME.config)
+        signature_bytes = xmss_sig.to_bytes(scheme.config)
 
         # Ensure the signature meets the consensus spec length (3100 bytes).
         #

src/lean_spec/subspecs/xmss/constants.py

@@ -127,11 +127,11 @@ def SIGNATURE_LEN_BYTES(self) -> int:  # noqa: N802
 TEST_CONFIG: Final = XmssConfig(
     MESSAGE_LENGTH=32,
     LOG_LIFETIME=8,
-    DIMENSION=16,
+    DIMENSION=4,
     BASE=4,
-    FINAL_LAYER=24,
-    TARGET_SUM=24,
-    MAX_TRIES=100_000,
+    FINAL_LAYER=6,
+    TARGET_SUM=6,
+    MAX_TRIES=1_000,
     PARAMETER_LEN=5,
     TWEAK_LEN_FE=2,
     MSG_LEN_FE=9,

src/lean_spec/subspecs/xmss/containers.py

@@ -12,6 +12,7 @@
 
 if TYPE_CHECKING:
     from .constants import XmssConfig
+    from .subtree import HashSubTree
 
 PRFKey = Annotated[bytes, Field(min_length=PRF_KEY_LENGTH, max_length=PRF_KEY_LENGTH)]
 """
@@ -112,10 +113,11 @@ class HashTreeLayer(StrictBaseModel):
 
 class HashTree(StrictBaseModel):
     """
-    The pre-computed, stored portion of the sparse Merkle tree.
+    A simple representation of a sparse Merkle tree.
 
-    This structure is part of the `SecretKey` and contains all the necessary nodes
-    to generate an authentication path for any signature within the key's active lifetime.
+    This structure contains the necessary nodes to generate an authentication path
+    for any signature within a key's active lifetime. For production use with
+    long lifetimes, prefer `HashSubTree` with the top-bottom tree approach.
     """
 
     depth: int
@@ -370,11 +372,61 @@ class SecretKey(StrictBaseModel):
 
     prf_key: PRFKey
     """The master secret key used to derive all one-time secrets."""
-    tree: HashTree
-    """The pre-computed sparse Merkle tree needed to generate authentication paths."""
+
     parameter: Parameter
     """The public parameter `P`, stored for convenience during signing."""
+
     activation_epoch: int
-    """The first epoch for which this secret key is valid."""
+    """
+    The first epoch for which this secret key is valid.
+
+    Note: With top-bottom trees, this is aligned to a multiple of `sqrt(LIFETIME)`
+    to ensure efficient tree partitioning.
+    """
+
     num_active_epochs: int
-    """The number of consecutive epochs this key can be used for."""
+    """
+    The number of consecutive epochs this key can be used for.
+
+    Note: With top-bottom trees, this is rounded up to be a multiple of
+    `sqrt(LIFETIME)`, with a minimum of `2 * sqrt(LIFETIME)`.
+    """
+
+    top_tree: HashSubTree | None = None
+    """
+    The top tree containing the root and top `LOG_LIFETIME/2` layers.
+
+    This tree is always kept in memory and contains the roots of all bottom trees
+    in its lowest layer. Its root is the public key's Merkle root.
+    """
+
+    left_bottom_tree_index: int | None = None
+    """
+    The index of the left bottom tree in the sliding window.
+
+    Bottom trees are numbered 0, 1, 2, ... where tree `i` covers epochs
+    `[i * sqrt(LIFETIME), (i+1) * sqrt(LIFETIME))`.
+
+    The prepared interval is:
+    [left_bottom_tree_index * sqrt(LIFETIME), (left_bottom_tree_index + 2) * sqrt(LIFETIME))
+
+    """
+
+    left_bottom_tree: HashSubTree | None = None
+    """
+    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 | None = None
+    """
+    The right bottom tree in the sliding window.
+
+    This covers epochs:
+    [(left_bottom_tree_index + 1) * sqrt(LIFETIME), (left_bottom_tree_index + 2) * sqrt(LIFETIME))
+
+    Together with `left_bottom_tree`, this provides a prepared interval of
+    exactly `2 * sqrt(LIFETIME)` consecutive epochs.
+    """