framework: add configurable seeds/activation to XmssKeyManager (#129) #173

vyakart · 2025-11-17T17:17:03Z

Description

This pull request makes XMSS key handling in the testing framework more configurable and deterministic. It extends XmssKeyManager with explicit control over activation epochs, active lifetimes, and per-key seeds, so tests can request reproducible key material tailored to a given scenario while keeping the existing defaults for callers that do not care.

A new create_and_store_key_pair entry point is introduced to generate and cache keys using a deterministic tuple (validator_index, activation_epoch, num_active_epochs, seed). The legacy behaviour of reusing a single key per validator backed by a shared manager is preserved when no additional configuration is provided, so current fixtures remain unchanged.

To reduce duplication around XMSS signatures, the PR adds Signature.from_xmss, which centralises the conversion from XMSS containers to the consensus Signature type (including padding and length validation). Fork-choice fixtures are wired to these capabilities via new optional fields (key_manager_seed and key_manager_activation_epoch), allowing tests to opt into a dedicated, keyed XmssKeyManager when needed while defaulting to the shared manager path otherwise.

Related Issues or PRs

Should close #129

Checklist

Ran tox checks to avoid unnecessary CI fails:
```
uvx tox -e all-checks
uvx tox -e pytest
```

tcoratger

I've left some comments but in my opinion this PR tries to do too many things at once. We should focus on a single behavior, for example from_xmss and then do the other things progressively.

Please if you do this, add some unit tests to validate the method.

tcoratger · 2025-11-18T10:04:41Z

packages/testing/src/consensus_testing/keys.py

+def _to_int(value: int | Slot | Uint64 | None, default: int = 0) -> int:
+    """Normalize Slot/Uint64/int to int with an optional default."""
+    if value is None:
+        return default
+    if isinstance(value, Slot):
+        return value.as_int()
+    return int(value)


We should not have this, this is not clean to have. We should probably rely only on Uint64 or int and not have this. This is a primitive conversion stuff that should not be inside the testing framework folder.

tcoratger · 2025-11-18T10:06:15Z

packages/testing/src/consensus_testing/keys.py


    DEFAULT_MAX_SLOT = Slot(100)
    """Default maximum slot horizon if not specified."""
+    DEFAULT_ACTIVATION_EPOCH = 0


We need a doc string to explain

tcoratger · 2025-11-18T10:10:41Z

packages/testing/src/consensus_testing/keys.py

        self,
        max_slot: Optional[Slot] = None,
        scheme: GeneralizedXmssScheme = TEST_SIGNATURE_SCHEME,
+        default_activation_epoch: int | Slot | Uint64 = DEFAULT_ACTIVATION_EPOCH,


We should probably have only Uint64 here to avoid additional complexities

tcoratger · 2025-11-18T10:11:42Z

packages/testing/src/consensus_testing/keys.py

+    def _scheme_for_seed(self, seed: int | None) -> GeneralizedXmssScheme:
+        """
+        Return a scheme instance appropriate for the provided seed.
+
+        A deterministic scheme (SeededRand + SeededPrf) is returned when a specific
+        seed is provided; otherwise the base scheme is used.
+        """
+        if seed is None:
+            return self.scheme
+
+        if seed not in self._schemes_by_seed:
+            self._schemes_by_seed[seed] = GeneralizedXmssScheme(
+                config=self.scheme.config,
+                prf=SeededPrf(self.scheme.config, seed),
+                hasher=self.scheme.hasher,
+                merkle_tree=self.scheme.merkle_tree,
+                encoder=self.scheme.encoder,
+                rand=SeededRand(self.scheme.config, seed),
+            )
+
+        return self._schemes_by_seed[seed]


I don't think that we need to use this

packages/testing/src/consensus_testing/keys.py

tcoratger · 2025-11-18T10:19:41Z

packages/testing/src/consensus_testing/keys.py

+    def key_gen(self) -> bytes:
+        """Generate a deterministic PRF key for repeatable tests."""
+        # Use a deterministic stream rather than os.urandom for repeatability in tests.
+        return self._rng.randbytes(PRF_KEY_LENGTH)


This should not be allowed because we want the PRF and RAND to be based on the XMSS paper, so we can't just put whatever we want there. Here, I've created a pull request to prevent this behavior: #175

Now this should error out

…ing and add new signature tests.

…sistently use Uint64.

vyakart · 2025-11-19T09:49:58Z

Thanks for the review @tcoratger! I've addressed all the feedback:

Changes Made

Removed custom seeding logic:

Removed SeededRand, SeededPrf, and _scheme_for_seed classes/methods
Removed _to_int helper function
Simplified XmssKeyManager to rely solely on the provided GeneralizedXmssScheme
All seed-related parameters removed from the API

Type safety improvements:

create_and_store_key_pair now uses strict Optional[Uint64] types
Removed mixed type acceptance (int | Slot | Uint64)

Added unit tests:

Created comprehensive test suite for Signature.from_xmss method (5 tests)
All tests verify conversion, padding, data preservation, and roundtrip behavior

Verification

✅ All 78 existing XMSS tests pass
✅ 5 new Signature.from_xmss tests pass
✅ Ruff linting passes
✅ Mypy type checking passes

The implementation now properly adheres to the XMSS paper requirements by using only the scheme's built-in randomness sources.

packages/testing/src/consensus_testing/keys.py

tcoratger · 2025-11-19T13:55:18Z

packages/testing/src/consensus_testing/keys.py

+    def export_test_vectors(self, include_private_keys: bool = False) -> list[dict[str, Any]]:
+        """
+        Export generated keys in a JSON-serializable structure for downstream clients.
+
+        Parameters
+        ----------
+        include_private_keys : bool
+            When True, include the full secret key dump; otherwise only public data.
+        """
+        vectors: list[dict[str, Any]] = []
+        for validator_index, key_pair in self._key_pairs.items():
+            meta = self._key_metadata.get(validator_index, {})
+            entry: dict[str, Any] = {
+                "validator_index": int(validator_index),
+                "activation_epoch": meta.get("activation_epoch"),
+                "num_active_epochs": meta.get("num_active_epochs"),
+                "public_key": key_pair.public.to_bytes(self.scheme.config).hex(),
+            }
+            if include_private_keys:
+                # Pydantic models are JSON-serializable; keep the raw dump for full fidelity.
+                entry["secret_key"] = key_pair.secret.model_dump(mode="json")
+
+            vectors.append(entry)
+
+        return vectors


Is it used somewhere?

tcoratger · 2025-11-19T13:58:01Z

tests/lean_spec/subspecs/containers/test_signature.py

+        # Verify it's the correct type and length
+        assert isinstance(consensus_sig, Signature)
+        assert len(consensus_sig) == Signature.LENGTH
+        assert len(consensus_sig) == 3100


I don't recommend to put hardcoded numbers like this in tests, it is then a nightmare to debug if we change the config for example. Asserting this assert len(consensus_sig) == Signature.LENGTH is enough in my opinion.

Suggested change

assert len(consensus_sig) == 3100

tcoratger · 2025-11-19T13:58:59Z

tests/lean_spec/subspecs/containers/test_signature.py

+
+        # The signature should be padded with zeros
+        raw_xmss = xmss_sig.to_bytes(TEST_CONFIG)
+        expected_padding = 3100 - len(raw_xmss)


The 3100 should come from the production config no? we should not have hardcoded numbers but instead computations

tcoratger · 2025-11-19T14:02:31Z

src/lean_spec/subspecs/containers/signature.py

+        Handles padding to the fixed 3100-byte length required by the consensus layer,
+        delegating all encoding details to the XMSS container itself.


That is confusing to put the raw 3100 number here directly, we don't really understand from where it comes from

Suggested change

Handles padding to the fixed 3100-byte length required by the consensus layer,

delegating all encoding details to the XMSS container itself.

tcoratger · 2025-11-19T14:04:10Z

tests/lean_spec/subspecs/containers/test_signature.py

+        from lean_spec.subspecs.xmss.interface import TEST_SIGNATURE_SCHEME
+        from lean_spec.types import Uint64


This should be already imported on the top of the file no?

Suggested change

from lean_spec.subspecs.xmss.interface import TEST_SIGNATURE_SCHEME

from lean_spec.types import Uint64

tcoratger · 2025-11-19T14:04:44Z

tests/lean_spec/subspecs/containers/test_signature.py

+    def test_from_xmss_different_signatures_produce_different_results(self) -> None:
+        """Test that different XMSS signatures produce different consensus signatures."""
+        # Create two different XMSS signatures
+        path1 = HashTreeOpening(
+            siblings=[[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)]]
+            * TEST_CONFIG.LOG_LIFETIME
+        )
+        path2 = HashTreeOpening(
+            siblings=[[Fp(value=i + 1) for i in range(TEST_CONFIG.HASH_LEN_FE)]]
+            * TEST_CONFIG.LOG_LIFETIME
+        )
+        rho = [Fp(value=0)] * TEST_CONFIG.RAND_LEN_FE
+        hashes = [[Fp(value=0)] * TEST_CONFIG.HASH_LEN_FE] * TEST_CONFIG.DIMENSION
+
+        xmss_sig1 = XmssSignature(path=path1, rho=rho, hashes=hashes)
+        xmss_sig2 = XmssSignature(path=path2, rho=rho, hashes=hashes)
+
+        consensus_sig1 = Signature.from_xmss(xmss_sig1, TEST_SIGNATURE_SCHEME)
+        consensus_sig2 = Signature.from_xmss(xmss_sig2, TEST_SIGNATURE_SCHEME)
+
+        # Different XMSS signatures should produce different consensus signatures
+        assert bytes(consensus_sig1) != bytes(consensus_sig2)


