pre-generated test keys for faster CI

.github/workflows/ci.yml

@@ -18,6 +18,8 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
+      with:
+        submodules: recursive
 
     - name: Set up Zig
       uses: mlugg/[email protected]
@@ -82,6 +84,8 @@ jobs:
         os: [ubuntu-latest, macos-latest]
     steps:
     - uses: actions/checkout@v4
+      with:
+        submodules: recursive
 
     - name: Free disk space (Ubuntu)
       if: runner.os == 'Linux'
@@ -142,6 +146,8 @@ jobs:
         os: [ubuntu-latest, macos-latest]
     steps:
     - uses: actions/checkout@v4
+      with:
+        submodules: recursive
 
     - name: Free disk space (Ubuntu)
       if: runner.os == 'Linux'
@@ -282,6 +288,8 @@ jobs:
         os: [ubuntu-latest, macos-latest]
     steps:
     - uses: actions/checkout@v4
+      with:
+        submodules: recursive
 
     - name: Free disk space (Ubuntu)
       if: runner.os == 'Linux'
@@ -327,6 +335,7 @@ jobs:
     steps:
     - uses: actions/checkout@v4
       with:
+        submodules: recursive
         fetch-depth: 0  # Fetch full history to get git commit info
 
     - name: Free disk space (Ubuntu)

.gitmodules

@@ -4,3 +4,6 @@
 [submodule "leanSpec"]
 	path = leanSpec
 	url = https://github.com/leanEthereum/leanSpec
+[submodule "test-keys"]
+	path = test-keys
+	url = https://github.com/blockblaz/zeam-test-keys.git

build.zig

@@ -156,6 +156,8 @@ pub fn build(b: *Builder) !void {
     build_options.addOption(bool, "has_openvm", prover == .openvm or prover == .all);
     const use_poseidon = b.option(bool, "use_poseidon", "Use Poseidon SSZ hasher (default: false)") orelse false;
     build_options.addOption(bool, "use_poseidon", use_poseidon);
+    // Absolute path to test-keys for pre-generated validator keys
+    build_options.addOption([]const u8, "test_keys_path", b.pathFromRoot("test-keys/hash-sig-keys"));
     const build_options_module = build_options.createModule();
 
     // add zeam-utils
@@ -240,6 +242,7 @@ pub fn build(b: *Builder) !void {
         .target = target,
         .optimize = optimize,
     });
+    zeam_key_manager.addImport("build_options", build_options_module);
     zeam_key_manager.addImport("@zeam/xmss", zeam_xmss);
     zeam_key_manager.addImport("@zeam/types", zeam_types);
     zeam_key_manager.addImport("@zeam/utils", zeam_utils);
@@ -421,6 +424,10 @@ pub fn build(b: *Builder) !void {
     tools_cli_exe.root_module.addImport("enr", enr);
     tools_cli_exe.root_module.addImport("build_options", build_options_module);
     tools_cli_exe.root_module.addImport("simargs", simargs);
+    tools_cli_exe.root_module.addImport("@zeam/xmss", zeam_xmss);
+    tools_cli_exe.root_module.addImport("@zeam/types", zeam_types);
+    tools_cli_exe.step.dependOn(&build_rust_lib_steps.step);
+    addRustGlueLib(b, tools_cli_exe, target, prover);
 
     const install_tools_cli = b.addInstallArtifact(tools_cli_exe, .{});
     tools_step.dependOn(&install_tools_cli.step);
@@ -602,6 +609,10 @@ pub fn build(b: *Builder) !void {
         .root_module = tools_cli_exe.root_module,
     });
     tools_cli_tests.root_module.addImport("enr", enr);
+    tools_cli_tests.root_module.addImport("@zeam/xmss", zeam_xmss);
+    tools_cli_tests.root_module.addImport("@zeam/types", zeam_types);
+    tools_cli_tests.step.dependOn(&build_rust_lib_steps.step);
+    addRustGlueLib(b, tools_cli_tests, target, prover);
     const run_tools_cli_test = b.addRunArtifact(tools_cli_tests);
     setTestRunLabelFromCompile(b, run_tools_cli_test, tools_cli_tests);
     tools_test_step.dependOn(&run_tools_cli_test.step);

pkgs/key-manager/src/lib.zig

