1extern crate alloc;
9
10use alloc::string::ToString;
11use alloc::vec::Vec;
12
13use lib_q_core::Result;
14
15use crate::{
16 ZkpProof,
17 ZkpProver,
18 ZkpVerifier,
19};
20
21#[derive(Debug, Clone)]
23pub struct MerklePath {
24 pub path_bits: Vec<bool>,
26 pub siblings: Vec<crate::air::merkle_inclusion::MerkleHash>,
28}
29
30pub fn build_merkle_tree(leaves: &[&[u8]]) -> Result<crate::merkle::PoseidonMerkleTree> {
40 crate::merkle::PoseidonMerkleTree::from_leaves(leaves)
41}
42
43pub fn merkle_path_from_tree(
47 tree: &crate::merkle::PoseidonMerkleTree,
48 leaf_index: usize,
49) -> Result<MerklePath> {
50 let (path_bits, siblings) = tree.path(leaf_index)?;
51 Ok(MerklePath {
52 path_bits,
53 siblings,
54 })
55}
56
57pub fn prove_membership(leaf: &[u8], path: &MerklePath) -> Result<ZkpProof> {
94 prove_membership_with_config(leaf, path, crate::stark::default_config())
95}
96
97pub fn prove_membership_with_config(
103 leaf: &[u8],
104 path: &MerklePath,
105 config: crate::stark::DefaultConfig,
106) -> Result<ZkpProof> {
107 use crate::ProofMetadata;
108 use crate::air::{
109 MerkleInclusionAir,
110 MerkleProofInput,
111 TraceGenerator,
112 };
113 use crate::stark::StarkProver;
114
115 if path.path_bits.len() > 64 {
116 return Err(lib_q_core::Error::InvalidState {
117 operation: "prove_membership_with_config".into(),
118 reason: "Tree depth exceeds maximum of 64".into(),
119 });
120 }
121
122 let tree_depth = path.path_bits.len();
123 let air =
124 MerkleInclusionAir::new(tree_depth).map_err(|e| lib_q_core::Error::InternalError {
125 operation: "prove_membership_with_config".into(),
126 details: e.to_string(),
127 })?;
128
129 let input = MerkleProofInput {
130 leaf: leaf.to_vec(),
131 leaf_hash_direct: None,
132 path_bits: path.path_bits.clone(),
133 siblings: path.siblings.clone(),
134 };
135
136 let trace = air
137 .generate_trace(&input)
138 .map_err(|e| lib_q_core::Error::InternalError {
139 operation: "prove_membership_with_config".into(),
140 details: e.to_string(),
141 })?;
142
143 let public_values = air.public_values(&input);
144
145 let prover = StarkProver::new(config);
146 let proof = prover.prove(&air, trace, &public_values).map_err(|e| {
147 lib_q_core::Error::InternalError {
148 operation: "STARK proof generation".to_string(),
149 details: e.to_string(),
150 }
151 })?;
152
153 let metadata = ProofMetadata::MerkleInclusion {
155 tree_depth: tree_depth as u8,
156 };
157
158 ZkpProof::from_stark_proof(&proof, metadata)
159}
160
161pub fn verify_membership_with_depth(
191 proof: &ZkpProof,
192 root: &[u8],
193 expected_tree_depth: usize,
194) -> Result<bool> {
195 verify_membership_with_depth_and_config(
196 proof,
197 root,
198 expected_tree_depth,
199 crate::stark::default_config(),
200 )
201}
202
203pub fn verify_membership_with_depth_and_config(
206 proof: &ZkpProof,
207 root: &[u8],
208 expected_tree_depth: usize,
209 config: crate::stark::DefaultConfig,
210) -> Result<bool> {
211 use crate::ProofMetadata;
212 use crate::air::{
213 MerkleInclusionAir,
214 poseidon_slice_to_field,
215 };
216 use crate::stark::StarkVerifier;
217
218 if proof.proof_type != crate::ProofType::Stark {
219 return Ok(false);
220 }
221
222 if proof.data.is_empty() {
223 return Ok(false);
224 }
225
226 if let ProofMetadata::MerkleInclusion { tree_depth } = &proof.metadata &&
228 *tree_depth as usize != expected_tree_depth
229 {
230 return Ok(false);
232 }
233
234 if expected_tree_depth == 0 || expected_tree_depth > 64 {
236 return Err(lib_q_core::Error::InvalidState {
237 operation: "verify_membership_with_depth_and_config".into(),
238 reason: "Tree depth must be between 1 and 64".into(),
239 });
240 }
241
242 let air = MerkleInclusionAir::new(expected_tree_depth).map_err(|e| {
244 lib_q_core::Error::InternalError {
245 operation: "verify_membership_with_depth_and_config".into(),
246 details: e.to_string(),
247 }
248 })?;
249
250 let root_poseidon =
252 crate::air::merkle_root_from_bytes(root).map_err(|e| lib_q_core::Error::InternalError {
253 operation: "verify_membership_with_depth_and_config".into(),
254 details: e.to_string(),
255 })?;
256 let expected_public_values = poseidon_slice_to_field(&[root_poseidon]);
257
258 let stark_proof = proof.to_stark_proof()?;
260 let verifier = StarkVerifier::new(config);
261
262 match verifier.verify(&air, &stark_proof, &expected_public_values) {
263 Ok(()) => Ok(true),
264 Err(_) => Ok(false),
265 }
266}
267
268pub fn verify_membership(proof: &ZkpProof, root: &[u8]) -> Result<bool> {
298 verify_membership_with_config(proof, root, crate::stark::default_config())
299}
300
301pub fn verify_membership_with_config(
304 proof: &ZkpProof,
305 root: &[u8],
306 config: crate::stark::DefaultConfig,
307) -> Result<bool> {
308 use crate::ProofMetadata;
309 use crate::air::{
310 MerkleInclusionAir,
311 poseidon_slice_to_field,
312 };
313 use crate::stark::StarkVerifier;
314
315 if proof.proof_type != crate::ProofType::Stark {
316 return Ok(false);
317 }
318
319 if proof.data.is_empty() {
320 return Ok(false);
321 }
322
323 let ProofMetadata::MerkleInclusion { tree_depth } = &proof.metadata else {
325 return Ok(false);
327 };
328
329 let depth = *tree_depth as usize;
330
331 if depth == 0 || depth > 64 {
333 return Ok(false);
334 }
335
336 let root_poseidon = match crate::air::merkle_root_from_bytes(root) {
338 Ok(r) => r,
339 Err(_) => return Ok(false),
340 };
341 let expected_public_values = poseidon_slice_to_field(&[root_poseidon]);
342
343 let air = MerkleInclusionAir::new(depth).map_err(|e| lib_q_core::Error::InternalError {
345 operation: "verify_membership_with_config".into(),
346 details: e.to_string(),
347 })?;
348
349 let stark_proof = proof.to_stark_proof()?;
351 let verifier = StarkVerifier::new(config);
352
353 match verifier.verify(&air, &stark_proof, &expected_public_values) {
354 Ok(()) => Ok(true),
355 Err(_) => Ok(false),
356 }
357}
358
359pub fn prove_preimage(secret: &[u8]) -> Result<ZkpProof> {
381 let mut prover = ZkpProver::new();
382 let public_statement = b""; prover.prove_secret_value(secret, public_statement)
384}
385
386pub fn verify_preimage(proof: &ZkpProof, expected_hash: &[u8]) -> Result<bool> {
399 let verifier = ZkpVerifier::new();
400 verifier.verify_secret_value(proof, expected_hash)
401}
402
403pub fn prove_preimage_nist(secret: &[u8]) -> Result<ZkpProof> {
416 let mut prover = ZkpProver::new();
417 prover.prove_secret_value_nist(secret, b"")
418}
419
420pub fn verify_preimage_nist(proof: &ZkpProof, expected_hash: &[u8]) -> Result<bool> {
424 let verifier = ZkpVerifier::new();
425 verifier.verify_secret_value_nist(proof, expected_hash)
426}
427
428#[cfg(test)]
429mod tests {
430 extern crate alloc;
431
432 use alloc::vec;
433
434 use super::*;
435 use crate::{
436 ProofMetadata,
437 ProofType,
438 ZkpProof,
439 };
440
441 fn empty_stark_proof_with_metadata(metadata: ProofMetadata) -> ZkpProof {
442 ZkpProof {
443 data: vec![],
444 proof_type: ProofType::Stark,
445 security_level: 1,
446 metadata,
447 }
448 }
449
450 #[test]
451 fn test_build_merkle_tree_rejects_empty() {
452 let result = build_merkle_tree(&[]);
453 assert!(result.is_err());
454 }
455
456 #[test]
457 fn test_merkle_path_from_tree_rejects_out_of_bounds() {
458 let tree = build_merkle_tree(&[b"a".as_slice(), b"b".as_slice()]).expect("tree");
459 let result = merkle_path_from_tree(&tree, 2);
460 assert!(result.is_err());
461 }
462
463 #[test]
464 fn test_verify_membership_rejects_missing_or_empty_proof_data() {
465 let proof_missing_metadata = empty_stark_proof_with_metadata(ProofMetadata::None);
466 assert!(!verify_membership(&proof_missing_metadata, &[0u8; 32]).unwrap());
467
468 let proof_with_metadata =
469 empty_stark_proof_with_metadata(ProofMetadata::MerkleInclusion { tree_depth: 4 });
470 assert!(!verify_membership(&proof_with_metadata, &[0u8; 32]).unwrap());
471 }
472
473 #[test]
474 fn test_verify_membership_with_depth_validates_depth_inputs() {
475 let mut proof =
476 empty_stark_proof_with_metadata(ProofMetadata::MerkleInclusion { tree_depth: 4 });
477 proof.data = vec![1u8; 8];
478
479 assert!(!verify_membership_with_depth(&proof, &[0u8; 32], 3).unwrap());
480
481 let mut zero_depth_proof =
482 empty_stark_proof_with_metadata(ProofMetadata::MerkleInclusion { tree_depth: 0 });
483 zero_depth_proof.data = vec![1u8; 8];
484 assert!(verify_membership_with_depth(&zero_depth_proof, &[0u8; 32], 0).is_err());
485
486 let mut too_deep_proof =
487 empty_stark_proof_with_metadata(ProofMetadata::MerkleInclusion { tree_depth: 65 });
488 too_deep_proof.data = vec![1u8; 8];
489 assert!(verify_membership_with_depth(&too_deep_proof, &[0u8; 32], 65).is_err());
490 }
491
492 #[test]
493 fn test_verify_membership_rejects_invalid_depth_and_root_encoding() {
494 let mut zero_depth_proof =
495 empty_stark_proof_with_metadata(ProofMetadata::MerkleInclusion { tree_depth: 0 });
496 zero_depth_proof.data = vec![1u8; 8];
497 assert!(!verify_membership(&zero_depth_proof, &[0u8; 32]).unwrap());
498
499 let mut bad_root_proof =
500 empty_stark_proof_with_metadata(ProofMetadata::MerkleInclusion { tree_depth: 4 });
501 bad_root_proof.data = vec![1u8; 8];
502 assert!(!verify_membership(&bad_root_proof, &[1u8; 4]).unwrap());
503 }
504
505 #[test]
506 fn test_prove_membership_rejects_depth_over_64() {
507 let path = MerklePath {
508 path_bits: vec![false; 65],
509 siblings: vec![crate::air::MerkleHash::from_bytes(&[0u8; 32]).unwrap(); 65],
510 };
511 let result = prove_membership(b"leaf", &path);
512 assert!(result.is_err());
513 }
514
515 #[test]
516 fn test_verify_membership_rejects_empty_stark_data() {
517 let proof = ZkpProof {
518 data: vec![],
519 proof_type: ProofType::Stark,
520 security_level: 1,
521 metadata: ProofMetadata::None,
522 };
523 assert!(!verify_membership(&proof, &[0u8; 32]).unwrap());
524 assert!(!verify_membership_with_depth(&proof, &[0u8; 32], 4).unwrap());
525 }
526
527 #[test]
528 fn test_preimage_api_roundtrip_and_mismatch() {
529 let secret = b"api-preimage-secret";
530 let proof = prove_preimage(secret).expect("proof");
531
532 let ok = verify_preimage(&proof, secret).expect("verify");
533 assert!(ok);
534
535 let bad = verify_preimage(&proof, b"wrong-secret").expect("verify wrong");
536 assert!(!bad);
537 }
538
539 #[test]
540 fn test_preimage_nist_api_roundtrip_and_mismatch() {
541 use digest::{
542 ExtendableOutput,
543 Update,
544 };
545 use lib_q_sha3::CShake256;
546
547 use crate::air::hash_preimage_nist::CSHAKE_DOMAIN;
548
549 let secret = b"api-preimage-secret-nist";
550 let proof = prove_preimage_nist(secret).expect("proof");
551 let mut expected_hash = [0u8; 32];
552 {
553 let mut hasher = CShake256::new_with_function_name(&[], CSHAKE_DOMAIN);
554 hasher.update(secret);
555 hasher.finalize_xof_into(&mut expected_hash);
556 }
557
558 let ok = verify_preimage_nist(&proof, &expected_hash).expect("verify");
559 assert!(ok);
560
561 let bad = verify_preimage_nist(&proof, &[0u8; 32]).expect("verify wrong");
562 assert!(!bad);
563 }
564}