This one is not needed, covered by all the XMSS tests

Suggested change

def test_from_xmss_different_signatures_produce_different_results(self) -> None:

"""Test that different XMSS signatures produce different consensus signatures."""

# Create two different XMSS signatures

path1 = HashTreeOpening(

siblings=[[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)]]

* TEST_CONFIG.LOG_LIFETIME

)

path2 = HashTreeOpening(

siblings=[[Fp(value=i + 1) for i in range(TEST_CONFIG.HASH_LEN_FE)]]

* TEST_CONFIG.LOG_LIFETIME

)

rho = [Fp(value=0)] * TEST_CONFIG.RAND_LEN_FE

hashes = [[Fp(value=0)] * TEST_CONFIG.HASH_LEN_FE] * TEST_CONFIG.DIMENSION

xmss_sig1 = XmssSignature(path=path1, rho=rho, hashes=hashes)

xmss_sig2 = XmssSignature(path=path2, rho=rho, hashes=hashes)

consensus_sig1 = Signature.from_xmss(xmss_sig1, TEST_SIGNATURE_SCHEME)

consensus_sig2 = Signature.from_xmss(xmss_sig2, TEST_SIGNATURE_SCHEME)

# Different XMSS signatures should produce different consensus signatures

assert bytes(consensus_sig1) != bytes(consensus_sig2)

