[frontend] Add xmss verifier

Author: paulcadman

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

@@ -3,3 +3,4 @@ pub mod codeword;
 pub mod merkle_tree;
 pub mod tweak;
 pub mod winternitz_ots;
+pub mod xmss;

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

@@ -2,8 +2,10 @@
 mod base;
 mod chain;
 mod message;
+mod public_key;
 mod tree;
 
 pub use chain::{CHAIN_TWEAK, FIXED_MESSAGE_OVERHEAD, build_chain_tweak, verify_chain_tweak};
 pub use message::{MESSAGE_TWEAK, build_message_tweak, verify_message_tweak};
+pub use public_key::{PUBLIC_KEY_TWEAK, build_public_key_tweak, verify_public_key_tweak};
 pub use tree::{TREE_MESSAGE_OVERHEAD, TREE_TWEAK, build_tree_tweak, verify_tree_tweak};

crates/frontend/src/circuits/hash_based_sig/tweak/public_key.rs

@@ -0,0 +1,351 @@
+use super::base::verify_tweaked_keccak;
+// Note: PublicKeyTweak reuses TREE_TWEAK (0x01) for consistency with XMSS spec
+pub use super::tree::TREE_TWEAK as PUBLIC_KEY_TWEAK;
+use crate::{
+	circuits::{concat::Term, keccak::Keccak},
+	compiler::{CircuitBuilder, Wire},
+};
+
+/// A circuit that verifies a public key hash for XMSS.
+///
+/// This circuit verifies Keccak-256 of a message that's been tweaked with
+/// multiple public key hashes: `Keccak256(param || 0x01 || pk_hash_0 || pk_hash_1 || ... ||
+/// pk_hash_n)`
+///
+/// # Arguments
+///
+/// * `builder` - Circuit builder for constructing constraints
+/// * `param_wires` - The cryptographic parameter wires, where each wire holds 8 bytes as a 64-bit
+///   LE-packed value
+/// * `param_len` - The actual parameter length in bytes
+/// * `pk_hashes` - The public key end hashes (32 bytes each as 4x64-bit LE-packed wires)
+/// * `digest` - Output: The computed Keccak-256 digest (32 bytes as 4x64-bit LE-packed wires)
+///
+/// # Returns
+///
+/// A `Keccak` circuit that needs to be populated with the tweaked message and digest
+pub fn verify_public_key_tweak(
+	builder: &CircuitBuilder,
+	param_wires: Vec<Wire>,
+	param_len: usize,
+	pk_hashes: &[[Wire; 4]],
+	digest: [Wire; 4],
+) -> Keccak {
+	let num_hashes = pk_hashes.len();
+	let message_len = param_len + 1 + (num_hashes * 32); // +1 for tweak byte
+	assert_eq!(param_wires.len(), param_len.div_ceil(8));
+
+	// Build additional terms for all public key hashes
+	let mut additional_terms = Vec::new();
+
+	// Add all public key hashes
+	for pk_hash in pk_hashes {
+		let hash_term = Term {
+			len: builder.add_constant_64(32),
+			data: pk_hash.to_vec(),
+			max_len: 32,
+		};
+		additional_terms.push(hash_term);
+	}
+
+	verify_tweaked_keccak(
+		builder,
+		param_wires,
+		param_len,
+		PUBLIC_KEY_TWEAK,
+		additional_terms,
+		message_len,
+		digest,
+	)
+}
+
+/// Build the tweaked message for public key hashing.
+///
+/// Constructs the complete message for Keccak-256 hashing by concatenating:
+/// `param || 0x01 || pk_hash_0 || pk_hash_1 || ... || pk_hash_n`
+///
+/// This function is typically used when populating witness data for the
+/// `verify_public_key_tweak` circuit.
+///
+/// # Arguments
+///
+/// * `param_bytes` - The cryptographic parameter bytes
+/// * `pk_hashes` - Array of 32-byte public key hashes
+///
+/// # Returns
+///
+/// A vector containing the complete tweaked message ready for hashing
+pub fn build_public_key_tweak(param_bytes: &[u8], pk_hashes: &[[u8; 32]]) -> Vec<u8> {
+	let mut message = Vec::new();
+	message.extend_from_slice(param_bytes);
+	message.push(PUBLIC_KEY_TWEAK);
+	for pk_hash in pk_hashes {
+		message.extend_from_slice(pk_hash);
+	}
+	message
+}
+
+#[cfg(test)]
+mod tests {
+	use proptest::prelude::*;
+	use sha3::{Digest, Keccak256};
+
+	use super::*;
+	use crate::{
+		compiler::{CircuitBuilder, circuit::Circuit},
+		constraint_verifier::verify_constraints,
+		util::pack_bytes_into_wires_le,
+	};
+
+	/// Helper struct for PublicKeyTweak testing
+	struct PublicKeyTestCircuit {
+		circuit: Circuit,
+		keccak: Keccak,
+		param_wires: Vec<Wire>,
+		param_len: usize,
+		pk_hashes: Vec<[Wire; 4]>,
+	}
+
+	impl PublicKeyTestCircuit {
+		fn new(param_len: usize, num_hashes: usize) -> Self {
+			let builder = CircuitBuilder::new();
+
+			let digest: [Wire; 4] = std::array::from_fn(|_| builder.add_inout());
+
+			let num_param_wires = param_len.div_ceil(8);
+			let param_wires: Vec<Wire> =
+				(0..num_param_wires).map(|_| builder.add_inout()).collect();
+
+			let pk_hashes: Vec<[Wire; 4]> = (0..num_hashes)
+				.map(|_| std::array::from_fn(|_| builder.add_inout()))
+				.collect();
+
+			let keccak = verify_public_key_tweak(
+				&builder,
+				param_wires.clone(),
+				param_len,
+				&pk_hashes,
+				digest,
+			);
+
+			let circuit = builder.build();
+
+			Self {
+				circuit,
+				keccak,
+				param_wires,
+				param_len,
+				pk_hashes,
+			}
+		}
+
+		/// Populate witness and verify constraints with given test data
+		fn populate_and_verify(
+			&self,
+			param_bytes: &[u8],
+			pk_hashes_bytes: &[[u8; 32]],
+			message: &[u8],
+			digest: [u8; 32],
+		) -> Result<(), Box<dyn std::error::Error>> {
+			let mut w = self.circuit.new_witness_filler();
+
+			// Populate param
+			assert_eq!(param_bytes.len(), self.param_len);
+			pack_bytes_into_wires_le(&mut w, &self.param_wires, param_bytes);
+
+			// Populate public key hashes
+			assert_eq!(pk_hashes_bytes.len(), self.pk_hashes.len());
+			for (wires, bytes) in self.pk_hashes.iter().zip(pk_hashes_bytes.iter()) {
+				pack_bytes_into_wires_le(&mut w, wires, bytes);
+			}
+
+			// Populate message for Keccak
+			let expected_len = self.param_len + 1 + (self.pk_hashes.len() * 32);
+			assert_eq!(
+				message.len(),
+				expected_len,
+				"Message length {} doesn't match expected length {}",
+				message.len(),
+				expected_len
+			);
+			self.keccak.populate_message(&mut w, message);
+
+			// Populate digest
+			self.keccak.populate_digest(&mut w, digest);
+
+			self.circuit.populate_wire_witness(&mut w)?;
+			let cs = self.circuit.constraint_system();
+			verify_constraints(cs, &w.into_value_vec())?;
+			Ok(())
+		}
+	}
+
+	#[test]
+	fn test_public_key_tweak_basic() {
+		let test_circuit = PublicKeyTestCircuit::new(32, 3);
+
+		let param_bytes = b"test_parameter_32_bytes_long!!!!";
+		let pk_hashes = [
+			*b"first_public_key_hash_32_bytes!!",
+			*b"second_public_key_hash_32_bytes!",
+			*b"third_public_key_hash_32_bytes!!",
+		];
+
+		let message = build_public_key_tweak(param_bytes, &pk_hashes);
+
+		let expected_digest = Keccak256::digest(&message);
+
+		test_circuit
+			.populate_and_verify(param_bytes, &pk_hashes, &message, expected_digest.into())
+			.unwrap();
+	}
+
+	#[test]
+	fn test_public_key_tweak_with_18_byte_param() {
+		// Test with 18-byte param as per XMSS specifications
+		let test_circuit = PublicKeyTestCircuit::new(18, 2);
+
+		let param_bytes: &[u8; 18] = b"test_param_18bytes";
+		let pk_hashes = [
+			*b"first_public_key_hash_32_bytes!!",
+			*b"second_public_key_hash_32_bytes!",
+		];
+
+		let message = build_public_key_tweak(param_bytes, &pk_hashes);
+
+		let expected_digest = Keccak256::digest(&message);
+
+		test_circuit
+			.populate_and_verify(param_bytes, &pk_hashes, &message, expected_digest.into())
+			.unwrap();
+	}
+
+	#[test]
+	fn test_public_key_tweak_single_hash() {
+		let test_circuit = PublicKeyTestCircuit::new(32, 1);
+
+		let param_bytes = b"test_parameter_32_bytes_long!!!!";
+		let pk_hashes = [*b"single_public_key_hash_32_bytes!"];
+
+		let message = build_public_key_tweak(param_bytes, &pk_hashes);
+
+		let expected_digest = Keccak256::digest(&message);
+
+		test_circuit
+			.populate_and_verify(param_bytes, &pk_hashes, &message, expected_digest.into())
+			.unwrap();
+	}
+
+	#[test]
+	fn test_public_key_tweak_many_hashes() {
+		// Test with many hashes (e.g., Winternitz with 72 chains)
+		let num_hashes = 72;
+		let test_circuit = PublicKeyTestCircuit::new(18, num_hashes);
+
+		let param_bytes: &[u8; 18] = b"test_param_18bytes";
+
+		// Generate deterministic hashes for testing
+		let mut pk_hashes = Vec::new();
+		for i in 0..num_hashes {
+			let mut hash = [0u8; 32];
+			hash[0] = i as u8;
+			hash[1] = (i >> 8) as u8;
+			for j in 2..32 {
+				hash[j] = ((i * j) % 256) as u8;
+			}
+			pk_hashes.push(hash);
+		}
+
+		let message = build_public_key_tweak(param_bytes, &pk_hashes);
+
+		let expected_digest = Keccak256::digest(&message);
+
+		test_circuit
+			.populate_and_verify(param_bytes, &pk_hashes, &message, expected_digest.into())
+			.unwrap();
+	}
+
+	#[test]
+	fn test_public_key_tweak_wrong_digest() {
+		let test_circuit = PublicKeyTestCircuit::new(32, 2);
+
+		let param_bytes = b"test_parameter_32_bytes_long!!!!";
+		let pk_hashes = [
+			*b"first_public_key_hash_32_bytes!!",
+			*b"second_public_key_hash_32_bytes!",
+		];
+
+		let message = build_public_key_tweak(param_bytes, &pk_hashes);
+
+		// Populate with WRONG digest - this should cause verification to fail
+		let wrong_digest = [0u8; 32];
+
+		let result =
+			test_circuit.populate_and_verify(param_bytes, &pk_hashes, &message, wrong_digest);
+
+		assert!(result.is_err(), "Expected error for wrong digest");
+	}
+
+	#[test]
+	fn test_public_key_tweak_ensures_tweak_byte() {
+		// This test verifies that the PUBLIC_KEY_TWEAK byte (0x01) is correctly inserted
+		let test_circuit = PublicKeyTestCircuit::new(16, 1);
+
+		let param_bytes = b"param_16_bytes!!";
+		let pk_hashes = [*b"single_public_key_hash_32_bytes!"];
+
+		let message = build_public_key_tweak(param_bytes, &pk_hashes);
+
+		// Verify the tweak byte is at the correct position
+		assert_eq!(message[16], PUBLIC_KEY_TWEAK);
+		assert_eq!(message.len(), 16 + 1 + 32); // param + tweak + one hash
+
+		let expected_digest = Keccak256::digest(&message);
+
+		test_circuit
+			.populate_and_verify(param_bytes, &pk_hashes, &message, expected_digest.into())
+			.unwrap();
+	}
+
+	proptest! {
+		#[test]
+		fn test_public_key_tweak_property_based(
+			param_len in 1usize..=100,
+			num_hashes in 1usize..=10,
+		) {
+			use rand::SeedableRng;
+			use rand::prelude::StdRng;
+
+			let mut rng = StdRng::seed_from_u64(0);
+
+			// Generate random param bytes
+			let mut param_bytes = vec![0u8; param_len];
+			rng.fill_bytes(&mut param_bytes);
+
+			// Generate random public key hashes
+			let mut pk_hashes = Vec::new();
+			for _ in 0..num_hashes {
+				let mut hash = [0u8; 32];
+				rng.fill_bytes(&mut hash);
+				pk_hashes.push(hash);
+			}
+
+			// Create circuit
+			let test_circuit = PublicKeyTestCircuit::new(param_len, num_hashes);
+
+			// Build message and compute digest
+			let message = build_public_key_tweak(&param_bytes, &pk_hashes);
+
+			// Verify message structure
+			prop_assert_eq!(message.len(), param_len + 1 + (num_hashes * 32));
+			prop_assert_eq!(message[param_len], PUBLIC_KEY_TWEAK);
+
+			let expected_digest: [u8; 32] = Keccak256::digest(&message).into();
+
+			// Verify circuit
+			test_circuit
+				.populate_and_verify(&param_bytes, &pk_hashes, &message, expected_digest)
+				.unwrap();
+		}
+	}
+}