feat: support spec test using test sig scheme

build.zig

@@ -321,6 +321,7 @@ pub fn build(b: *Builder) !void {
     zeam_spectests.addImport("build_options", build_options_module);
     zeam_spectests.addImport("@zeam/state-transition", zeam_state_transition);
     zeam_spectests.addImport("@zeam/node", zeam_beam_node);
+    zeam_spectests.addImport("@zeam/xmss", zeam_xmss);
 
     // Add the cli executable
     const cli_exe = b.addExecutable(.{
@@ -572,6 +573,10 @@ pub fn build(b: *Builder) !void {
     spectests.root_module.addImport("@zeam/metrics", zeam_metrics);
     spectests.root_module.addImport("@zeam/state-transition", zeam_state_transition);
     spectests.root_module.addImport("ssz", ssz);
+    spectests.root_module.addImport("@zeam/xmss", zeam_xmss);
+
+    spectests.step.dependOn(&build_rust_lib_steps.step);
+    addRustGlueLib(b, spectests, target, prover);
 
     manager_tests.step.dependOn(&build_rust_lib_steps.step);
 

pkgs/key-manager/src/lib.zig

@@ -4,9 +4,105 @@ const types = @import("@zeam/types");
 const zeam_utils = @import("@zeam/utils");
 const zeam_metrics = @import("@zeam/metrics");
 const Allocator = std.mem.Allocator;
+const JsonValue = std.json.Value;
+
+pub const XmssTestConfig = struct {
+    scheme: xmss.HashSigScheme,
+    signature_ssz_len: usize,
+    allow_placeholder_aggregated_proof: bool,
+
+    pub fn fromLeanEnv(lean_env: ?[]const u8) XmssTestConfig {
+        const scheme = schemeFromLeanEnv(lean_env);
+        return .{
+            .scheme = scheme,
+            .signature_ssz_len = xmss.signatureSszLenForScheme(scheme),
+            .allow_placeholder_aggregated_proof = scheme == .@"test",
+        };
+    }
+};
+
+pub const FixtureKeyError = error{
+    DuplicateKeyIndex,
+    InvalidKeyFile,
+    InvalidKeyIndex,
+    InvalidPublicKey,
+    NoKeysFound,
+};
+
+fn schemeFromLeanEnv(lean_env: ?[]const u8) xmss.HashSigScheme {
+    const env = lean_env orelse return .prod;
+    if (std.ascii.eqlIgnoreCase(env, "test")) return .@"test";
+    return .prod;
+}
+
+fn parseKeyIndex(file_name: []const u8) !usize {
+    if (!std.mem.endsWith(u8, file_name, ".json")) {
+        return FixtureKeyError.InvalidKeyIndex;
+    }
+    const stem = file_name[0 .. file_name.len - ".json".len];
+    if (stem.len == 0) {
+        return FixtureKeyError.InvalidKeyIndex;
+    }
+    return std.fmt.parseInt(usize, stem, 10) catch FixtureKeyError.InvalidKeyIndex;
+}
+
+const fixture_key_file_max_bytes: usize = 2 * 1024 * 1024;
+
+fn readPublicKeyFromJson(
+    allocator: Allocator,
+    dir: std.fs.Dir,
+    file_name: []const u8,
+) !types.Bytes52 {
+    const payload = dir.readFileAlloc(allocator, file_name, fixture_key_file_max_bytes) catch {
+        return FixtureKeyError.InvalidKeyFile;
+    };
+    defer allocator.free(payload);
+
+    var parsed = std.json.parseFromSlice(JsonValue, allocator, payload, .{ .ignore_unknown_fields = true }) catch {
+        return FixtureKeyError.InvalidKeyFile;
+    };
+    defer parsed.deinit();
+
+    const obj = switch (parsed.value) {
+        .object => |map| map,
+        else => return FixtureKeyError.InvalidKeyFile,
+    };
+    const pub_val = obj.get("public") orelse return FixtureKeyError.InvalidKeyFile;
+    const pub_hex = switch (pub_val) {
+        .string => |s| s,
+        else => return FixtureKeyError.InvalidKeyFile,
+    };
+
+    return parsePublicKeyHex(pub_hex);
+}
+
+fn parsePublicKeyHex(input: []const u8) !types.Bytes52 {
+    const public_key_hex_len: usize = 2 * @sizeOf(types.Bytes52);
+    const hex_str = if (std.mem.startsWith(u8, input, "0x")) input[2..] else input;
+    if (hex_str.len != public_key_hex_len) {
+        return FixtureKeyError.InvalidPublicKey;
+    }
+    var bytes: types.Bytes52 = undefined;
+    _ = std.fmt.hexToBytes(&bytes, hex_str) catch {
+        return FixtureKeyError.InvalidPublicKey;
+    };
+    return bytes;
+}
 
 const KeyManagerError = error{
     ValidatorKeyNotFound,
+    PrivateKeyMissing,
+    PublicKeyBufferTooSmall,
+};
+
+const PublicKeyEntry = struct {
+    bytes: types.Bytes52,
+    public_key: xmss.PublicKey,
+};
+
+const KeyEntry = union(enum) {
+    keypair: xmss.KeyPair,
+    public_key: PublicKeyEntry,
 };
 
 const CachedKeyPair = struct {
@@ -53,32 +149,80 @@ fn getOrCreateCachedKeyPair(
 }
 
 pub const KeyManager = struct {
-    keys: std.AutoHashMap(usize, xmss.KeyPair),
+    keys: std.AutoHashMap(usize, KeyEntry),
     allocator: Allocator,
     owns_keypairs: bool,
 
     const Self = @This();
 
     pub fn init(allocator: Allocator) Self {
         return Self{
-            .keys = std.AutoHashMap(usize, xmss.KeyPair).init(allocator),
+            .keys = std.AutoHashMap(usize, KeyEntry).init(allocator),
             .allocator = allocator,
             .owns_keypairs = true,
         };
     }
 
     pub fn deinit(self: *Self) void {
-        if (self.owns_keypairs) {
-            var it = self.keys.iterator();
-            while (it.next()) |entry| {
-                entry.value_ptr.deinit();
-            }
-        }
+        self.clearEntries();
         self.keys.deinit();
     }
 
     pub fn addKeypair(self: *Self, validator_id: usize, keypair: xmss.KeyPair) !void {
-        try self.keys.put(validator_id, keypair);
+        if (self.keys.getPtr(validator_id)) |entry| {
+            self.deinitEntry(entry);
+            entry.* = .{ .keypair = keypair };
+            return;
+        }
+        try self.keys.put(validator_id, .{ .keypair = keypair });
+    }
+
+    pub fn addPublicKey(self: *Self, validator_id: usize, pubkey_bytes: types.Bytes52) !void {
+        var public_key = try xmss.PublicKey.fromBytes(pubkey_bytes[0..]);
+        errdefer public_key.deinit();
+
+        const entry_value = PublicKeyEntry{ .bytes = pubkey_bytes, .public_key = public_key };
+        if (self.keys.getPtr(validator_id)) |entry| {
+            self.deinitEntry(entry);
+            entry.* = .{ .public_key = entry_value };
+            return;
+        }
+        try self.keys.put(validator_id, .{ .public_key = entry_value });
+    }
+
+    pub fn loadLeanSpecKeys(self: *Self, keys_root: []const u8, lean_env: ?[]const u8) !void {
+        const scheme_dir_name = switch (schemeFromLeanEnv(lean_env)) {
+            .@"test" => "test_scheme",
+            .prod => "prod_scheme",
+        };
+        const scheme_dir_path = try std.fs.path.join(self.allocator, &.{ keys_root, scheme_dir_name });
+        defer self.allocator.free(scheme_dir_path);
+        try self.loadKeysFromDir(scheme_dir_path);
+    }
+
+    pub fn loadKeysFromDir(self: *Self, keys_dir_path: []const u8) !void {
+        var dir = try std.fs.cwd().openDir(keys_dir_path, .{ .iterate = true });
+        defer dir.close();
+
+        self.clearEntries();
+
+        var it = dir.iterate();
+        while (try it.next()) |entry| {
+            if (entry.kind != .file) continue;
+            const index = parseKeyIndex(entry.name) catch continue;
+            const pubkey = try readPublicKeyFromJson(self.allocator, dir, entry.name);
+            if (self.keys.get(index) != null) {
+                return FixtureKeyError.DuplicateKeyIndex;
+            }
+            self.addPublicKey(index, pubkey) catch |err| switch (err) {
+                error.OutOfMemory => return err,
+                else => return FixtureKeyError.InvalidPublicKey,
+            };
+        }
+
+        if (self.keys.count() == 0) {
+            return FixtureKeyError.NoKeysFound;
+        }
     }
 
     pub fn loadFromKeypairDir(_: *Self, _: []const u8) !void {
@@ -109,8 +253,17 @@ pub const KeyManager = struct {
         validator_index: usize,
         buffer: []u8,
     ) !usize {
-        const keypair = self.keys.get(validator_index) orelse return KeyManagerError.ValidatorKeyNotFound;
-        return try keypair.pubkeyToBytes(buffer);
+        const entry = self.keys.getPtr(validator_index) orelse return KeyManagerError.ValidatorKeyNotFound;
+        return switch (entry.*) {
+            .keypair => |keypair| try keypair.pubkeyToBytes(buffer),
+            .public_key => |pubkey| blk: {
+                if (buffer.len < pubkey.bytes.len) {
+                    return KeyManagerError.PublicKeyBufferTooSmall;
+                }
+                @memcpy(buffer[0..pubkey.bytes.len], pubkey.bytes[0..]);
+                break :blk pubkey.bytes.len;
+            },
+        };
     }
 
     /// Extract all validator public keys into an array
@@ -136,8 +289,11 @@ pub const KeyManager = struct {
         self: *const Self,
         validator_index: usize,
     ) !*const xmss.HashSigPublicKey {
-        const keypair = self.keys.get(validator_index) orelse return KeyManagerError.ValidatorKeyNotFound;
-        return keypair.public_key;
+        const entry = self.keys.getPtr(validator_index) orelse return KeyManagerError.ValidatorKeyNotFound;
+        return switch (entry.*) {
+            .keypair => |keypair| keypair.public_key,
+            .public_key => |pubkey| pubkey.public_key.handle,
+        };
     }
 
     /// Sign an attestation and return the raw signature handle (for aggregation)
@@ -148,7 +304,11 @@ pub const KeyManager = struct {
         allocator: Allocator,
     ) !xmss.Signature {
         const validator_index: usize = @intCast(attestation.validator_id);
-        const keypair = self.keys.get(validator_index) orelse return KeyManagerError.ValidatorKeyNotFound;
+        const entry = self.keys.getPtr(validator_index) orelse return KeyManagerError.ValidatorKeyNotFound;
+        const keypair = switch (entry.*) {
+            .keypair => |*kp| kp,
+            .public_key => return KeyManagerError.PrivateKeyMissing,
+        };
 
         const signing_timer = zeam_metrics.lean_pq_signature_attestation_signing_time_seconds.start();
         var message: [32]u8 = undefined;
@@ -160,6 +320,25 @@ pub const KeyManager = struct {
 
         return signature;
     }
+
+    fn clearEntries(self: *Self) void {
+        var it = self.keys.iterator();
+        while (it.next()) |entry| {
+            self.deinitEntry(entry.value_ptr);
+        }
+        self.keys.clearRetainingCapacity();
+    }
+
+    fn deinitEntry(self: *Self, entry: *KeyEntry) void {
+        switch (entry.*) {
+            .keypair => |*keypair| {
+                if (self.owns_keypairs) {
+                    keypair.deinit();
+                }
+            },
+            .public_key => |*pubkey| pubkey.public_key.deinit(),
+        }
+    }
 };
 
 pub fn getTestKeyManager(

pkgs/spectest/src/fixture_kind.zig

@@ -1,20 +1,23 @@
 pub const FixtureKind = enum {
     state_transition,
     fork_choice,
+    verify_signatures,
 
     pub fn runnerModule(self: FixtureKind) []const u8 {
         return switch (self) {
             .state_transition => "state_transition",
             .fork_choice => "fork_choice",
+            .verify_signatures => "verify_signatures",
         };
     }
 
     pub fn handlerSubdir(self: FixtureKind) []const u8 {
         return switch (self) {
             .state_transition => "state_transition",
             .fork_choice => "fc",
+            .verify_signatures => "verify_signatures",
         };
     }
 };
 
-pub const all = [_]FixtureKind{ .state_transition, .fork_choice };
+pub const all = [_]FixtureKind{ .state_transition, .fork_choice, .verify_signatures };

pkgs/spectest/src/json_expect.zig

@@ -245,6 +245,23 @@ pub fn expectArrayValue(
     };
 }
 
+pub fn expectArrayField(
+    comptime FixtureError: type,
+    obj: std.json.ObjectMap,
+    field_names: []const []const u8,
+    context: Context,
+    label: []const u8,
+) FixtureError!std.json.Array {
+    const value = getField(obj, field_names) orelse {
+        std.debug.print(
+            "fixture {s} case {s}{}: missing field {s}\n",
+            .{ context.fixture_label, context.case_name, context.formatStep(), label },
+        );
+        return FixtureError.InvalidFixture;
+    };
+    return expectArrayValue(FixtureError, value, context, label);
+}
+
 pub fn appendBytesDataField(
     comptime FixtureError: type,
     comptime T: type,

pkgs/spectest/src/runner/fork_choice_runner.zig

@@ -633,7 +633,10 @@ fn processBlockStep(
     }
     try types.sszClone(ctx.allocator, types.BeamState, parent_state_ptr.*, new_state_ptr);
 
-    state_transition.apply_transition(ctx.allocator, new_state_ptr, block, .{ .logger = ctx.fork_logger, .validateResult = false }) catch |err| {
+    state_transition.apply_transition(ctx.allocator, new_state_ptr, block, .{
+        .logger = ctx.fork_logger,
+        .validateResult = false,
+    }) catch |err| {
         std.debug.print(
             "fixture {s} case {s}{}: state transition failed {s}\n",
             .{ fixture_path, case_name, formatStep(step_index), @errorName(err) },

pkgs/spectest/src/runner/state_transition_runner.zig

@@ -275,7 +275,11 @@ fn runCase(
             }
         }
 
-        state_transition.apply_transition(allocator, &pre_state, block, .{ .logger = logger }) catch |err| {
+        const validate_result = expect_exception != null;
+        state_transition.apply_transition(allocator, &pre_state, block, .{
+            .logger = logger,
+            .validateResult = validate_result,
+        }) catch |err| {
             encountered_error = true;
             if (expect_exception == null) {
                 std.debug.print(