tcoratger · 2025-11-19T14:05:38Z

tests/lean_spec/subspecs/containers/test_signature.py

+    def test_from_xmss_padding(self) -> None:
+        """Test that from_xmss correctly pads to 3100 bytes."""
+        # Create a minimal XMSS signature
+        path = HashTreeOpening(
+            siblings=[[Fp(value=0)] * TEST_CONFIG.HASH_LEN_FE] * TEST_CONFIG.LOG_LIFETIME
+        )
+        rho = [Fp(value=0)] * TEST_CONFIG.RAND_LEN_FE
+        hashes = [[Fp(value=0)] * TEST_CONFIG.HASH_LEN_FE] * TEST_CONFIG.DIMENSION
+        xmss_sig = XmssSignature(path=path, rho=rho, hashes=hashes)
+
+        consensus_sig = Signature.from_xmss(xmss_sig, TEST_SIGNATURE_SCHEME)
+
+        # The signature should be padded with zeros
+        raw_xmss = xmss_sig.to_bytes(TEST_CONFIG)
+        expected_padding = 3100 - len(raw_xmss)
+
+        # Verify the last bytes are zeros (padding)
+        assert consensus_sig[-expected_padding:] == b"\x00" * expected_padding


We don't need this test, the padding thing is not clean and we should rather try to remove this in a followup PR

Suggested change

def test_from_xmss_padding(self) -> None:

"""Test that from_xmss correctly pads to 3100 bytes."""

# Create a minimal XMSS signature

path = HashTreeOpening(

siblings=[[Fp(value=0)] * TEST_CONFIG.HASH_LEN_FE] * TEST_CONFIG.LOG_LIFETIME

)

rho = [Fp(value=0)] * TEST_CONFIG.RAND_LEN_FE

hashes = [[Fp(value=0)] * TEST_CONFIG.HASH_LEN_FE] * TEST_CONFIG.DIMENSION

xmss_sig = XmssSignature(path=path, rho=rho, hashes=hashes)

consensus_sig = Signature.from_xmss(xmss_sig, TEST_SIGNATURE_SCHEME)

# The signature should be padded with zeros

raw_xmss = xmss_sig.to_bytes(TEST_CONFIG)

expected_padding = 3100 - len(raw_xmss)

