From 36b077c1a678e6b04f33351cd8e6c608feea2edd Mon Sep 17 00:00:00 2001 From: Thomas Coratger <60488569+tcoratger@users.noreply.github.com> Date: Sat, 6 Jun 2026 08:41:20 +0200 Subject: [PATCH] refactor(testing): split the key-generation CLI out of the key manager The key module mixed two unrelated concerns: the runtime signing manager that every fixture and unit test imports, and a standalone generate/download command-line tool that only runs from CI or a developer shell. The tool section also carried mid-function imports of sys and time, violating the top-of-file import rule. Changes: - the new keys_cli module holds the download URLs, the CLI-only defaults (validator count, generation max slot), the parallel key-generation workers, the archive download, and the entry point; its imports all live at the top of the file - the key module keeps the signing manager, the dummy signature, the scheme map, and the key-directory resolver; its module docstring points at the companion CLI - the missing-keys error hints and the CI workflows invoke the new module path; the fill CLI imports the download from it Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/generate-keys.yml | 4 +- .github/workflows/prod-vectors.yml | 2 +- .../testing/src/consensus_testing/keys.py | 239 +----------------- .../testing/src/consensus_testing/keys_cli.py | 239 ++++++++++++++++++ packages/testing/src/framework/cli/fill.py | 3 +- 5 files changed, 250 insertions(+), 237 deletions(-) create mode 100644 packages/testing/src/consensus_testing/keys_cli.py diff --git a/.github/workflows/generate-keys.yml b/.github/workflows/generate-keys.yml index 71848ce4b..7a95f3b61 100644 --- a/.github/workflows/generate-keys.yml +++ b/.github/workflows/generate-keys.yml @@ -38,10 +38,10 @@ jobs: run: uv sync --no-progress - name: Generate test scheme keys - run: uv run python -m consensus_testing.keys --scheme test + run: uv run python -m consensus_testing.keys_cli --scheme test - name: Generate prod scheme keys - run: LEAN_ENV=prod uv run python -m consensus_testing.keys --scheme prod + run: LEAN_ENV=prod uv run python -m consensus_testing.keys_cli --scheme prod - name: Publish to leansig-test-keys run: | diff --git a/.github/workflows/prod-vectors.yml b/.github/workflows/prod-vectors.yml index 14986a55c..6a84a528b 100644 --- a/.github/workflows/prod-vectors.yml +++ b/.github/workflows/prod-vectors.yml @@ -83,7 +83,7 @@ jobs: - name: Download keys if: steps.key-cache.outputs.cache-hit != 'true' - run: uv run python -m consensus_testing.keys --download --scheme prod + run: uv run python -m consensus_testing.keys_cli --download --scheme prod - name: Save key cache if: steps.key-cache.outputs.cache-hit != 'true' diff --git a/packages/testing/src/consensus_testing/keys.py b/packages/testing/src/consensus_testing/keys.py index f0e5fcae1..50f1a35ea 100755 --- a/packages/testing/src/consensus_testing/keys.py +++ b/packages/testing/src/consensus_testing/keys.py @@ -3,18 +3,9 @@ Management of XMSS key pairs for test validators. -Keys are pre-generated and cached on disk to avoid expensive generation during tests. - -Downloading Pre-generated Keys: - - python -m consensus_testing.keys --download --scheme test # test scheme - python -m consensus_testing.keys --download --scheme prod # prod scheme - -Regenerating Keys: - - python -m consensus_testing.keys # defaults - python -m consensus_testing.keys --count 20 # more validators - python -m consensus_testing.keys --max-slot 200 # longer lifetime +Keys are pre-generated and cached on disk to avoid expensive generation +during tests. +The companion CLI module generates or downloads them. File format: @@ -27,16 +18,8 @@ from __future__ import annotations -import argparse import json -import os -import shutil -import tarfile -import tempfile -import urllib.request from collections.abc import Iterator, Mapping -from concurrent.futures import ProcessPoolExecutor -from functools import partial from pathlib import Path from typing import ClassVar, Literal @@ -67,22 +50,11 @@ AttestationData, SingleMessageAggregate, ) -from lean_spec.spec.ssz import Bytes32, Uint64 +from lean_spec.spec.ssz import Bytes32 KeyRole = Literal["attestation", "proposal"] """Discriminator for which signing role's key to load from a validator key pair.""" -KEY_DOWNLOAD_URLS = { - "test": "https://github.com/leanEthereum/leansig-test-keys/releases/download/latest/test_scheme.tar.gz", - "prod": "https://github.com/leanEthereum/leansig-test-keys/releases/download/latest/prod_scheme.tar.gz", -} -""" -GitHub release URLs for pre-generated key archives. - -Keyed by scheme name ("test" or "prod"). -Each URL points to a tar.gz containing per-validator JSON files. -""" - LEAN_ENV_TO_SCHEMES: dict[str, GeneralizedXmssScheme] = { "test": TEST_SIGNATURE_SCHEME, "prod": PROD_SIGNATURE_SCHEME, @@ -135,21 +107,6 @@ def create_dummy_signature() -> Signature: Slot 10 is high enough for most unit tests while keeping key generation fast. """ -NUM_VALIDATORS: int = 8 -""" -Default number of validator key pairs. - -Eight validators is enough to exercise committee logic and 2/3 supermajority -thresholds while keeping key generation and test execution fast. -""" - -CLI_DEFAULT_MAX_SLOT = Slot(100) -""" -Maximum slot when generating keys via CLI (inclusive). - -One hundred slots provides ample signing headroom for typical test scenarios. -""" - def get_keys_directory(scheme_name: str) -> Path: """ @@ -284,7 +241,7 @@ def _scan_indices(self) -> set[ValidatorIndex]: if not self._keys_directory.exists(): raise FileNotFoundError( f"Keys directory not found: {self._keys_directory} - " - f"Run: python -m consensus_testing.keys --scheme {self.scheme_name}" + f"Run: python -m consensus_testing.keys_cli --scheme {self.scheme_name}" ) # Each JSON file is named by its validator index (e.g. "0.json"). @@ -297,7 +254,7 @@ def _scan_indices(self) -> set[ValidatorIndex]: if not self._available_indices: raise FileNotFoundError( f"No key files found in: {self._keys_directory} - " - f"Run: python -m consensus_testing.keys --scheme {self.scheme_name}" + f"Run: python -m consensus_testing.keys_cli --scheme {self.scheme_name}" ) return self._available_indices @@ -593,187 +550,3 @@ def build_attestation_proofs( ) for aggregate in aggregated_attestations ] - - -def _generate_single_keypair( - scheme: GeneralizedXmssScheme, num_slots: int, index: int -) -> ValidatorKeyPair: - """ - Generate attestation and proposal key pairs for one validator. - - Defined at module level so it can be pickled for multiprocessing. - - Args: - scheme: XMSS scheme instance to use for key generation. - num_slots: Total number of slots the keys must cover. - index: Validator index (used only for progress logging). - - Returns: - Complete key pair with both attestation and proposal keys. - """ - import sys - import time - - # Generate two independent key pairs: one for attestations, one for proposals. - # - # Separate keys allow signing both roles within the same slot - # without exhausting a one-time leaf. - start = time.monotonic() - print(f"[key #{index}] generating attestation key...", flush=True) - attestation_keypair = scheme.key_gen(Slot(0), Uint64(num_slots)) - - elapsed = time.monotonic() - start - print( - f"[key #{index}] attestation key done ({elapsed:.0f}s), generating proposal key...", - flush=True, - ) - proposal_keypair = scheme.key_gen(Slot(0), Uint64(num_slots)) - - elapsed = time.monotonic() - start - print(f"[key #{index}] done ({elapsed:.0f}s)", file=sys.stderr, flush=True) - - return ValidatorKeyPair( - attestation_keypair=attestation_keypair, proposal_keypair=proposal_keypair - ) - - -def _generate_keys(lean_env: str, count: int, max_slot: int) -> None: - """ - Generate XMSS key pairs in parallel and write each to a separate file. - - Each validator gets its own JSON file to keep individual files small, - which matters especially for production-scheme keys. - - Args: - lean_env: Scheme name (e.g. "test" or "prod"). - count: Number of validator key pairs to generate. - max_slot: Maximum signable slot (key lifetime = max_slot + 1 slots). - """ - scheme = LEAN_ENV_TO_SCHEMES[lean_env] - keys_directory = get_keys_directory(lean_env) - num_slots = max_slot + 1 - num_workers = os.cpu_count() or 1 - - print( - f"Generating {count} XMSS key pairs for {lean_env} environment " - f"({num_slots} slots) using {num_workers} cores..." - ) - - # Ensure the output directory exists. - keys_directory.mkdir(parents=True, exist_ok=True) - - # Remove stale key files from previous runs that may have generated - # a different number of keys. - for old_file in keys_directory.glob("*.json"): - old_file.unlink() - - # Generate key pairs in parallel across all CPU cores. - # Results arrive in index order thanks to executor.map. - import time - - gen_start = time.monotonic() - with ProcessPoolExecutor(max_workers=num_workers) as executor: - worker_func = partial(_generate_single_keypair, scheme, num_slots) - for validator_index, key_pair in enumerate(executor.map(worker_func, range(count))): - elapsed = time.monotonic() - gen_start - print( - f"[{validator_index + 1}/{count}] saved key #{validator_index} " - f"({elapsed:.0f}s elapsed)" - ) - - key_file = keys_directory / f"{validator_index}.json" - key_file.write_text(key_pair.model_dump_json(indent=2)) - - total = time.monotonic() - gen_start - print(f"Saved {count} key pairs to {keys_directory}/ ({total:.0f}s total)") - - -def download_keys(scheme: str) -> None: - """ - Download pre-generated key pairs from a GitHub release. - - Downloads a tar.gz archive for the specified scheme, removes any - existing keys for that scheme, and extracts the archive in place. - - Args: - scheme: Scheme name ("test" or "prod"). - """ - base_directory = Path(__file__).parent / "test_keys" - url = KEY_DOWNLOAD_URLS[scheme] - - print(f"Downloading {scheme} keys from {url}...") - - # Reserve a temp path; we open it explicitly below so the writer can close - # before the reader opens. - temporary_fd, temporary_name = tempfile.mkstemp(suffix=".tar.gz") - os.close(temporary_fd) - tmp_path = Path(temporary_name) - - try: - # Close the writer before opening the reader. - # Otherwise Python's userspace buffer can withhold the tail of the gzip stream. - # That produces a near-end decompression failure that looks like a truncated download. - with urllib.request.urlopen(url) as response, tmp_path.open("wb") as out: - shutil.copyfileobj(response, out) - - # Remove any existing keys for this scheme before extracting. - target_directory = base_directory / f"{scheme}_scheme" - if target_directory.exists(): - shutil.rmtree(target_directory) - base_directory.mkdir(parents=True, exist_ok=True) - - # Extract the archive into the base directory. - # The archive root is the scheme directory itself. - with tarfile.open(tmp_path, "r:gz") as tar: - tar.extractall(path=base_directory, filter="data") - - print(f"Extracted {scheme} keys to {target_directory}/") - finally: - # Always clean up the temporary download file. - tmp_path.unlink(missing_ok=True) - - print("Download complete!") - - -def main() -> None: - """CLI entry point for generating or downloading test keys.""" - parser = argparse.ArgumentParser( - description="Generate XMSS key pairs for consensus testing", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - parser.add_argument( - "--download", - action="store_true", - help="Download pre-generated keys from a GitHub release", - ) - parser.add_argument( - "--scheme", - choices=LEAN_ENV_TO_SCHEMES.keys(), - default="test", - help="XMSS scheme to use", - ) - parser.add_argument( - "--count", - type=int, - default=NUM_VALIDATORS, - help="Number of validator key pairs", - ) - parser.add_argument( - "--max-slot", - type=int, - default=int(CLI_DEFAULT_MAX_SLOT), - help="Maximum slot (key lifetime = max_slot + 1)", - ) - args = parser.parse_args() - - # Download pre-generated keys instead of generating locally. - if args.download: - download_keys(scheme=args.scheme) - return - - # Generate fresh keys with the specified parameters. - _generate_keys(lean_env=args.scheme, count=args.count, max_slot=args.max_slot) - - -if __name__ == "__main__": - main() diff --git a/packages/testing/src/consensus_testing/keys_cli.py b/packages/testing/src/consensus_testing/keys_cli.py new file mode 100644 index 000000000..6c12debbc --- /dev/null +++ b/packages/testing/src/consensus_testing/keys_cli.py @@ -0,0 +1,239 @@ +""" +CLI for generating or downloading the pre-generated XMSS test keys. + +Downloading Pre-generated Keys: + + python -m consensus_testing.keys_cli --download --scheme test # test scheme + python -m consensus_testing.keys_cli --download --scheme prod # prod scheme + +Regenerating Keys: + + python -m consensus_testing.keys_cli # defaults + python -m consensus_testing.keys_cli --count 20 # more validators + python -m consensus_testing.keys_cli --max-slot 200 # longer lifetime +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import sys +import tarfile +import tempfile +import time +import urllib.request +from concurrent.futures import ProcessPoolExecutor +from functools import partial +from pathlib import Path + +from consensus_testing.keys import LEAN_ENV_TO_SCHEMES, get_keys_directory +from lean_spec.spec.crypto.xmss.containers import ValidatorKeyPair +from lean_spec.spec.crypto.xmss.interface import GeneralizedXmssScheme +from lean_spec.spec.forks import Slot +from lean_spec.spec.ssz import Uint64 + +KEY_DOWNLOAD_URLS = { + "test": "https://github.com/leanEthereum/leansig-test-keys/releases/download/latest/test_scheme.tar.gz", + "prod": "https://github.com/leanEthereum/leansig-test-keys/releases/download/latest/prod_scheme.tar.gz", +} +""" +GitHub release URLs for pre-generated key archives. + +Keyed by scheme name ("test" or "prod"). +Each URL points to a tar.gz containing per-validator JSON files. +""" + +NUM_VALIDATORS: int = 8 +""" +Default number of validator key pairs. + +Eight validators is enough to exercise committee logic and 2/3 supermajority +thresholds while keeping key generation and test execution fast. +""" + +CLI_DEFAULT_MAX_SLOT = Slot(100) +""" +Maximum slot when generating keys via CLI (inclusive). + +One hundred slots provides ample signing headroom for typical test scenarios. +""" + + +def _generate_single_keypair( + scheme: GeneralizedXmssScheme, num_slots: int, index: int +) -> ValidatorKeyPair: + """ + Generate attestation and proposal key pairs for one validator. + + Defined at module level so it can be pickled for multiprocessing. + + Args: + scheme: XMSS scheme instance to use for key generation. + num_slots: Total number of slots the keys must cover. + index: Validator index (used only for progress logging). + + Returns: + Complete key pair with both attestation and proposal keys. + """ + # Generate two independent key pairs: one for attestations, one for proposals. + # + # Separate keys allow signing both roles within the same slot + # without exhausting a one-time leaf. + start = time.monotonic() + print(f"[key #{index}] generating attestation key...", flush=True) + attestation_keypair = scheme.key_gen(Slot(0), Uint64(num_slots)) + + elapsed = time.monotonic() - start + print( + f"[key #{index}] attestation key done ({elapsed:.0f}s), generating proposal key...", + flush=True, + ) + proposal_keypair = scheme.key_gen(Slot(0), Uint64(num_slots)) + + elapsed = time.monotonic() - start + print(f"[key #{index}] done ({elapsed:.0f}s)", file=sys.stderr, flush=True) + + return ValidatorKeyPair( + attestation_keypair=attestation_keypair, proposal_keypair=proposal_keypair + ) + + +def _generate_keys(lean_env: str, count: int, max_slot: int) -> None: + """ + Generate XMSS key pairs in parallel and write each to a separate file. + + Each validator gets its own JSON file to keep individual files small, + which matters especially for production-scheme keys. + + Args: + lean_env: Scheme name (e.g. "test" or "prod"). + count: Number of validator key pairs to generate. + max_slot: Maximum signable slot (key lifetime = max_slot + 1 slots). + """ + scheme = LEAN_ENV_TO_SCHEMES[lean_env] + keys_directory = get_keys_directory(lean_env) + num_slots = max_slot + 1 + num_workers = os.cpu_count() or 1 + + print( + f"Generating {count} XMSS key pairs for {lean_env} environment " + f"({num_slots} slots) using {num_workers} cores..." + ) + + # Ensure the output directory exists. + keys_directory.mkdir(parents=True, exist_ok=True) + + # Remove stale key files from previous runs that may have generated + # a different number of keys. + for old_file in keys_directory.glob("*.json"): + old_file.unlink() + + # Generate key pairs in parallel across all CPU cores. + # Results arrive in index order thanks to executor.map. + gen_start = time.monotonic() + with ProcessPoolExecutor(max_workers=num_workers) as executor: + worker_func = partial(_generate_single_keypair, scheme, num_slots) + for validator_index, key_pair in enumerate(executor.map(worker_func, range(count))): + elapsed = time.monotonic() - gen_start + print( + f"[{validator_index + 1}/{count}] saved key #{validator_index} " + f"({elapsed:.0f}s elapsed)" + ) + + key_file = keys_directory / f"{validator_index}.json" + key_file.write_text(key_pair.model_dump_json(indent=2)) + + total = time.monotonic() - gen_start + print(f"Saved {count} key pairs to {keys_directory}/ ({total:.0f}s total)") + + +def download_keys(scheme: str) -> None: + """ + Download pre-generated key pairs from a GitHub release. + + Downloads a tar.gz archive for the specified scheme, removes any + existing keys for that scheme, and extracts the archive in place. + + Args: + scheme: Scheme name ("test" or "prod"). + """ + base_directory = get_keys_directory(scheme).parent + url = KEY_DOWNLOAD_URLS[scheme] + + print(f"Downloading {scheme} keys from {url}...") + + # Reserve a temp path; we open it explicitly below so the writer can close + # before the reader opens. + temporary_fd, temporary_name = tempfile.mkstemp(suffix=".tar.gz") + os.close(temporary_fd) + tmp_path = Path(temporary_name) + + try: + # Close the writer before opening the reader. + # Otherwise Python's userspace buffer can withhold the tail of the gzip stream. + # That produces a near-end decompression failure that looks like a truncated download. + with urllib.request.urlopen(url) as response, tmp_path.open("wb") as out: + shutil.copyfileobj(response, out) + + # Remove any existing keys for this scheme before extracting. + target_directory = base_directory / f"{scheme}_scheme" + if target_directory.exists(): + shutil.rmtree(target_directory) + base_directory.mkdir(parents=True, exist_ok=True) + + # Extract the archive into the base directory. + # The archive root is the scheme directory itself. + with tarfile.open(tmp_path, "r:gz") as tar: + tar.extractall(path=base_directory, filter="data") + + print(f"Extracted {scheme} keys to {target_directory}/") + finally: + # Always clean up the temporary download file. + tmp_path.unlink(missing_ok=True) + + print("Download complete!") + + +def main() -> None: + """CLI entry point for generating or downloading test keys.""" + parser = argparse.ArgumentParser( + description="Generate XMSS key pairs for consensus testing", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--download", + action="store_true", + help="Download pre-generated keys from a GitHub release", + ) + parser.add_argument( + "--scheme", + choices=LEAN_ENV_TO_SCHEMES.keys(), + default="test", + help="XMSS scheme to use", + ) + parser.add_argument( + "--count", + type=int, + default=NUM_VALIDATORS, + help="Number of validator key pairs", + ) + parser.add_argument( + "--max-slot", + type=int, + default=int(CLI_DEFAULT_MAX_SLOT), + help="Maximum slot (key lifetime = max_slot + 1)", + ) + args = parser.parse_args() + + # Download pre-generated keys instead of generating locally. + if args.download: + download_keys(scheme=args.scheme) + return + + # Generate fresh keys with the specified parameters. + _generate_keys(lean_env=args.scheme, count=args.count, max_slot=args.max_slot) + + +if __name__ == "__main__": + main() diff --git a/packages/testing/src/framework/cli/fill.py b/packages/testing/src/framework/cli/fill.py index 96a51328b..84b30cc05 100644 --- a/packages/testing/src/framework/cli/fill.py +++ b/packages/testing/src/framework/cli/fill.py @@ -64,7 +64,8 @@ def fill( # Check and download keys if needed # Import here to avoid loading leanSpec modules before LEAN_ENV is set - from consensus_testing.keys import download_keys, get_keys_directory + from consensus_testing.keys import get_keys_directory + from consensus_testing.keys_cli import download_keys keys_directory = get_keys_directory(scheme.lower())