chore: remove default xmss signature scheme (leanEthereum#159)

Author: anshalshukla

packages/testing/src/consensus_testing/keys.py

@@ -7,7 +7,8 @@
 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 (
-    TEST_SIGNATURE_SCHEME as DEFAULT_SIGNATURE_SCHEME,
+    TEST_SIGNATURE_SCHEME,
+    GeneralizedXmssScheme,
 )
 from lean_spec.types import ValidatorIndex
 
@@ -36,7 +37,11 @@ class XmssKeyManager:
     DEFAULT_MAX_SLOT = Slot(100)
     """Default maximum slot horizon if not specified."""
 
-    def __init__(self, max_slot: Optional[Slot] = None) -> None:
+    def __init__(
+        self,
+        max_slot: Optional[Slot] = None,
+        scheme: GeneralizedXmssScheme = TEST_SIGNATURE_SCHEME,
+    ) -> None:
         """
         Initialize the key manager.
 
@@ -45,13 +50,17 @@ def __init__(self, max_slot: Optional[Slot] = None) -> None:
         max_slot : Slot, optional
             Highest slot number for which keys must remain valid.
             Defaults to `Slot(100)`.
+        scheme : GeneralizedXmssScheme, optional
+            The XMSS scheme to use.
+            Defaults to `TEST_SIGNATURE_SCHEME`.
 
         Notes:
         -----
         Internally, keys are stored in a single dictionary:
         `{ValidatorIndex → KeyPair}`.
         """
         self.max_slot = max_slot if max_slot is not None else self.DEFAULT_MAX_SLOT
+        self.scheme = scheme
         self._key_pairs: dict[ValidatorIndex, KeyPair] = {}
 
     def __getitem__(self, validator_index: ValidatorIndex) -> KeyPair:
@@ -98,7 +107,7 @@ def __getitem__(self, validator_index: ValidatorIndex) -> KeyPair:
         # The seed is set to 0 for deterministic test keys.
         from lean_spec.types import Uint64
 
-        pk, sk = DEFAULT_SIGNATURE_SCHEME.key_gen(Uint64(0), Uint64(num_active_epochs))
+        pk, sk = self.scheme.key_gen(Uint64(0), Uint64(num_active_epochs))
 
         # Store as a cohesive unit and return.
         key_pair = KeyPair(public=pk, secret=sk)
@@ -139,27 +148,22 @@ def sign_attestation(self, attestation: Attestation) -> Signature:
         # Each slot gets its own epoch to avoid key reuse.
         epoch = attestation.data.slot
 
-        # 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)
+        prepared_interval = self.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)
+            activation_interval = self.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)
+            sk = self.scheme.advance_preparation(sk)
 
             # Update the prepared interval for the next loop check
-            prepared_interval = scheme.get_prepared_interval(sk)
+            prepared_interval = self.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.
@@ -171,10 +175,10 @@ def sign_attestation(self, attestation: Attestation) -> Signature:
         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)
+        xmss_sig = self.scheme.sign(sk, epoch, message)
 
         # Convert the signature to the wire format (byte array).
-        signature_bytes = xmss_sig.to_bytes(scheme.config)
+        signature_bytes = xmss_sig.to_bytes(self.scheme.config)
 
         # Ensure the signature meets the consensus spec length (3100 bytes).
         #

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

@@ -27,7 +27,7 @@
 from lean_spec.subspecs.containers.state.state import State
 from lean_spec.subspecs.forkchoice import Store
 from lean_spec.subspecs.ssz import hash_tree_root
-from lean_spec.subspecs.xmss.interface import DEFAULT_SIGNATURE_SCHEME
+from lean_spec.subspecs.xmss.interface import TEST_SIGNATURE_SCHEME
 from lean_spec.types import Bytes32, Uint64, ValidatorIndex
 
 from ..keys import XmssKeyManager
@@ -186,15 +186,15 @@ def make_fixture(self) -> ForkChoiceTest:
         key_manager = (
             shared_key_manager
             if self.max_slot <= shared_key_manager.max_slot
-            else XmssKeyManager(max_slot=self.max_slot)
+            else XmssKeyManager(max_slot=self.max_slot, scheme=TEST_SIGNATURE_SCHEME)
         )
 
         # Update validator pubkeys to match key_manager's generated keys
         updated_validators = [
             validator.model_copy(
                 update={
                     "pubkey": key_manager[ValidatorIndex(i)].public.to_bytes(
-                        DEFAULT_SIGNATURE_SCHEME.config
+                        key_manager.scheme.config
                     )
                 }
             )

src/lean_spec/subspecs/containers/signature.py

@@ -6,16 +6,21 @@
 
 from ..xmss.containers import PublicKey
 from ..xmss.containers import Signature as XmssSignature
-from ..xmss.interface import DEFAULT_SIGNATURE_SCHEME
+from ..xmss.interface import TEST_SIGNATURE_SCHEME, GeneralizedXmssScheme
 
 
 class Signature(Bytes3100):
     """Represents aggregated signature produced by the leanVM (SNARKs in the future)."""
 
-    def verify(self, public_key: PublicKey, epoch: Uint64, message: bytes) -> bool:
+    def verify(
+        self,
+        public_key: PublicKey,
+        epoch: Uint64,
+        message: bytes,
+        scheme: GeneralizedXmssScheme = TEST_SIGNATURE_SCHEME,
+    ) -> bool:
         """Verify the signature using XMSS verification algorithm."""
         try:
-            scheme = DEFAULT_SIGNATURE_SCHEME
             # Signature container is always 3100 bytes, but scheme config may expect less.
             # Slice to the expected size if needed, assumes padding to the right.
             signature_data = bytes(self)[: scheme.config.SIGNATURE_LEN_BYTES]

src/lean_spec/subspecs/containers/validator.py

@@ -3,7 +3,7 @@
 from lean_spec.types import Bytes52, Container
 
 from ..xmss.containers import PublicKey
-from ..xmss.interface import DEFAULT_SIGNATURE_SCHEME
+from ..xmss.interface import TEST_SIGNATURE_SCHEME, GeneralizedXmssScheme
 
 
 class Validator(Container):
@@ -12,6 +12,6 @@ class Validator(Container):
     pubkey: Bytes52
     """XMSS one-time signature public key."""
 
-    def get_pubkey(self) -> PublicKey:
+    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), DEFAULT_SIGNATURE_SCHEME.config)
+        return PublicKey.from_bytes(bytes(self.pubkey), scheme.config)

src/lean_spec/subspecs/xmss/interface.py

@@ -589,6 +589,3 @@ def advance_preparation(self, sk: SecretKey) -> SecretKey:
     TEST_RAND,
 )
 """A lightweight instance for test environments."""
-
-DEFAULT_SIGNATURE_SCHEME = PROD_SIGNATURE_SCHEME
-"""The default signature scheme to use."""

tests/conftest.py

@@ -7,7 +7,3 @@
 # Create a profile named "no_deadline" with deadline disabled.
 settings.register_profile("no_deadline", deadline=None)
 settings.load_profile("no_deadline")
-
-# Enable test mode for XMSS signature scheme
-# This uses smaller parameters (TEST_CONFIG) for faster test execution
-xmss.interface.DEFAULT_SIGNATURE_SCHEME = xmss.interface.TEST_SIGNATURE_SCHEME