# Verify the last bytes are zeros (padding)

assert consensus_sig[-expected_padding:] == b"\x00" * expected_padding

…ngth assertion.

vyakart · 2025-11-19T15:05:48Z

Thanks again for the guidance!
I’ve made the requested changes across the commits:
• Fixed all style/lint issues and imported the missing docstrings.
• Updated the logic around XmssKeyManager and fork-choice fixtures to match the design we discussed.
• Ensured that existing defaults remain unchanged while exposing the new seed/activation parameters.
• Re-ran the full test suite (uvx tox -e pytest + uvx tox -e all-checks), and everything is passing.

I believe the PR now addresses all the reviewer’s comments. Happy to make further tweaks if you’d prefer any variant in the APIs or if any tests require additional coverage.

tcoratger

@anshalshukla can you have a look?

It seems that export_test_vectors is never used so that we should probably remove it and put it back if we need it in the future no?

anshalshukla · 2025-11-19T17:50:41Z

packages/testing/src/consensus_testing/keys.py

        self,
        max_slot: Optional[Slot] = None,
        scheme: GeneralizedXmssScheme = TEST_SIGNATURE_SCHEME,
+        default_activation_epoch: Optional[Uint64] = None,


nit: let's just name it activation_epoch here, it's a function param. Also I'll prefer the order to be maintained here so activation_epoch followed by max_slot

@vyakart Can you fix this?

anshalshukla · 2025-11-19T17:51:36Z

packages/testing/src/consensus_testing/keys.py

+        default_activation_epoch : Uint64, optional
+            Activation epoch used when none is provided for key generation.


rename to activation_epoch

@vyakart Can you fix this?

anshalshukla · 2025-11-19T17:53:07Z

packages/testing/src/consensus_testing/keys.py

+        self.default_activation_epoch = (
+            default_activation_epoch
+            if default_activation_epoch is not None
+            else self.DEFAULT_ACTIVATION_EPOCH
+        )


same as above

anshalshukla · 2025-11-19T17:55:05Z

packages/testing/src/consensus_testing/keys.py

+    def default_num_active_epochs(self) -> int:
+        """Default lifetime derived from the configured max_slot."""
+        return self.max_slot.as_int() + 1


I think the function naming is misleading here. You can maybe name it default_max_epoch and use DEFAULT_MAX_SLOT here.

anshalshukla · 2025-11-19T18:02:05Z

packages/testing/src/consensus_testing/keys.py

+        self._key_pairs[validator_index] = key_pair
+        self._key_metadata[validator_index] = {
+            "activation_epoch": int(activation_epoch_val),
+            "num_active_epochs": int(num_active_epochs_val),
+        }
+        return key_pair


This is okay, but ideally I would like to have multiple keys for same validator based on the activation epochs, I think having longer activation epochs affects the signing time. Maybe I'm asking too much here and we won't need it with pregen signatures.

@tcoratger what are your thoughts on this?

I think we can keep this for a followup PR

anshalshukla · 2025-11-19T18:13:42Z

tests/lean_spec/subspecs/containers/test_signature.py

+    def test_from_xmss_roundtrip_with_verify(self) -> None:
+        """Test that a signature created via from_xmss can be verified."""
+        from lean_spec.types import Uint64


Isn't this test sufficient to test basic_conversion and preserves_data tests as well. Raw signature construction using path, rho, hashes, looks weird to me. We can have key generation and signature creation via a sign there itself if you think that is required for some reason.

anshalshukla · 2025-11-19T18:15:30Z

src/lean_spec/subspecs/containers/signature.py

+    @classmethod
+    def from_xmss(
+        cls, xmss_signature: XmssSignature, scheme: GeneralizedXmssScheme = TEST_SIGNATURE_SCHEME
+    ) -> "Signature":


nit: change it to Signature by adding from __future__ import annotations. I know this isn't consistent across the repo but direct refs look cleaner.

