From 1af9e21006bb0d5191133490136270995653d147 Mon Sep 17 00:00:00 2001 From: zclawz Date: Tue, 26 May 2026 07:28:27 +0000 Subject: [PATCH 1/3] hashsig-glue: replace Box allocs with placement-init pattern MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rust no longer heap-allocates the outer KeyPair / Signature / PublicKey wrapper structs. Instead, each init function accepts a caller-supplied buffer and writes into it with std::ptr::write; the matching _deinit runs std::ptr::drop_in_place (Rust Drop) without freeing the buffer. Zig side allocates the buffers via C malloc (always malloc-aligned, safe for all Rust #[repr(C)] types) and frees them after calling _deinit. New Rust exports hashsig_sizeof_keypair / _signature / _public_key — allocation size hashsig_alignof_keypair / _signature / _public_key — required alignment hashsig_keypair_generate_into (was: hashsig_keypair_generate → *mut Box) hashsig_keypair_from_ssz_into (was: hashsig_keypair_from_ssz → *mut Box) hashsig_keypair_deinit (was: hashsig_keypair_free → Box::from_raw) hashsig_sign_into (was: hashsig_sign → *mut Box) hashsig_signature_deinit (was: hashsig_signature_free) hashsig_signature_from_ssz_into hashsig_public_key_from_ssz_into hashsig_public_key_deinit Removed Rust exports (were Box-allocating): hashsig_keypair_generate, hashsig_keypair_from_ssz, hashsig_keypair_free hashsig_sign, hashsig_signature_free, hashsig_signature_from_ssz hashsig_public_key_from_ssz, hashsig_public_key_free Note: the inner XMSS types (leansig XmssPublicKey / XmssSignature etc.) still carry their own heap-allocated data (Vec inside leansig); those are inherent to the leansig API and are freed correctly when _deinit runs Drop in-place. The intermediate Vec in to_bytes() serialisation helpers is similarly a leansig API constraint; those are short-lived temporaries that do not escape the FFI call. Addresses: https://github.com/blockblaz/zeam/pull/918#discussion_r3301879001 --- pkgs/xmss/src/hashsig.zig | 575 ++++++++++++++++++++--------------- rust/hashsig-glue/src/lib.rs | 320 ++++++++++--------- 2 files changed, 511 insertions(+), 384 deletions(-) diff --git a/pkgs/xmss/src/hashsig.zig b/pkgs/xmss/src/hashsig.zig index 86bde3baf..f7a25e8f2 100644 --- a/pkgs/xmss/src/hashsig.zig +++ b/pkgs/xmss/src/hashsig.zig @@ -3,100 +3,110 @@ const Allocator = std.mem.Allocator; pub const aggregate = @import("aggregation.zig"); -/// Opaque pointer to the Rust KeyPair struct +// Opaque Zig types that map to the corresponding Rust structs. +// Zig never looks inside these; sizes come from the hashsig_sizeof_* queries. pub const HashSigKeyPair = opaque {}; - -/// Opaque pointer to the Rust Signature struct pub const HashSigSignature = opaque {}; - -/// Opaque pointer to the Rust PublicKey struct pub const HashSigPublicKey = opaque {}; - -/// Opaque pointer to the Rust PrivateKey struct pub const HashSigPrivateKey = opaque {}; -/// Generate a new key pair -extern fn hashsig_keypair_generate( +// ─── Layout queries ──────────────────────────────────────────────────────────── +// Rust exports these so Zig can pre-allocate the right amount of space. +// Call once at startup (or comptime-cache them); values are stable for the +// lifetime of a given Rust build. + +extern fn hashsig_sizeof_keypair() callconv(.c) usize; +extern fn hashsig_alignof_keypair() callconv(.c) usize; +extern fn hashsig_sizeof_signature() callconv(.c) usize; +extern fn hashsig_alignof_signature() callconv(.c) usize; +extern fn hashsig_sizeof_public_key() callconv(.c) usize; +extern fn hashsig_alignof_public_key() callconv(.c) usize; + +// ─── Placement-init (no Rust Box allocation) ────────────────────────────────── +// Each `_into` function writes a fully-initialised struct into caller-supplied +// storage (allocated by Zig/C malloc, which always satisfies Rust's alignment). +// The matching `_deinit` runs Rust Drop in-place WITHOUT freeing the buffer; +// the caller owns the buffer and must free it afterwards. + +extern fn hashsig_keypair_generate_into( + out: *HashSigKeyPair, seed_phrase: [*:0]const u8, activation_epoch: usize, num_active_epochs: usize, -) callconv(.c) ?*HashSigKeyPair; +) callconv(.c) c_int; -/// Reconstruct a key pair from SSZ-encoded bytes -extern fn hashsig_keypair_from_ssz( +extern fn hashsig_keypair_from_ssz_into( + out: *HashSigKeyPair, private_key_ssz: [*]const u8, private_key_len: usize, public_key_ssz: [*]const u8, public_key_len: usize, -) callconv(.c) ?*HashSigKeyPair; - -/// Free a key pair -extern fn hashsig_keypair_free(keypair: ?*HashSigKeyPair) callconv(.c) void; +) callconv(.c) c_int; -/// Get pointer to public key from keypair (valid as long as keypair is alive) -extern fn hashsig_keypair_get_public_key(keypair: *const HashSigKeyPair) callconv(.c) ?*const HashSigPublicKey; - -/// Get pointer to private key from keypair (valid as long as keypair is alive) -extern fn hashsig_keypair_get_private_key(keypair: *const HashSigKeyPair) callconv(.c) ?*const HashSigPrivateKey; +/// Destroy the KeyPair in-place (runs Rust Drop). Does NOT free the buffer. +extern fn hashsig_keypair_deinit(kp: *HashSigKeyPair) callconv(.c) void; -/// Sign a message using private key directly -extern fn hashsig_sign( +extern fn hashsig_sign_into( + out: *HashSigSignature, private_key: *const HashSigPrivateKey, message_ptr: [*]const u8, epoch: u32, -) callconv(.c) ?*HashSigSignature; +) callconv(.c) c_int; -/// Verify a signature using public key directly -extern fn hashsig_verify( - public_key: *const HashSigPublicKey, - message_ptr: [*]const u8, - epoch: u32, - signature: *const HashSigSignature, -) callconv(.c) i32; +/// Destroy the Signature in-place. Does NOT free the buffer. +extern fn hashsig_signature_deinit(sig: *HashSigSignature) callconv(.c) void; + +extern fn hashsig_signature_from_ssz_into( + out: *HashSigSignature, + sig_bytes: [*]const u8, + sig_len: usize, +) callconv(.c) c_int; + +extern fn hashsig_public_key_from_ssz_into( + out: *HashSigPublicKey, + pubkey_bytes: [*]const u8, + pubkey_len: usize, +) callconv(.c) c_int; + +/// Destroy the PublicKey in-place. Does NOT free the buffer. +extern fn hashsig_public_key_deinit(pk: *HashSigPublicKey) callconv(.c) void; + +// ─── Accessor / sub-key views ───────────────────────────────────────────────── +// These return pointers INTO the caller-owned KeyPair buffer. Valid for the +// lifetime of that buffer. + +extern fn hashsig_keypair_get_public_key(keypair: *const HashSigKeyPair) callconv(.c) ?*const HashSigPublicKey; +extern fn hashsig_keypair_get_private_key(keypair: *const HashSigKeyPair) callconv(.c) ?*const HashSigPrivateKey; + +// ─── Serialisation (caller-supplies-output, unchanged) ──────────────────────── -/// Serialize a public key pointer to bytes extern fn hashsig_public_key_to_bytes( public_key: *const HashSigPublicKey, buffer: [*]u8, buffer_len: usize, ) callconv(.c) usize; -/// Serialize a private key pointer to bytes extern fn hashsig_private_key_to_bytes( private_key: *const HashSigPrivateKey, buffer: [*]u8, buffer_len: usize, ) callconv(.c) usize; -/// Free a signature -extern fn hashsig_signature_free(signature: ?*HashSigSignature) callconv(.c) void; - -/// Construct a signature from SSZ bytes -extern fn hashsig_signature_from_ssz( - sig_bytes: [*]const u8, - sig_len: usize, -) callconv(.c) ?*HashSigSignature; - -/// Construct a public key from SSZ bytes -extern fn hashsig_public_key_from_ssz( - pubkey_bytes: [*]const u8, - pubkey_len: usize, -) callconv(.c) ?*HashSigPublicKey; - -/// Free a standalone public key -extern fn hashsig_public_key_free(pubkey: ?*HashSigPublicKey) callconv(.c) void; - -/// Get the message length constant -extern fn hashsig_message_length() callconv(.c) usize; - -/// Serialize a signature to bytes using SSZ encoding extern fn hashsig_signature_to_bytes( signature: *const HashSigSignature, buffer: [*]u8, buffer_len: usize, ) callconv(.c) usize; -/// Verify XMSS signature from SSZ-encoded bytes +// ─── Verify ─────────────────────────────────────────────────────────────────── + +extern fn hashsig_verify( + public_key: *const HashSigPublicKey, + message_ptr: [*]const u8, + epoch: u32, + signature: *const HashSigSignature, +) callconv(.c) i32; + extern fn hashsig_verify_ssz( pubkey_bytes: [*]const u8, pubkey_len: usize, @@ -106,9 +116,6 @@ extern fn hashsig_verify_ssz( signature_len: usize, ) callconv(.c) i32; -/// Verify XMSS signature against the leanSpec test scheme (LOG_LIFETIME=8, -/// DIMENSION=4). Used by spec-test fixtures whose `leanEnv=test` produces -/// ~424-byte signatures that the production scheme cannot parse. extern fn hashsig_test_verify_ssz( pubkey_bytes: [*]const u8, pubkey_len: usize, @@ -118,29 +125,54 @@ extern fn hashsig_test_verify_ssz( signature_len: usize, ) callconv(.c) i32; -pub const HashSigError = error{ KeyGenerationFailed, SigningFailed, VerificationFailed, InvalidSignature, SerializationFailed, InvalidMessageLength, DeserializationFailed, OutOfMemory, ValidatorIndexOutOfRange }; +extern fn hashsig_message_length() callconv(.c) usize; + +// ─── Error set ──────────────────────────────────────────────────────────────── + +pub const HashSigError = error{ + KeyGenerationFailed, + SigningFailed, + VerificationFailed, + InvalidSignature, + SerializationFailed, + InvalidMessageLength, + DeserializationFailed, + OutOfMemory, + ValidatorIndexOutOfRange, +}; + +// ─── Helpers ────────────────────────────────────────────────────────────────── + +/// Allocate `sz` bytes from the C heap (malloc). C malloc always returns +/// memory aligned to at least max_align_t (≥8 bytes on LP64), which satisfies +/// the alignment of any Rust `#[repr(C)]` type used in hashsig-glue. +fn cAlloc(sz: usize) HashSigError![*]u8 { + const raw = std.c.malloc(sz) orelse return HashSigError.OutOfMemory; + return @ptrCast(raw); +} + +fn cFree(ptr: [*]u8) void { + std.c.free(@ptrCast(ptr)); +} + +// ─── Byte-level verify helpers ──────────────────────────────────────────────── -/// Verify signature using SSZ-encoded bytes +/// Verify signature using SSZ-encoded bytes (no object allocation). pub fn verifySsz( pubkey_bytes: []const u8, message: []const u8, epoch: u32, signature_bytes: []const u8, ) HashSigError!void { - if (message.len != 32) { - return HashSigError.InvalidMessageLength; - } - - const result = hashsig_verify_ssz( + if (message.len != 32) return HashSigError.InvalidMessageLength; + switch (hashsig_verify_ssz( pubkey_bytes.ptr, pubkey_bytes.len, message.ptr, epoch, signature_bytes.ptr, signature_bytes.len, - ); - - switch (result) { + )) { 1 => {}, 0 => return HashSigError.VerificationFailed, -1 => return HashSigError.InvalidSignature, @@ -148,28 +180,22 @@ pub fn verifySsz( } } -/// Verify signature against the leanSpec test scheme. Mirrors verifySsz but -/// dispatches to the test-config FFI symbol. Used by spec-test fixtures. +/// Verify signature against the leanSpec test scheme (LOG_LIFETIME=8). pub fn verifySszTest( pubkey_bytes: []const u8, message: []const u8, epoch: u32, signature_bytes: []const u8, ) HashSigError!void { - if (message.len != 32) { - return HashSigError.InvalidMessageLength; - } - - const result = hashsig_test_verify_ssz( + if (message.len != 32) return HashSigError.InvalidMessageLength; + switch (hashsig_test_verify_ssz( pubkey_bytes.ptr, pubkey_bytes.len, message.ptr, epoch, signature_bytes.ptr, signature_bytes.len, - ); - - switch (result) { + )) { 1 => {}, 0 => return HashSigError.VerificationFailed, -1 => return HashSigError.InvalidSignature, @@ -177,247 +203,335 @@ pub fn verifySszTest( } } -/// Wrapper for the hash signature key pair +// ─── KeyPair ───────────────────────────────────────────────────────────────── + +/// Wrapper for the XMSS key pair. +/// +/// The Rust `KeyPair` struct lives in a C-heap buffer owned by this wrapper; +/// Rust never Box-allocates it. `deinit` runs Rust Drop in-place then frees +/// the buffer. pub const KeyPair = struct { - handle: *HashSigKeyPair, + /// C-heap buffer holding the Rust KeyPair struct. + _buf: [*]u8, + /// Pointer into `_buf` for the embedded public key. public_key: *const HashSigPublicKey, + /// Pointer into `_buf` for the embedded private key. private_key: *const HashSigPrivateKey, + /// Zig allocator used for ephemeral work (e.g. null-terminated seed string). allocator: Allocator, const Self = @This(); - /// Generate a new key pair + /// Generate a new key pair. pub fn generate( allocator: Allocator, seed_phrase: []const u8, activation_epoch: usize, num_active_epochs: usize, ) HashSigError!Self { - // Create null-terminated string for C - const c_seed = try allocator.dupeZ(u8, seed_phrase); + const buf = try cAlloc(hashsig_sizeof_keypair()); + errdefer cFree(buf); + + const kp: *HashSigKeyPair = @ptrCast(buf); + + const c_seed = allocator.dupeZ(u8, seed_phrase) catch return HashSigError.OutOfMemory; defer allocator.free(c_seed); - const handle = hashsig_keypair_generate( - c_seed.ptr, - activation_epoch, - num_active_epochs, - ) orelse { + if (hashsig_keypair_generate_into(kp, c_seed.ptr, activation_epoch, num_active_epochs) != 0) return HashSigError.KeyGenerationFailed; - }; - const public_key = hashsig_keypair_get_public_key(handle) orelse { - hashsig_keypair_free(handle); + const public_key = hashsig_keypair_get_public_key(kp) orelse { + hashsig_keypair_deinit(kp); return HashSigError.KeyGenerationFailed; }; - - const private_key = hashsig_keypair_get_private_key(handle) orelse { - hashsig_keypair_free(handle); + const private_key = hashsig_keypair_get_private_key(kp) orelse { + hashsig_keypair_deinit(kp); return HashSigError.KeyGenerationFailed; }; return Self{ - .handle = handle, + ._buf = buf, .public_key = public_key, .private_key = private_key, .allocator = allocator, }; } - /// Reconstruct a key pair from SSZ-encoded bytes + /// Reconstruct a key pair from SSZ-encoded bytes. pub fn fromSsz( allocator: Allocator, private_key_ssz: []const u8, public_key_ssz: []const u8, ) HashSigError!Self { - if (private_key_ssz.len == 0 or public_key_ssz.len == 0) { + if (private_key_ssz.len == 0 or public_key_ssz.len == 0) return HashSigError.DeserializationFailed; - } - const handle = hashsig_keypair_from_ssz( + const buf = try cAlloc(hashsig_sizeof_keypair()); + errdefer cFree(buf); + + const kp: *HashSigKeyPair = @ptrCast(buf); + + if (hashsig_keypair_from_ssz_into( + kp, private_key_ssz.ptr, private_key_ssz.len, public_key_ssz.ptr, public_key_ssz.len, - ) orelse { - return HashSigError.DeserializationFailed; - }; + ) != 0) return HashSigError.DeserializationFailed; - const public_key = hashsig_keypair_get_public_key(handle) orelse { - hashsig_keypair_free(handle); + const public_key = hashsig_keypair_get_public_key(kp) orelse { + hashsig_keypair_deinit(kp); return HashSigError.DeserializationFailed; }; - - const private_key = hashsig_keypair_get_private_key(handle) orelse { - hashsig_keypair_free(handle); + const private_key = hashsig_keypair_get_private_key(kp) orelse { + hashsig_keypair_deinit(kp); return HashSigError.DeserializationFailed; }; return Self{ - .handle = handle, + ._buf = buf, .public_key = public_key, .private_key = private_key, .allocator = allocator, }; } - /// Sign a message - /// Caller owns the returned signature and must free it with deinit() - pub fn sign( - self: *const Self, - message: []const u8, - epoch: u32, - ) HashSigError!Signature { - const msg_len = hashsig_message_length(); - if (message.len != msg_len) { - return HashSigError.InvalidMessageLength; - } - - const sig_handle = hashsig_sign( - self.private_key, - message.ptr, - epoch, - ) orelse { - return HashSigError.SigningFailed; - }; - - return Signature{ .handle = sig_handle }; + /// Sign a message. Caller owns the returned `Signature` and must call + /// `Signature.deinit` when done. + pub fn sign(self: *const Self, message: []const u8, epoch: u32) HashSigError!Signature { + if (message.len != hashsig_message_length()) return HashSigError.InvalidMessageLength; + return Signature.fromPrivKey(self.private_key, message, epoch); } - /// Verify a signature - pub fn verify( - self: *const Self, - message: []const u8, - signature: *const Signature, - epoch: u32, - ) HashSigError!void { - const msg_len = hashsig_message_length(); - if (message.len != msg_len) { - return HashSigError.InvalidMessageLength; - } - - const result = hashsig_verify( - self.public_key, - message.ptr, - epoch, - signature.handle, - ); - - if (result != 1) { - return HashSigError.VerificationFailed; + /// Verify a signature. + pub fn verify(self: *const Self, message: []const u8, signature: *const Signature, epoch: u32) HashSigError!void { + if (message.len != hashsig_message_length()) return HashSigError.InvalidMessageLength; + switch (hashsig_verify(self.public_key, message.ptr, epoch, signature.handle)) { + 1 => {}, + else => return HashSigError.VerificationFailed, } } - /// Get the required message length + /// Get the required message length. pub fn messageLength() usize { return hashsig_message_length(); } - /// Serialize public key to bytes (SSZ format) + /// Serialize the public key to SSZ bytes. pub fn pubkeyToBytes(self: *const Self, buffer: []u8) HashSigError!usize { - const bytes_written = hashsig_public_key_to_bytes( - self.public_key, - buffer.ptr, - buffer.len, - ); - - if (bytes_written == 0) { - return HashSigError.SerializationFailed; - } - - return bytes_written; + const n = hashsig_public_key_to_bytes(self.public_key, buffer.ptr, buffer.len); + if (n == 0) return HashSigError.SerializationFailed; + return n; } - /// Serialize private key to bytes (SSZ format) + /// Serialize the private key to SSZ bytes. pub fn privkeyToBytes(self: *const Self, buffer: []u8) HashSigError!usize { - const bytes_written = hashsig_private_key_to_bytes( - self.private_key, - buffer.ptr, - buffer.len, - ); - - if (bytes_written == 0) { - return HashSigError.SerializationFailed; - } - - return bytes_written; + const n = hashsig_private_key_to_bytes(self.private_key, buffer.ptr, buffer.len); + if (n == 0) return HashSigError.SerializationFailed; + return n; } - /// Free the key pair + /// Destroy the key pair and free its storage. pub fn deinit(self: *Self) void { - hashsig_keypair_free(self.handle); + hashsig_keypair_deinit(@ptrCast(self._buf)); + cFree(self._buf); } }; -/// Wrapper for the hash signature +// ─── Signature ─────────────────────────────────────────────────────────────── + +/// Wrapper for an XMSS signature. +/// +/// The Rust `Signature` struct lives in a C-heap buffer owned by this wrapper. +/// `handle` is an opaque pointer into that buffer, used for FFI calls. pub const Signature = struct { + /// C-heap buffer holding the Rust Signature struct. + _buf: [*]u8, + /// Opaque FFI view of the struct for use in verify / aggregate calls. handle: *HashSigSignature, const Self = @This(); - /// Deserialize a signature from SSZ bytes + /// Sign with a private key; returns a Signature owning its own storage. + /// Internal helper used by `KeyPair.sign`. + fn fromPrivKey(private_key: *const HashSigPrivateKey, message: []const u8, epoch: u32) HashSigError!Self { + const buf = try cAlloc(hashsig_sizeof_signature()); + errdefer cFree(buf); + + const sig: *HashSigSignature = @ptrCast(buf); + + if (hashsig_sign_into(sig, private_key, message.ptr, epoch) != 0) + return HashSigError.SigningFailed; + + return Self{ ._buf = buf, .handle = sig }; + } + + /// Deserialize a Signature from SSZ bytes. pub fn fromBytes(bytes: []const u8) HashSigError!Self { - if (bytes.len == 0) { - return HashSigError.DeserializationFailed; - } + if (bytes.len == 0) return HashSigError.DeserializationFailed; - const handle = hashsig_signature_from_ssz( - bytes.ptr, - bytes.len, - ) orelse { + const buf = try cAlloc(hashsig_sizeof_signature()); + errdefer cFree(buf); + + const sig: *HashSigSignature = @ptrCast(buf); + + if (hashsig_signature_from_ssz_into(sig, bytes.ptr, bytes.len) != 0) return HashSigError.DeserializationFailed; - }; - return Self{ .handle = handle }; + return Self{ ._buf = buf, .handle = sig }; } - /// Serialize signature to bytes (SSZ format) - /// Returns the number of bytes written to the buffer + /// Serialize the signature to SSZ bytes. pub fn toBytes(self: *const Self, buffer: []u8) HashSigError!usize { - const bytes_written = hashsig_signature_to_bytes( - self.handle, - buffer.ptr, - buffer.len, - ); - - if (bytes_written == 0) { - return HashSigError.SerializationFailed; - } - - return bytes_written; + const n = hashsig_signature_to_bytes(self.handle, buffer.ptr, buffer.len); + if (n == 0) return HashSigError.SerializationFailed; + return n; } - /// Free the signature + /// Destroy the signature and free its storage. pub fn deinit(self: *Self) void { - hashsig_signature_free(self.handle); + hashsig_signature_deinit(self.handle); + cFree(self._buf); } }; -/// Wrapper for standalone public keys reconstructed from SSZ bytes +// ─── PublicKey ──────────────────────────────────────────────────────────────── + +/// Wrapper for a standalone XMSS public key (e.g. deserialized for cache use). +/// +/// `handle` is an opaque pointer into the C-heap buffer for FFI calls. pub const PublicKey = struct { + /// C-heap buffer holding the Rust PublicKey struct. + _buf: [*]u8, + /// Opaque FFI view of the struct. handle: *HashSigPublicKey, const Self = @This(); + /// Deserialize a public key from SSZ bytes. pub fn fromBytes(bytes: []const u8) HashSigError!Self { - if (bytes.len == 0) { - return HashSigError.DeserializationFailed; - } + if (bytes.len == 0) return HashSigError.DeserializationFailed; - const handle = hashsig_public_key_from_ssz( - bytes.ptr, - bytes.len, - ) orelse { + const buf = try cAlloc(hashsig_sizeof_public_key()); + errdefer cFree(buf); + + const pk: *HashSigPublicKey = @ptrCast(buf); + + if (hashsig_public_key_from_ssz_into(pk, bytes.ptr, bytes.len) != 0) return HashSigError.DeserializationFailed; - }; - return Self{ .handle = handle }; + return Self{ ._buf = buf, .handle = pk }; } + /// Destroy the public key and free its storage. pub fn deinit(self: *Self) void { - hashsig_public_key_free(self.handle); + hashsig_public_key_deinit(self.handle); + cFree(self._buf); } }; +// ─── PublicKeyCache ────────────────────────────────────────────────────────── + +/// Lock-free cache for validator public keys, indexed by validator index. +/// +/// Each slot is a single `usize` atomic holding `*HashSigPublicKey` (cast +/// to `usize`); 0 is the empty sentinel. Reads are a single atomic load +/// — no mutex on the hot path. Population is lazy: a miss runs +/// `PublicKey.fromBytes` and CAS-installs the handle; lost-race writers +/// free their handle and adopt the winner's. +/// +/// Replaces the previous `std.AutoHashMap` + `pubkey_cache_lock` design +/// (P1 of #863). The cache backing is sized to `numValidators()` at +/// chain init; out-of-range indices fall through to a non-cached +/// deserialise. Validator-set growth (post-genesis additions) is not +/// supported here yet — we expect that when leanSpec adds it, the +/// fork-boundary handler will rebuild the cache with the new size. +pub const PublicKeyCache = struct { + /// One atomic per validator index; stores `@intFromPtr(handle)` or 0. + slots: []std.atomic.Value(usize), + allocator: Allocator, + + const Self = @This(); + const EMPTY: usize = 0; + + pub fn init(allocator: Allocator, capacity: usize) !Self { + const slots = try allocator.alloc(std.atomic.Value(usize), capacity); + for (slots) |*s| s.* = std.atomic.Value(usize).init(EMPTY); + return .{ .slots = slots, .allocator = allocator }; + } + + pub fn deinit(self: *Self) void { + for (self.slots) |*s| { + const ptr_int = s.load(.monotonic); + if (ptr_int != EMPTY) { + // `handle` and `_buf` both point to the same C-heap allocation + // (handle = @ptrCast(_buf) at construction time). + const raw: [*]u8 = @ptrFromInt(ptr_int); + var pk = PublicKey{ + ._buf = raw, + .handle = @ptrCast(raw), + }; + pk.deinit(); + } + } + self.allocator.free(self.slots); + } + + /// Get a cached public key handle, deserialising from bytes on miss + /// and CAS-installing the result. Returns the raw + /// `*const HashSigPublicKey` for FFI use; the cache retains + /// ownership of the handle for its full lifetime. + /// + /// Returns `HashSigError.ValidatorIndexOutOfRange` when + /// `validator_index >= capacity`. The cache is sized at + /// `BeamChain.init` from `genesis.numValidators()`; lean spec does + /// not currently grow the validator set after genesis. If/when + /// post-genesis growth lands, the fork-boundary handler must + /// rebuild the cache with the new size — until then we fail loudly + /// rather than fall back to a leaky uncached deserialise (PR #884 + /// review by @zclawz). + pub fn getOrPut(self: *Self, validator_index: usize, pubkey_bytes: []const u8) HashSigError!*const HashSigPublicKey { + if (validator_index >= self.slots.len) { + return HashSigError.ValidatorIndexOutOfRange; + } + + const slot = &self.slots[validator_index]; + const existing = slot.load(.acquire); + if (existing != EMPTY) return @ptrFromInt(existing); + + var pk = try PublicKey.fromBytes(pubkey_bytes); + // The `handle` == `_buf` (both point to the start of the C-heap allocation + // for the Rust PublicKey struct), so storing `handle` as the atomic value + // is sufficient to reconstruct both fields in `deinit`. + const new_int = @intFromPtr(pk.handle); + + if (slot.cmpxchgStrong(EMPTY, new_int, .release, .acquire)) |loser| { + // Another thread populated this slot first. Free our + // freshly-deserialised handle and adopt the winner's. + pk.deinit(); + return @ptrFromInt(loser); + } + return pk.handle; + } + + /// Check if a validator's public key is already cached. + pub fn contains(self: *const Self, validator_index: usize) bool { + if (validator_index >= self.slots.len) return false; + return self.slots[validator_index].load(.monotonic) != EMPTY; + } + + /// Best-effort count of populated slots. + pub fn count(self: *const Self) usize { + var n: usize = 0; + for (self.slots) |*s| { + if (s.load(.monotonic) != EMPTY) n += 1; + } + return n; + } +}; + +// ─── Tests ──────────────────────────────────────────────────────────────────── + test "HashSig: generate keypair" { const allocator = std.testing.allocator; @@ -431,23 +545,18 @@ test "HashSig: generate keypair" { test "HashSig: SSZ keypair roundtrip" { const allocator = std.testing.allocator; - // Generate original keypair var keypair = try KeyPair.generate(allocator, "test_ssz_roundtrip", 0, 5); defer keypair.deinit(); - // Serialize to SSZ var pk_buffer: [256]u8 = undefined; const pk_len = try keypair.pubkeyToBytes(&pk_buffer); - // We need a large buffer for private key (it contains many one-time keys) - // Allocating on heap to be safe with stack size - const sk_buffer = try allocator.alloc(u8, 1024 * 1024 * 10); // 10MB should be enough + const sk_buffer = try allocator.alloc(u8, 1024 * 1024 * 10); defer allocator.free(sk_buffer); const sk_len = try keypair.privkeyToBytes(sk_buffer); std.debug.print("\nPK size: {d}, SK size: {d}\n", .{ pk_len, sk_len }); - // Reconstruct from SSZ var restored_keypair = try KeyPair.fromSsz( allocator, sk_buffer[0..sk_len], @@ -455,15 +564,12 @@ test "HashSig: SSZ keypair roundtrip" { ); defer restored_keypair.deinit(); - // Verify functionality with restored keypair const message = [_]u8{42} ** 32; const epoch: u32 = 0; - // Sign with restored keypair var signature = try restored_keypair.sign(&message, epoch); defer signature.deinit(); - // Verify with original keypair (should work as they are same keys) try keypair.verify(&message, &signature, epoch); } @@ -473,32 +579,26 @@ test "HashSig: sign and verify" { var keypair = try KeyPair.generate(allocator, "test_seed", 0, 2); defer keypair.deinit(); - // Create a message of the correct length const msg_len = KeyPair.messageLength(); const message = try allocator.alloc(u8, msg_len); defer allocator.free(message); - // Fill with test data for (message, 0..) |*byte, i| { byte.* = @intCast(i % 256); } const epoch: u32 = 0; - // Sign the message var signature = try keypair.sign(message, epoch); defer signature.deinit(); - // Verify the signature try keypair.verify(message, &signature, epoch); - // Test with wrong epoch keypair.verify(message, &signature, epoch + 100) catch |err| { try std.testing.expect(err == HashSigError.VerificationFailed); }; - // Test with wrong message - message[0] = message[0] + 1; // Modify message + message[0] = message[0] + 1; keypair.verify(message, &signature, epoch) catch |err| { try std.testing.expect(err == HashSigError.VerificationFailed); }; @@ -513,10 +613,7 @@ test "HashSig: invalid message length" { const wrong_message = try allocator.alloc(u8, 10); defer allocator.free(wrong_message); - const epoch: u32 = 0; - - // Should fail with invalid message length - const result = keypair.sign(wrong_message, epoch); + const result = keypair.sign(wrong_message, 0); try std.testing.expectError(HashSigError.InvalidMessageLength, result); } @@ -529,21 +626,17 @@ test "HashSig: SSZ serialize and verify" { const message = [_]u8{1} ** 32; const epoch: u32 = 0; - // Sign var signature = try keypair.sign(&message, epoch); defer signature.deinit(); - // Serialize signature var sig_buffer: [4000]u8 = undefined; const sig_size = try signature.toBytes(&sig_buffer); std.debug.print("\nSignature size: {d} bytes\n", .{sig_size}); - // Serialize public key var pubkey_buffer: [256]u8 = undefined; const pubkey_size = try keypair.pubkeyToBytes(&pubkey_buffer); std.debug.print("Public key size: {d} bytes\n", .{pubkey_size}); - // Verify using SSZ try verifySsz( pubkey_buffer[0..pubkey_size], &message, @@ -563,7 +656,6 @@ test "HashSig: verify fails with zero signature" { const message = [_]u8{1} ** 32; const epoch: u32 = 0; - // Serialize public key var pubkey_buffer: [256]u8 = undefined; const pubkey_size = try keypair.pubkeyToBytes(&pubkey_buffer); @@ -574,10 +666,8 @@ test "HashSig: verify fails with zero signature" { const signature_size = try signature.toBytes(&signature_buffer); - // Create invalid signature with all zeros var zero_sig_buffer = [_]u8{0} ** 4000; - // Invalid signature length - should fail with InvalidSignature const invalid_signature_result = verifySsz( pubkey_buffer[0..pubkey_size], &message, @@ -588,7 +678,6 @@ test "HashSig: verify fails with zero signature" { try std.testing.expectError(HashSigError.InvalidSignature, invalid_signature_result); const invalid_message = [_]u8{2} ** 32; - // Verification should fail - should fail with VerificationFailed const verification_failed_result = verifySsz( pubkey_buffer[0..pubkey_size], &invalid_message, diff --git a/rust/hashsig-glue/src/lib.rs b/rust/hashsig-glue/src/lib.rs index c2f0bfe98..f64b40646 100644 --- a/rust/hashsig-glue/src/lib.rs +++ b/rust/hashsig-glue/src/lib.rs @@ -152,88 +152,134 @@ fn xmss_secret_key_to_ssz(sk: &HashSigPrivateKey) -> Vec { // FFI Functions for Zig interop -/// Generate a new key pair -/// Returns a pointer to the KeyPair or null on error +// ─── Layout queries ────────────────────────────────────────────────────────── +// +// Zig callers query these at startup so they can pre-allocate correctly-sized +// and correctly-aligned storage, then pass it to the `_into` init functions +// below. Rust never Box-allocates the outer KeyPair / Signature / PublicKey +// structs; allocation and lifetime are owned entirely by the Zig side. + +/// Size in bytes of the KeyPair struct. +#[no_mangle] +pub extern "C" fn hashsig_sizeof_keypair() -> usize { + std::mem::size_of::() +} + +/// Required alignment in bytes of the KeyPair struct. +#[no_mangle] +pub extern "C" fn hashsig_alignof_keypair() -> usize { + std::mem::align_of::() +} + +/// Size in bytes of the Signature struct. +#[no_mangle] +pub extern "C" fn hashsig_sizeof_signature() -> usize { + std::mem::size_of::() +} + +/// Required alignment in bytes of the Signature struct. +#[no_mangle] +pub extern "C" fn hashsig_alignof_signature() -> usize { + std::mem::align_of::() +} + +/// Size in bytes of the PublicKey struct. +#[no_mangle] +pub extern "C" fn hashsig_sizeof_public_key() -> usize { + std::mem::size_of::() +} + +/// Required alignment in bytes of the PublicKey struct. +#[no_mangle] +pub extern "C" fn hashsig_alignof_public_key() -> usize { + std::mem::align_of::() +} + +// ─── Placement-init (no Rust heap allocation for the outer struct) ──────────── +// +// Each `_into` function writes a fully-initialised struct into caller-supplied +// storage. The inner XMSS types still carry their own heap-allocated data +// (leansig `Vec`), but the outer wrapper struct itself lives in Zig memory. +// Callers must later call the matching `_deinit` to run Rust Drop in-place, +// then free the buffer themselves. + +/// Generate a new key pair into caller-supplied storage. +/// +/// `out` must point to at least `hashsig_sizeof_keypair()` bytes aligned to at +/// least `hashsig_alignof_keypair()` (C `malloc` always satisfies this). On +/// success the struct is written in-place and the caller must later call +/// `hashsig_keypair_deinit` before freeing `out`. Returns 0 on success, -1 on +/// error; on error the storage is left uninitialised. +/// /// # Safety -/// This is meant to be called from zig, so the pointers will always dereference correctly +/// `out` must be a valid, properly-aligned, non-null pointer to at least +/// `hashsig_sizeof_keypair()` bytes. `seed_phrase` must be a valid C string. #[no_mangle] -pub unsafe extern "C" fn hashsig_keypair_generate( +pub unsafe extern "C" fn hashsig_keypair_generate_into( + out: *mut KeyPair, seed_phrase: *const c_char, activation_epoch: usize, num_active_epochs: usize, -) -> *mut KeyPair { - let seed_phrase = unsafe { CStr::from_ptr(seed_phrase).to_string_lossy().into_owned() }; - - // Hash the seed phrase to get a 32-byte seed +) -> i32 { + if out.is_null() || seed_phrase.is_null() { + return -1; + } + let seed_phrase = CStr::from_ptr(seed_phrase).to_string_lossy().into_owned(); let mut hasher = Sha256::new(); hasher.update(seed_phrase.as_bytes()); let seed = hasher.finalize().into(); - let (public_key, private_key) = PrivateKey::generate( &mut StdRng::from_seed(seed), activation_epoch as u32, num_active_epochs as u32, ); - - let keypair = Box::new(KeyPair { - public_key, - private_key, - }); - - Box::into_raw(keypair) + std::ptr::write(out, KeyPair { public_key, private_key }); + 0 } -/// Reconstruct a key pair from SSZ-encoded secret and public keys -/// Returns a pointer to the KeyPair or null on error +/// Reconstruct a key pair from SSZ-encoded bytes into caller-supplied storage. +/// Returns 0 on success, -1 on error. +/// /// # Safety -/// This is meant to be called from zig, so the pointers will always dereference correctly +/// `out` must be a valid, properly-aligned, non-null pointer to at least +/// `hashsig_sizeof_keypair()` bytes. Key byte pointers must be valid for their +/// respective lengths. #[no_mangle] -pub unsafe extern "C" fn hashsig_keypair_from_ssz( +pub unsafe extern "C" fn hashsig_keypair_from_ssz_into( + out: *mut KeyPair, private_key_ptr: *const u8, private_key_len: usize, public_key_ptr: *const u8, public_key_len: usize, -) -> *mut KeyPair { - if private_key_ptr.is_null() || public_key_ptr.is_null() { - return ptr::null_mut(); - } - - unsafe { - let sk_slice = slice::from_raw_parts(private_key_ptr, private_key_len); - let pk_slice = slice::from_raw_parts(public_key_ptr, public_key_len); - - let private_key: HashSigPrivateKey = match xmss_secret_key_from_ssz(sk_slice) { - Ok(key) => key, - Err(_) => { - return ptr::null_mut(); - } - }; - - let public_key: HashSigPublicKey = match xmss_public_key_from_ssz(pk_slice) { - Ok(key) => key, - Err(_) => { - return ptr::null_mut(); - } - }; - - let keypair = Box::new(KeyPair { - public_key: PublicKey::new(public_key), - private_key: PrivateKey::new(private_key), - }); - - Box::into_raw(keypair) +) -> i32 { + if out.is_null() || private_key_ptr.is_null() || public_key_ptr.is_null() { + return -1; } + let sk_slice = slice::from_raw_parts(private_key_ptr, private_key_len); + let pk_slice = slice::from_raw_parts(public_key_ptr, public_key_len); + let private_key = match xmss_secret_key_from_ssz(sk_slice) { + Ok(k) => PrivateKey::new(k), + Err(_) => return -1, + }; + let public_key = match xmss_public_key_from_ssz(pk_slice) { + Ok(k) => PublicKey::new(k), + Err(_) => return -1, + }; + std::ptr::write(out, KeyPair { public_key, private_key }); + 0 } -/// Free a key pair +/// Destroy a KeyPair in-place (runs Rust `Drop`). Does NOT free the storage — +/// the caller owns and must free it after this call. Safe to call with null +/// (no-op). +/// /// # Safety -/// This is meant to be called from zig, so the pointers will always dereference correctly +/// `kp` must either be null or point to a fully-initialised `KeyPair` previously +/// set up by `hashsig_keypair_generate_into` or `hashsig_keypair_from_ssz_into`. #[no_mangle] -pub unsafe extern "C" fn hashsig_keypair_free(keypair: *mut KeyPair) { - if !keypair.is_null() { - unsafe { - let _ = Box::from_raw(keypair); - } +pub unsafe extern "C" fn hashsig_keypair_deinit(kp: *mut KeyPair) { + if !kp.is_null() { + std::ptr::drop_in_place(kp); } } @@ -269,117 +315,109 @@ pub unsafe extern "C" fn hashsig_keypair_get_private_key( &(*keypair).private_key } -/// Construct a standalone public key from SSZ-encoded bytes. -/// Returns a pointer to PublicKey or null on error. +/// Deserialize a PublicKey from SSZ bytes into caller-supplied storage. +/// Returns 0 on success, -1 on error; on error the storage is left uninitialised. +/// /// # Safety -/// Inputs must be valid pointers and buffers. +/// `out` must be a valid, properly-aligned, non-null pointer to at least +/// `hashsig_sizeof_public_key()` bytes. `bytes` must be valid for `len` bytes. #[no_mangle] -pub unsafe extern "C" fn hashsig_public_key_from_ssz( - public_key_ptr: *const u8, - public_key_len: usize, -) -> *mut PublicKey { - if public_key_ptr.is_null() { - return ptr::null_mut(); - } - - unsafe { - let pk_slice = slice::from_raw_parts(public_key_ptr, public_key_len); - let public_key: HashSigPublicKey = match xmss_public_key_from_ssz(pk_slice) { - Ok(key) => key, - Err(_) => { - return ptr::null_mut(); - } - }; - - Box::into_raw(Box::new(PublicKey::new(public_key))) +pub unsafe extern "C" fn hashsig_public_key_from_ssz_into( + out: *mut PublicKey, + bytes: *const u8, + len: usize, +) -> i32 { + if out.is_null() || bytes.is_null() || len == 0 { + return -1; } + let slice = slice::from_raw_parts(bytes, len); + let pk = match xmss_public_key_from_ssz(slice) { + Ok(k) => PublicKey::new(k), + Err(_) => return -1, + }; + std::ptr::write(out, pk); + 0 } -/// Free a public key created via hashsig_public_key_from_ssz. +/// Destroy a PublicKey in-place. Does NOT free the storage. +/// Safe to call with null (no-op). +/// /// # Safety -/// Pointer must be valid or null. +/// `pk` must either be null or point to a fully-initialised `PublicKey`. #[no_mangle] -pub unsafe extern "C" fn hashsig_public_key_free(public_key: *mut PublicKey) { - if !public_key.is_null() { - unsafe { - let _ = Box::from_raw(public_key); - } +pub unsafe extern "C" fn hashsig_public_key_deinit(pk: *mut PublicKey) { + if !pk.is_null() { + std::ptr::drop_in_place(pk); } } -/// Sign a message using a private key directly -/// Returns pointer to Signature on success, null on error +/// Sign a message, placing the result into caller-supplied Signature storage. +/// +/// `out` must point to at least `hashsig_sizeof_signature()` bytes aligned to +/// at least `hashsig_alignof_signature()`. Returns 0 on success, -1 on error. +/// /// # Safety -/// This is meant to be called from zig, so it's safe as the pointer will always exist +/// `out` must be a valid, properly-aligned, non-null pointer. `private_key` +/// and `message_ptr` must be valid non-null pointers; `message_ptr` must point +/// to at least `MESSAGE_LENGTH` bytes. #[no_mangle] -pub unsafe extern "C" fn hashsig_sign( +pub unsafe extern "C" fn hashsig_sign_into( + out: *mut Signature, private_key: *const PrivateKey, message_ptr: *const u8, epoch: u32, -) -> *mut Signature { - if private_key.is_null() || message_ptr.is_null() { - return ptr::null_mut(); - } - - unsafe { - let private_key_ref = &*private_key; - let message_slice = slice::from_raw_parts(message_ptr, MESSAGE_LENGTH); - - // Convert slice to array - let message_array: &[u8; MESSAGE_LENGTH] = match message_slice.try_into() { - Ok(arr) => arr, - Err(_) => { - return ptr::null_mut(); - } - }; - - let signature = match private_key_ref.sign(message_array, epoch) { - Ok(sig) => sig, - Err(_) => { - return ptr::null_mut(); - } - }; - - Box::into_raw(Box::new(signature)) +) -> i32 { + if out.is_null() || private_key.is_null() || message_ptr.is_null() { + return -1; } + let private_key_ref = &*private_key; + let message_slice = slice::from_raw_parts(message_ptr, MESSAGE_LENGTH); + let message_array: &[u8; MESSAGE_LENGTH] = match message_slice.try_into() { + Ok(arr) => arr, + Err(_) => return -1, + }; + let signature = match private_key_ref.sign(message_array, epoch) { + Ok(sig) => sig, + Err(_) => return -1, + }; + std::ptr::write(out, signature); + 0 } -/// Free a signature +/// Destroy a Signature in-place. Does NOT free the storage. +/// Safe to call with null (no-op). +/// /// # Safety -/// This is meant to be called from zig, so it's safe as the pointer will always exist +/// `sig` must either be null or point to a fully-initialised `Signature`. #[no_mangle] -pub unsafe extern "C" fn hashsig_signature_free(signature: *mut Signature) { - if !signature.is_null() { - unsafe { - let _ = Box::from_raw(signature); - } +pub unsafe extern "C" fn hashsig_signature_deinit(sig: *mut Signature) { + if !sig.is_null() { + std::ptr::drop_in_place(sig); } } -/// Construct a signature from SSZ-encoded bytes. -/// Returns a pointer to Signature or null on error. +/// Deserialize a Signature from SSZ bytes into caller-supplied storage. +/// Returns 0 on success, -1 on error. +/// /// # Safety -/// Inputs must be valid pointers and buffers. +/// `out` must be a valid, properly-aligned, non-null pointer to at least +/// `hashsig_sizeof_signature()` bytes. `bytes` must be valid for `len` bytes. #[no_mangle] -pub unsafe extern "C" fn hashsig_signature_from_ssz( - signature_ptr: *const u8, - signature_len: usize, -) -> *mut Signature { - if signature_ptr.is_null() || signature_len == 0 { - return ptr::null_mut(); - } - - unsafe { - let sig_slice = slice::from_raw_parts(signature_ptr, signature_len); - let signature: HashSigSignature = match xmss_signature_from_ssz(sig_slice) { - Ok(sig) => sig, - Err(_) => { - return ptr::null_mut(); - } - }; - - Box::into_raw(Box::new(Signature { inner: signature })) +pub unsafe extern "C" fn hashsig_signature_from_ssz_into( + out: *mut Signature, + bytes: *const u8, + len: usize, +) -> i32 { + if out.is_null() || bytes.is_null() || len == 0 { + return -1; } + let slice = slice::from_raw_parts(bytes, len); + let sig = match xmss_signature_from_ssz(slice) { + Ok(s) => Signature::new(s), + Err(_) => return -1, + }; + std::ptr::write(out, sig); + 0 } /// Verify a signature using a public key directly From 6a719b5b0344773c279fa40e79113b01459eb0ba Mon Sep 17 00:00:00 2001 From: zclawz Date: Tue, 26 May 2026 07:40:04 +0000 Subject: [PATCH 2/3] hashsig-glue: apply cargo fmt --- rust/hashsig-glue/src/lib.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/rust/hashsig-glue/src/lib.rs b/rust/hashsig-glue/src/lib.rs index f64b40646..5e7a34f71 100644 --- a/rust/hashsig-glue/src/lib.rs +++ b/rust/hashsig-glue/src/lib.rs @@ -233,7 +233,13 @@ pub unsafe extern "C" fn hashsig_keypair_generate_into( activation_epoch as u32, num_active_epochs as u32, ); - std::ptr::write(out, KeyPair { public_key, private_key }); + std::ptr::write( + out, + KeyPair { + public_key, + private_key, + }, + ); 0 } @@ -265,7 +271,13 @@ pub unsafe extern "C" fn hashsig_keypair_from_ssz_into( Ok(k) => PublicKey::new(k), Err(_) => return -1, }; - std::ptr::write(out, KeyPair { public_key, private_key }); + std::ptr::write( + out, + KeyPair { + public_key, + private_key, + }, + ); 0 } From 2efa790c66ea5f40c8704ad7d1ee2ea7580bd06a Mon Sep 17 00:00:00 2001 From: zclawz Date: Tue, 26 May 2026 08:10:37 +0000 Subject: [PATCH 3/3] xmss: fix PublicKeyCache::deinit in lib.zig to supply _buf field PublicKey gained a _buf field in the placement-init refactor (PR #935) but the PublicKeyCache::deinit in lib.zig still used the old single-field struct init. Match the already-corrected pattern from hashsig.zig: reconstruct both _buf and handle from the stored atomic integer (they both point to the start of the same C-heap allocation). Fixes build/test CI failures on macOS. --- pkgs/xmss/src/lib.zig | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/pkgs/xmss/src/lib.zig b/pkgs/xmss/src/lib.zig index 33636e34b..dbc53762d 100644 --- a/pkgs/xmss/src/lib.zig +++ b/pkgs/xmss/src/lib.zig @@ -57,7 +57,13 @@ pub const PublicKeyCache = struct { for (self.slots) |*s| { const ptr_int = s.load(.monotonic); if (ptr_int != EMPTY) { - var pk = PublicKey{ .handle = @ptrFromInt(ptr_int) }; + // `handle` and `_buf` both point to the same C-heap allocation + // (handle = @ptrCast(_buf) at construction time). + const raw: [*]u8 = @ptrFromInt(ptr_int); + var pk = PublicKey{ + ._buf = raw, + .handle = @ptrCast(raw), + }; pk.deinit(); } }