[frontend] Add Winternitz OTS verifier

Author: paulcadman

crates/frontend/src/circuits/hash_based_sig/mod.rs

@@ -1,3 +1,4 @@
 pub mod chain;
 pub mod codeword;
 pub mod tweak;
+pub mod winternitz_ots;

crates/frontend/src/circuits/hash_based_sig/winternitz_ots.rs

@@ -0,0 +1,350 @@
+use super::{chain::verify_chain, codeword::codeword, tweak::verify_message_tweak};
+use crate::{
+	circuits::keccak::Keccak,
+	compiler::{CircuitBuilder, Wire},
+};
+
+/// Verifies a Winternitz One-Time Signature.
+///
+/// This circuit implements verification for the Winternitz OTS scheme, which combines:
+/// 1. Message hashing with tweaking
+/// 2. Codeword extraction from the hash
+/// 3. Hash chain verification for each coordinate
+///
+/// # Arguments
+///
+/// * `builder` - Circuit builder for constructing constraints
+/// * `param` - Cryptographic parameter (32 bytes as 4x64-bit LE wires)
+/// * `message` - Message to verify (32 bytes as 4x64-bit LE wires)
+/// * `nonce` - Random nonce used in signing (23 bytes)
+/// * `signature_hashes` - Signature hash values (one per chain)
+/// * `public_key_hashes` - Public key end hashes (one per chain)
+/// * `spec` - Winternitz specification parameters
+///
+/// # Returns
+///
+/// A tuple containing:
+/// - Keccak hasher for the message tweak
+/// - Vector of Keccak hashers for the chain tweaks
+pub fn verify_winternitz_ots(
+	builder: &CircuitBuilder,
+	param: &[Wire],
+	message: &[Wire],
+	nonce: &[Wire],
+	signature_hashes: &[[Wire; 4]],
+	public_key_hashes: &[[Wire; 4]],
+	spec: &WinternitzSpec,
+) -> (Keccak, Vec<Keccak>) {
+	assert_eq!(param.len(), 4, "param must be 32 bytes as 4 wires");
+	assert_eq!(message.len(), 4, "message must be 32 bytes as 4 wires");
+
+	// Step 1: Hash the message with tweaking (param || TWEAK_MESSAGE || nonce || message)
+	let message_hash_output: [Wire; 4] = std::array::from_fn(|_| builder.add_witness());
+	// Note: nonce is 23 bytes, but packed into 3 wires (24 bytes)
+	let nonce_len = 23; // Actual nonce length in bytes
+	let message_len = 32; // Message is 32 bytes (4 wires * 8 bytes)
+	let message_hasher = verify_message_tweak(
+		builder,
+		param.to_vec(),
+		param.len() * 8,
+		nonce.to_vec(),
+		nonce_len,
+		message.to_vec(),
+		message_len,
+		message_hash_output,
+	);
+
+	// Step 2: Extract codeword from message hash
+	// Only use the first spec.message_hash_len bytes
+	let message_hash_bytes = spec.message_hash_len;
+	let message_hash_wires_needed = message_hash_bytes.div_ceil(8);
+	let message_hash_for_codeword = &message_hash_output[..message_hash_wires_needed];
+
+	let coordinates = codeword(
+		builder,
+		spec.dimension,
+		spec.coordinate_resolution_bits,
+		spec.target_sum,
+		message_hash_for_codeword,
+	);
+
+	assert_eq!(coordinates.len(), spec.dimension, "Codeword dimension mismatch");
+	assert_eq!(signature_hashes.len(), spec.dimension, "Signature hashes count mismatch");
+	assert_eq!(public_key_hashes.len(), spec.dimension, "Public key hashes count mismatch");
+
+	// Step 3: Verify hash chains for each coordinate
+	let mut all_chain_hashers = Vec::new();
+	let max_chain_len = spec.chain_len();
+
+	for chain_index in 0..spec.dimension {
+		let chain_index_wire = builder.add_constant_64(chain_index as u64);
+
+		// For each chain, verify that hashing from signature_hash for `coordinate` steps
+		// produces the public_key_hash
+		let chain_hashers = verify_chain(
+			builder,
+			param,
+			param.len() * 8, // param_len in bytes
+			chain_index_wire,
+			signature_hashes[chain_index],
+			coordinates[chain_index],
+			max_chain_len as u64,
+			public_key_hashes[chain_index],
+		);
+
+		all_chain_hashers.extend(chain_hashers);
+	}
+
+	(message_hasher, all_chain_hashers)
+}
+
+/// Specification for Winternitz OTS parameters
+pub struct WinternitzSpec {
+	/// Number of bytes from message hash to use
+	pub message_hash_len: usize,
+	/// Number of bits per coordinate in the codeword
+	pub coordinate_resolution_bits: usize,
+	/// Number of coordinates/chains
+	pub dimension: usize,
+	/// Expected sum of all coordinates
+	pub target_sum: u64,
+}
+
+impl WinternitzSpec {
+	/// Returns the chain length (2^coordinate_resolution_bits)
+	pub fn chain_len(&self) -> usize {
+		1 << self.coordinate_resolution_bits
+	}
+
+	/// Create a spec matching SPEC_1 from leansig-xmss
+	pub fn spec_1() -> Self {
+		Self {
+			message_hash_len: 18,
+			coordinate_resolution_bits: 2,
+			dimension: 72, // 18 * 8 / 2
+			target_sum: 119,
+		}
+	}
+
+	/// Create a spec matching SPEC_2 from leansig-xmss
+	pub fn spec_2() -> Self {
+		Self {
+			message_hash_len: 18,
+			coordinate_resolution_bits: 4,
+			dimension: 36, // 18 * 8 / 4
+			target_sum: 297,
+		}
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use rand::{RngCore, SeedableRng, rngs::StdRng};
+	use sha3::{Digest, Keccak256};
+
+	use super::{
+		super::tweak::{build_chain_tweak, build_message_tweak},
+		*,
+	};
+	use crate::{constraint_verifier::verify_constraints, util::pack_bytes_into_wires_le};
+
+	fn hash_message_keccak(param: &[u8], nonce: &[u8], message: &[u8]) -> [u8; 32] {
+		let mut hasher = Keccak256::new();
+		hasher.update(param);
+		hasher.update([0x02]); // TWEAK_MESSAGE
+		hasher.update(nonce);
+		hasher.update(message);
+		hasher.finalize().into()
+	}
+
+	fn hash_chain_keccak(
+		param: &[u8],
+		chain_index: usize,
+		start_hash: &[u8; 32],
+		start_pos: usize,
+		num_hashes: usize,
+	) -> [u8; 32] {
+		let mut current = *start_hash;
+		for i in 0..num_hashes {
+			let position = start_pos + i + 1;
+			let mut hasher = Keccak256::new();
+			hasher.update(param);
+			hasher.update([0x00]); // TWEAK_CHAIN
+			hasher.update(current);
+			hasher.update(chain_index.to_le_bytes());
+			hasher.update(position.to_le_bytes());
+			current = hasher.finalize().into();
+		}
+		current
+	}
+
+	fn extract_coordinates(hash: &[u8], dimension: usize, resolution_bits: usize) -> Vec<u8> {
+		let mut coords = Vec::new();
+		let coords_per_byte = 8 / resolution_bits;
+		let mask = (1u8 << resolution_bits) - 1;
+
+		for i in 0..dimension {
+			let byte_idx = i / coords_per_byte;
+			let coord_idx = i % coords_per_byte;
+			let shift = coord_idx * resolution_bits;
+			let coord = (hash[byte_idx] >> shift) & mask;
+			coords.push(coord);
+		}
+		coords
+	}
+
+	#[test]
+	fn test_winternitz_ots_verification() {
+		let spec = WinternitzSpec::spec_1();
+		let builder = CircuitBuilder::new();
+
+		// Create input wires
+		let param: Vec<Wire> = (0..4).map(|_| builder.add_inout()).collect();
+		let message: Vec<Wire> = (0..4).map(|_| builder.add_inout()).collect();
+		let nonce: Vec<Wire> = (0..3).map(|_| builder.add_inout()).collect(); // 23 bytes = 3*8 - 1
+
+		let signature_hashes: Vec<[Wire; 4]> = (0..spec.dimension)
+			.map(|_| std::array::from_fn(|_| builder.add_inout()))
+			.collect();
+
+		let public_key_hashes: Vec<[Wire; 4]> = (0..spec.dimension)
+			.map(|_| std::array::from_fn(|_| builder.add_inout()))
+			.collect();
+
+		// Create the verification circuit
+		let (message_hasher, chain_hashers) = verify_winternitz_ots(
+			&builder,
+			&param,
+			&message,
+			&nonce,
+			&signature_hashes,
+			&public_key_hashes,
+			&spec,
+		);
+
+		let circuit = builder.build();
+		let mut w = circuit.new_witness_filler();
+
+		// Generate test data
+		let mut rng = StdRng::seed_from_u64(42);
+
+		// Set up param (32 bytes)
+		let mut param_bytes = [0u8; 32];
+		rng.fill_bytes(&mut param_bytes);
+		pack_bytes_into_wires_le(&mut w, &param, &param_bytes);
+
+		// Set up message (32 bytes)
+		let mut message_bytes = [0u8; 32];
+		rng.fill_bytes(&mut message_bytes);
+		pack_bytes_into_wires_le(&mut w, &message, &message_bytes);
+
+		// Grind for a valid nonce that produces the target sum
+		let mut nonce_bytes = [0u8; 23];
+		let mut coords = Vec::new();
+		let mut message_hash = [0u8; 32];
+
+		// Try up to 1000 times to find a valid nonce
+		let mut found = false;
+		for _ in 0..1000 {
+			rng.fill_bytes(&mut nonce_bytes);
+			message_hash = hash_message_keccak(&param_bytes, &nonce_bytes, &message_bytes);
+
+			coords = extract_coordinates(
+				&message_hash[..spec.message_hash_len],
+				spec.dimension,
+				spec.coordinate_resolution_bits,
+			);
+
+			let coord_sum: usize = coords.iter().map(|&c| c as usize).sum();
+			if coord_sum == spec.target_sum as usize {
+				found = true;
+				break;
+			}
+		}
+
+		assert!(found, "Failed to find valid nonce after 1000 attempts");
+
+		// Pack nonce into wires (24 bytes, last byte is 0)
+		let mut nonce_padded = [0u8; 24];
+		nonce_padded[..23].copy_from_slice(&nonce_bytes);
+		pack_bytes_into_wires_le(&mut w, &nonce, &nonce_padded);
+
+		// Generate signature and public key hashes
+		let mut sig_hashes = Vec::new();
+		let mut pk_hashes = Vec::new();
+
+		for (chain_idx, &coord) in coords.iter().enumerate() {
+			// Generate random signature hash
+			let mut sig_hash = [0u8; 32];
+			rng.fill_bytes(&mut sig_hash);
+			sig_hashes.push(sig_hash);
+
+			// Compute public key hash by hashing forward
+			let pk_hash = hash_chain_keccak(
+				&param_bytes,
+				chain_idx,
+				&sig_hash,
+				coord as usize,
+				spec.chain_len() - 1 - coord as usize,
+			);
+			pk_hashes.push(pk_hash);
+
+			// Pack into wires
+			pack_bytes_into_wires_le(&mut w, &signature_hashes[chain_idx], &sig_hash);
+			pack_bytes_into_wires_le(&mut w, &public_key_hashes[chain_idx], &pk_hash);
+		}
+
+		// Populate message hasher
+		// Note: param and message are already populated above
+		// Populate nonce (23 bytes into 3 wires)
+		let mut padded_nonce = [0u8; 24];
+		padded_nonce[..23].copy_from_slice(&nonce_bytes);
+		pack_bytes_into_wires_le(&mut w, &nonce, &padded_nonce);
+
+		let tweaked_message = build_message_tweak(&param_bytes, &nonce_bytes, &message_bytes);
+		message_hasher.populate_message(&mut w, &tweaked_message);
+		// Populate the digest (only first 32 bytes needed for SHA3-256)
+		message_hasher.populate_digest(&mut w, message_hash);
+
+		// Populate chain hashers
+		let mut hasher_idx = 0;
+		for (chain_idx, &coord) in coords.iter().enumerate() {
+			let mut current_hash = sig_hashes[chain_idx];
+
+			for step in 0..spec.chain_len() {
+				let position = step + coord as usize;
+				let position_plus_one = position + 1;
+
+				// Compute next hash
+				let next_hash =
+					hash_chain_keccak(&param_bytes, chain_idx, &current_hash, position, 1);
+
+				// Populate the Keccak hasher
+				let keccak = &chain_hashers[hasher_idx];
+
+				let chain_message = build_chain_tweak(
+					&param_bytes,
+					&current_hash,
+					chain_idx as u64,
+					position_plus_one as u64,
+				);
+				keccak.populate_message(&mut w, &chain_message);
+				// Populate the digest
+				keccak.populate_digest(&mut w, next_hash);
+
+				// Update current hash if used
+				if position_plus_one < spec.chain_len() {
+					current_hash = next_hash;
+				}
+
+				hasher_idx += 1;
+			}
+		}
+
+		// Populate witness and verify
+		circuit.populate_wire_witness(&mut w).unwrap();
+
+		let cs = circuit.constraint_system();
+		verify_constraints(cs, &w.into_value_vec()).unwrap();
+	}
+}