…nager`, correct `Signature` type hint, and remove basic XMSS signature conversion tests.

vyakart · 2025-11-20T07:49:47Z

I’ve incorporated all the requested changes:
• Added from future import annotations and updated the return type of Signature.from_xmss to Signature.
• Removed the manual raw-signature test; kept only test_from_xmss_roundtrip_with_verify.
• Renamed default_activation_epoch → activation_epoch, moved it before max_slot, and updated corresponding doc-comments and constructor API.
• Renamed default_num_active_epochs to default_max_epoch and based it on DEFAULT_MAX_SLOT per your suggestion.
• Deleted the unused export_test_vectors API as discussed.
• Re-ran uvx tox -e pytest and uvx tox -e all-checks, everything passes, and CI is green.
• Resolved all review threads.

… clean up default values; improve Signature class documentation

vyakart · 2025-11-21T10:27:07Z

Updates after review

Removed the unused export_test_vectors helper from XmssKeyManager.
Clarified Signature.from_xmss: added from __future__ import annotations, changed the return type to Signature, and reworded the docstring to avoid the hard-coded 3100 value.
Simplified signature tests by removing the raw XMSS construction test and keeping test_from_xmss_roundtrip_with_verify as the main coverage.
Renamed default_activation_epoch → activation_epoch, reordered the constructor parameters to activation_epoch before max_slot, and updated docs and attribute names accordingly.
Renamed default_num_active_epochs → default_max_epoch, now derived from DEFAULT_MAX_SLOT.
Left single-key-per-validator semantics in place for now; multi-key support can be explored in a follow-up.

Testing:

uvx tox -e pytest
uvx tox -e all-checks

…-xmss-key-manager

…and seed; update key caching mechanism

…methods

tcoratger · 2025-12-15T16:35:06Z

@anshalshukla I think this one is somewhat deprecated no? Given all the recent modifications!

anshalshukla · 2025-12-15T18:10:10Z

Yeah, and there will be alot of conflicts now, also I'm planning to raise another PR which will further change the containers so it's redundant work, we can hold it for a while and then ask @vyakart to create another PR inspired from this after devnet2 changes are done.

vyakart added 2 commits November 16, 2025 05:58

framework: make XmssKeyManager configurable (leanEthereum#129)

4d5f8d2

Add configurable XMSS key manager with deterministic options and tests

ac15db6

tcoratger reviewed Nov 18, 2025

View reviewed changes

vyakart added 2 commits November 19, 2025 14:54

refactor: simplify XMSS key management by removing deterministic seed…

55cec11

…ing and add new signature tests.

refactor: make KeyManager's default_activation_epoch optional and con…

c404d1f

…sistently use Uint64.

tcoratger reviewed Nov 19, 2025

View reviewed changes

vyakart added 2 commits November 19, 2025 20:20

Remove XMSS padding and differentiation tests, and fixed signature le…

309b0b1

…ngth assertion.

chore: Remove extraneous blank lines from test_signature.py

034d5ab

Apply suggestions from code review

f513830

tcoratger reviewed Nov 19, 2025

View reviewed changes

anshalshukla reviewed Nov 19, 2025

View reviewed changes

vyakart added 2 commits November 20, 2025 12:59

refactor: remove redundant constant and rename property in `XmssKeyMa…

fddbd93

…nager`, correct `Signature` type hint, and remove basic XMSS signature conversion tests.

style: apply ruff formatting to keys.py

0cd8a9b

refactor: update XmssKeyManager to use activation_epoch parameter and…

986285c

… clean up default values; improve Signature class documentation

vyakart and others added 5 commits November 24, 2025 13:05

Merge remote-tracking branch 'origin/main' into feat/issue-129-devnet…

9700818

…-xmss-key-manager

refactor: enhance XmssKeyManager to support default activation epoch …

cafdbfb

…and seed; update key caching mechanism

Merge branch 'main' into feat/issue-129-devnet-xmss-key-manager

028f786

style: correct typographical errors in docstrings for XmssKeyManager …

3f12f66

…methods

style: improve docstring clarity for XmssKeyManager methods

85a0b83

		Handles padding to the fixed 3100-byte length required by the consensus layer,
		delegating all encoding details to the XMSS container itself.

		from lean_spec.subspecs.xmss.interface import TEST_SIGNATURE_SCHEME
		from lean_spec.types import Uint64

		default_activation_epoch : Uint64, optional
		Activation epoch used when none is provided for key generation.

framework: add configurable seeds/activation to XmssKeyManager (#129) #173

Are you sure you want to change the base?

framework: add configurable seeds/activation to XmssKeyManager (#129) #173

Uh oh!

Conversation

vyakart commented Nov 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issues or PRs

Checklist

Uh oh!

tcoratger left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vyakart commented Nov 19, 2025

Changes Made

Verification

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vyakart commented Nov 19, 2025

Uh oh!

tcoratger left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vyakart commented Nov 20, 2025

Uh oh!

vyakart commented Nov 21, 2025

Updates after review

Uh oh!

vyakart commented Nov 17, 2025 •

edited

Loading