@@ -55,30 +55,39 @@ fn getOrCreateCachedKeyPair(
 pub const KeyManager = struct {
     keys: std.AutoHashMap(usize, xmss.KeyPair),
     allocator: Allocator,
-    owns_keypairs: bool,
+    /// Tracks which keypairs are owned (allocated by us) vs borrowed (cached).
+    owned_keys: std.AutoHashMap(usize, void),
 
     const Self = @This();
 
     pub fn init(allocator: Allocator) Self {
         return Self{
             .keys = std.AutoHashMap(usize, xmss.KeyPair).init(allocator),
             .allocator = allocator,
-            .owns_keypairs = true,
+            .owned_keys = std.AutoHashMap(usize, void).init(allocator),
         };
     }
 
     pub fn deinit(self: *Self) void {
-        if (self.owns_keypairs) {
-            var it = self.keys.iterator();
-            while (it.next()) |entry| {
+        var it = self.keys.iterator();
+        while (it.next()) |entry| {
+            if (self.owned_keys.contains(entry.key_ptr.*)) {
                 entry.value_ptr.deinit();
             }
         }
         self.keys.deinit();
+        self.owned_keys.deinit();
     }
 
+    /// Add an owned keypair that will be freed on deinit.
     pub fn addKeypair(self: *Self, validator_id: usize, keypair: xmss.KeyPair) !void {
         try self.keys.put(validator_id, keypair);
+        try self.owned_keys.put(validator_id, {});
+    }
+
+    /// Add a cached/borrowed keypair that will NOT be freed on deinit.
+    pub fn addCachedKeypair(self: *Self, validator_id: usize, keypair: xmss.KeyPair) !void {
+        try self.keys.put(validator_id, keypair);
     }
 
     pub fn loadFromKeypairDir(_: *Self, _: []const u8) !void {
@@ -162,24 +171,105 @@ pub const KeyManager = struct {
     }
 };
 
+/// Maximum size of a serialized XMSS private key (20MB).
+const MAX_SK_SIZE = 1024 * 1024 * 20;
+
+/// Maximum size of a serialized XMSS public key (256 bytes).
+const MAX_PK_SIZE = 256;
+
+/// Number of pre-generated test keys available in the test-keys submodule.
+const NUM_PREGENERATED_KEYS: usize = 32;
+
+const build_options = @import("build_options");
+
+/// Find the test-keys directory using the repo root path injected by build.zig.
+fn findTestKeysDir() ?[]const u8 {
+    const keys_path = build_options.test_keys_path;
+    if (keys_path.len == 0) return null;
+
+    // Verify it actually exists at runtime
+    if (std.fs.cwd().openDir(keys_path, .{})) |dir| {
+        var d = dir;
+        d.close();
+        return keys_path;
+    } else |_| {}
+
+    return null;
+}
+
+/// Load a single pre-generated key pair from SSZ files on disk.
+fn loadPreGeneratedKey(
+    allocator: Allocator,
+    keys_dir: []const u8,
+    index: usize,
+) !xmss.KeyPair {
+    // Build file paths
+    var sk_path_buf: [512]u8 = undefined;
+    const sk_path = std.fmt.bufPrint(&sk_path_buf, "{s}/validator_{d}_sk.ssz", .{ keys_dir, index }) catch unreachable;
+
+    var pk_path_buf: [512]u8 = undefined;
+    const pk_path = std.fmt.bufPrint(&pk_path_buf, "{s}/validator_{d}_pk.ssz", .{ keys_dir, index }) catch unreachable;
+
+    // Read private key
+    var sk_file = try std.fs.cwd().openFile(sk_path, .{});
+    defer sk_file.close();
+    const sk_data = try sk_file.readToEndAlloc(allocator, MAX_SK_SIZE);
+    defer allocator.free(sk_data);
+
+    // Read public key
+    var pk_file = try std.fs.cwd().openFile(pk_path, .{});
+    defer pk_file.close();
+    const pk_data = try pk_file.readToEndAlloc(allocator, MAX_PK_SIZE);
+    defer allocator.free(pk_data);
+
+    // Reconstruct keypair from SSZ
+    return xmss.KeyPair.fromSsz(allocator, sk_data, pk_data);
+}
+
 pub fn getTestKeyManager(
     allocator: Allocator,
     num_validators: usize,
     max_slot: usize,
 ) !KeyManager {
     var key_manager = KeyManager.init(allocator);
-    key_manager.owns_keypairs = false;
     errdefer key_manager.deinit();
 
-    var num_active_epochs = max_slot + 1;
-    // to reuse cached keypairs, gen for 10 since most tests ask for < 10 max slot including
-    // building mock chain for tests. otherwise getOrCreateCachedKeyPair might cleanup previous
-    //  key generated for smaller life time
-    if (num_active_epochs < 10) num_active_epochs = 10;
+    // Determine how many keys we can load from pre-generated files
+    const keys_dir = findTestKeysDir();
+    const num_preloaded = if (keys_dir != null)
+        @min(num_validators, NUM_PREGENERATED_KEYS)
+    else
+        0;
+
+    // Load pre-generated keys (fast path: near-instant from SSZ files)
+    var actually_loaded: usize = 0;
+    if (keys_dir) |dir| {
+        for (0..num_preloaded) |i| {
+            const keypair = loadPreGeneratedKey(allocator, dir, i) catch |err| {
+                std.debug.print("Failed to load pre-generated key {d}: {}\n", .{ i, err });
+                break;
+            };
+            key_manager.addKeypair(i, keypair) catch |err| {
+                std.debug.print("Failed to add pre-generated key {d}: {}\n", .{ i, err });
+                break;
+            };
+            actually_loaded += 1;
+        }
+        std.debug.print("Loaded {d} pre-generated test keys from {s}\n", .{ actually_loaded, dir });
+    } else {
+        std.debug.print("Pre-generated keys not found, generating all keys at runtime\n", .{});
+    }
+
+    // Generate remaining keys at runtime (for validators beyond the loaded set)
+    if (num_validators > actually_loaded) {
+        var num_active_epochs = max_slot + 1;
+        if (num_active_epochs < 10) num_active_epochs = 10;
 
-    for (0..num_validators) |i| {
-        const keypair = try getOrCreateCachedKeyPair(i, num_active_epochs);
-        try key_manager.addKeypair(i, keypair);
+        for (actually_loaded..num_validators) |i| {
+            const keypair = try getOrCreateCachedKeyPair(i, num_active_epochs);
+            try key_manager.addCachedKeypair(i, keypair);
+        }
+        std.debug.print("Generated {d} additional keys at runtime\n", .{num_validators - actually_loaded});
     }
 
     return key_manager;