From 16df613f46f039ecfff494024cbfddf1749e5281 Mon Sep 17 00:00:00 2001 From: Mehdi Date: Mon, 12 Jan 2026 10:54:49 -0500 Subject: [PATCH 1/4] Implemented requested feature for non-solana environments --- Cargo.toml | 4 +++ src/lib.rs | 80 ++++++++++++++++++++++++++++++++++++++++++++++-------- 2 files changed, 73 insertions(+), 11 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index ae2734d..e99b8fb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,3 +25,7 @@ members = [ [dev-dependencies] solana-program = "2.3" serde_json = "1.0.143" + +[features] +default = ["solana"] +solana = [] \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index c8d259e..88c33d2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,9 +14,67 @@ // along with this program. If not, see . // // SPDX-License-Identifier: AGPL-3.0-or-later - //! WOTS+ (Winternitz One-Time Signature Plus) implementation + +enum SignatureBufferEnum { + #[cfg(feature = "solana")] + Array { + buf: [u8; constants::SIGNATURE_SIZE], + len: usize, + }, + + #[cfg(not(feature = "solana"))] + Vec(Vec), +} + +impl SignatureBufferEnum { + fn new() -> Self { + #[cfg(feature = "solana")] + { + Self::Array { + buf: [0u8; constants::SIGNATURE_SIZE], + len: 0, + } + } + + #[cfg(not(feature = "solana"))] + { + Self::Vec(Vec::with_capacity(constants::SIGNATURE_SIZE)) + } + } + + fn push_slice(&mut self, data: &[u8]) { + #[cfg(feature = "solana")] + { + let Self::Array { buf, len } = self; + let end = *len + data.len(); + buf[*len..end].copy_from_slice(data); + *len = end; + } + + #[cfg(not(feature = "solana"))] + { + let Self::Vec(v) = self; + v.extend_from_slice(data); + } + } + + fn as_slice(&self) -> &[u8] { + #[cfg(feature = "solana")] + { + let Self::Array { buf, len } = self; + &buf[..*len] + } + + #[cfg(not(feature = "solana"))] + { + let Self::Vec(v) = self; + v.as_slice() + } + } +} + /// Hash function type for WOTS+ type HashFn = fn(&[u8]) -> [u8; 32]; @@ -252,7 +310,7 @@ impl WOTSPlus { let randomization_elements = self.generate_randomization_elements(&public_seed); let function_key = randomization_elements[0]; - let mut public_key_segments = Vec::with_capacity(constants::SIGNATURE_SIZE); + let mut public_key_segments = SignatureBufferEnum::new(); for i in 0..constants::NUM_SIGNATURE_CHUNKS { let mut to_hash = vec![0u8; constants::HASH_LEN * 2]; @@ -267,10 +325,10 @@ impl WOTSPlus { (constants::CHAIN_LEN - 1) as u16, ); - public_key_segments.extend_from_slice(&segment); + public_key_segments.push_slice(&segment); } - let public_key_hash = (self.hash_fn)(&public_key_segments); + let public_key_hash = (self.hash_fn)(public_key_segments.as_slice()); PublicKey { public_seed: *public_seed, @@ -353,7 +411,7 @@ impl WOTSPlus { let chain_segments = self.compute_message_hash_chain_indexes(message); - let mut public_key_segments = Vec::with_capacity(constants::SIGNATURE_SIZE); + let mut public_key_segments = SignatureBufferEnum::new(); // Compute each public key segment. These are done by taking the signature, which is prevChainOut at chainIdx, // and completing the hash chain via the chain function to recompute the public key segment. @@ -366,11 +424,11 @@ impl WOTSPlus { num_iterations, ); - public_key_segments.extend_from_slice(&segment); + public_key_segments.push_slice(&segment); } // Hash all public key segments together to recreate the original public key - let computed_hash = (self.hash_fn)(&public_key_segments); + let computed_hash = (self.hash_fn)(public_key_segments.as_slice()); // Compare computed hash with stored public key hash computed_hash == public_key.public_key_hash @@ -397,7 +455,7 @@ impl WOTSPlus { } let chain_segments = self.compute_message_hash_chain_indexes(message); - let mut public_key_segments = [0u8; constants::SIGNATURE_SIZE]; + let mut public_key_segments = SignatureBufferEnum::new(); // Compute each public key segment using the pre-computed randomization elements for (i, &chain_idx) in chain_segments.iter().enumerate() { @@ -409,12 +467,12 @@ impl WOTSPlus { num_iterations, ); - let offset = i * constants::HASH_LEN; - public_key_segments[offset..offset + constants::HASH_LEN].copy_from_slice(&segment); + // let offset = i * constants::HASH_LEN; + public_key_segments.push_slice(&segment); } // Hash all public key segments together and compare with the provided hash - let computed_hash = (self.hash_fn)(&public_key_segments); + let computed_hash = (self.hash_fn)(public_key_segments.as_slice()); computed_hash == *public_key_hash } } From 35788623a7f239ae7a468a01d77153bf828dcdc4 Mon Sep 17 00:00:00 2001 From: Mehdi Date: Mon, 12 Jan 2026 11:45:34 -0500 Subject: [PATCH 2/4] added basic tests for new solana/non-solana flows --- src/lib.rs | 69 ++++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 64 insertions(+), 5 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 88c33d2..237cb06 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,7 +17,7 @@ //! WOTS+ (Winternitz One-Time Signature Plus) implementation -enum SignatureBufferEnum { +enum SignatureBuffer { #[cfg(feature = "solana")] Array { buf: [u8; constants::SIGNATURE_SIZE], @@ -28,7 +28,7 @@ enum SignatureBufferEnum { Vec(Vec), } -impl SignatureBufferEnum { +impl SignatureBuffer { fn new() -> Self { #[cfg(feature = "solana")] { @@ -310,7 +310,7 @@ impl WOTSPlus { let randomization_elements = self.generate_randomization_elements(&public_seed); let function_key = randomization_elements[0]; - let mut public_key_segments = SignatureBufferEnum::new(); + let mut public_key_segments = SignatureBuffer::new(); for i in 0..constants::NUM_SIGNATURE_CHUNKS { let mut to_hash = vec![0u8; constants::HASH_LEN * 2]; @@ -411,7 +411,7 @@ impl WOTSPlus { let chain_segments = self.compute_message_hash_chain_indexes(message); - let mut public_key_segments = SignatureBufferEnum::new(); + let mut public_key_segments = SignatureBuffer::new(); // Compute each public key segment. These are done by taking the signature, which is prevChainOut at chainIdx, // and completing the hash chain via the chain function to recompute the public key segment. @@ -455,7 +455,7 @@ impl WOTSPlus { } let chain_segments = self.compute_message_hash_chain_indexes(message); - let mut public_key_segments = SignatureBufferEnum::new(); + let mut public_key_segments = SignatureBuffer::new(); // Compute each public key segment using the pre-computed randomization elements for (i, &chain_idx) in chain_segments.iter().enumerate() { @@ -548,6 +548,65 @@ mod tests { assert_eq!(recovered.public_key_hash, public_key.public_key_hash); } + #[test] + fn signatures_are_deterministic() { + let wots = WOTSPlus::new(mock_hash); + + let seed = [9u8; 32]; + let (_, sk) = wots.generate_key_pair(&seed); + let msg = [1u8; constants::MESSAGE_LEN]; + + let sig1 = wots.sign(&sk, &msg); + let sig2 = wots.sign(&sk, &msg); + + assert_eq!(sig1, sig2); + } + + + #[test] + fn sigbuf_appends_correctly() { + let mut buf = SignatureBuffer::new(); + + let a = [1u8; constants::HASH_LEN]; + let b = [2u8; constants::HASH_LEN]; + + buf.push_slice(&a); + buf.push_slice(&b); + + let out = buf.as_slice(); + assert_eq!(out.len(), 2 * constants::HASH_LEN); + assert_eq!(&out[..constants::HASH_LEN], &a); + assert_eq!(&out[constants::HASH_LEN..], &b); + } + + #[cfg(feature = "solana")] + #[test] + #[should_panic] + fn sigbuf_panics_on_overflow() { + let mut buf = SignatureBuffer::new(); + let chunk = [0u8; constants::HASH_LEN]; + + for _ in 0..(constants::NUM_SIGNATURE_CHUNKS + 1) { + buf.push_slice(&chunk); + } + } + + #[cfg(not(feature = "solana"))] + #[test] + fn sigbuf_allows_growth_on_heap() { + let mut buf = SignatureBuffer::new(); + let chunk = [0u8; constants::HASH_LEN]; + + for _ in 0..(constants::NUM_SIGNATURE_CHUNKS + 10) { + buf.push_slice(&chunk); + } + + assert_eq!( + buf.as_slice().len(), + (constants::NUM_SIGNATURE_CHUNKS + 10) * constants::HASH_LEN + ); + } + #[cfg(test)] mod tests { use super::*; From fce60b5ed54f3dd67957045105183f58c707e307 Mon Sep 17 00:00:00 2001 From: Mehdi Date: Wed, 14 Jan 2026 12:23:11 -0500 Subject: [PATCH 3/4] Addressed review comment, added overflow restrictions to both solana/non-solana flows --- src/lib.rs | 60 ++++++++++++++++++++++++++++-------------------------- 1 file changed, 31 insertions(+), 29 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 237cb06..67be076 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,59 +18,65 @@ enum SignatureBuffer { + #[cfg(feature = "solana")] + Vec(Vec), + + #[cfg(not(feature = "solana"))] Array { buf: [u8; constants::SIGNATURE_SIZE], len: usize, }, - - #[cfg(not(feature = "solana"))] - Vec(Vec), } impl SignatureBuffer { fn new() -> Self { #[cfg(feature = "solana")] { - Self::Array { - buf: [0u8; constants::SIGNATURE_SIZE], - len: 0, - } + Self::Vec(Vec::with_capacity(constants::SIGNATURE_SIZE)) } #[cfg(not(feature = "solana"))] { - Self::Vec(Vec::with_capacity(constants::SIGNATURE_SIZE)) + Self::Array { + buf: [0u8; constants::SIGNATURE_SIZE], + len: 0, + } } } fn push_slice(&mut self, data: &[u8]) { #[cfg(feature = "solana")] { - let Self::Array { buf, len } = self; - let end = *len + data.len(); - buf[*len..end].copy_from_slice(data); - *len = end; + let Self::Vec(v) = self; + let new_len = v.len() + data.len(); + assert!( + new_len <= constants::SIGNATURE_SIZE, + "SignatureBuffer overflow" + ); + v.extend_from_slice(data); } #[cfg(not(feature = "solana"))] { - let Self::Vec(v) = self; - v.extend_from_slice(data); + let Self::Array { buf, len } = self; + let end = *len + data.len(); + buf[*len..end].copy_from_slice(data); + *len = end; } } fn as_slice(&self) -> &[u8] { #[cfg(feature = "solana")] { - let Self::Array { buf, len } = self; - &buf[..*len] + let Self::Vec(v) = self; + v.as_slice() } #[cfg(not(feature = "solana"))] { - let Self::Vec(v) = self; - v.as_slice() + let Self::Array { buf, len } = self; + &buf[..*len] } } } @@ -579,10 +585,10 @@ mod tests { assert_eq!(&out[constants::HASH_LEN..], &b); } - #[cfg(feature = "solana")] + #[cfg(not(feature = "solana"))] #[test] #[should_panic] - fn sigbuf_panics_on_overflow() { + fn sigbuf_panics_on_overflow_non_solana() { let mut buf = SignatureBuffer::new(); let chunk = [0u8; constants::HASH_LEN]; @@ -591,22 +597,18 @@ mod tests { } } - #[cfg(not(feature = "solana"))] + #[cfg(feature = "solana")] #[test] - fn sigbuf_allows_growth_on_heap() { + #[should_panic] + fn sigbuf_panics_on_overflow_solana() { let mut buf = SignatureBuffer::new(); let chunk = [0u8; constants::HASH_LEN]; - for _ in 0..(constants::NUM_SIGNATURE_CHUNKS + 10) { + for _ in 0..(constants::NUM_SIGNATURE_CHUNKS + 1) { buf.push_slice(&chunk); } - - assert_eq!( - buf.as_slice().len(), - (constants::NUM_SIGNATURE_CHUNKS + 10) * constants::HASH_LEN - ); } - + #[cfg(test)] mod tests { use super::*; From 08936b6383abdbfad0578ceb5e3531dd46cfb3f9 Mon Sep 17 00:00:00 2001 From: Mehdi Date: Wed, 14 Jan 2026 18:40:22 -0500 Subject: [PATCH 4/4] clean ups, removed Vec::with_capacity() in sign --- src/lib.rs | 40 ++++++++++++++++++++++++++++++++++------ 1 file changed, 34 insertions(+), 6 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 67be076..d854f10 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -79,6 +79,32 @@ impl SignatureBuffer { &buf[..*len] } } + + /// Consume the buffer and return its contents as a Vec + fn as_signature_chunks(self) -> Vec<[u8; constants::HASH_LEN]> { + let slice = match self { + #[cfg(feature = "solana")] + SignatureBuffer::Vec(v) => v, + + #[cfg(not(feature = "solana"))] + SignatureBuffer::Array { buf, len } => buf[..len].to_vec(), + }; + // Enforce chunk alignment invariant + assert!( + slice.len() % constants::HASH_LEN == 0, + "SignatureBuffer length is not chunk-aligned" + ); + + slice + .chunks_exact(constants::HASH_LEN) + .map(|chunk| { + let mut arr = [0u8; constants::HASH_LEN]; + arr.copy_from_slice(chunk); + arr + }) + .collect() + } + } /// Hash function type for WOTS+ @@ -231,11 +257,13 @@ impl WOTSPlus { &self, public_seed: &[u8; constants::HASH_LEN] ) -> Vec<[u8; constants::HASH_LEN]> { - let mut elements = Vec::with_capacity(constants::NUM_SIGNATURE_CHUNKS); + + let mut elements = SignatureBuffer::new(); for i in 0..constants::NUM_SIGNATURE_CHUNKS { - elements.push(self.prf(public_seed, i as u16)); + elements.push_slice(&self.prf(public_seed, i as u16)); } - elements + elements.as_signature_chunks() + } /// XOR two 32-byte arrays @@ -376,7 +404,7 @@ impl WOTSPlus { let function_key = randomization_elements[0]; let chain_segments = self.compute_message_hash_chain_indexes(message); - let mut signature = Vec::with_capacity(constants::NUM_SIGNATURE_CHUNKS); + let mut signature = SignatureBuffer::new(); for (i, &chain_idx) in chain_segments.iter().enumerate() { let mut to_hash = vec![0u8; constants::HASH_LEN * 2]; @@ -390,10 +418,10 @@ impl WOTSPlus { 0, chain_idx as u16, ); - signature.push(sig_segment); + signature.push_slice(&sig_segment); } - signature + signature.as_signature_chunks() } /// Verify a WOTS+ signature