Skip to main content

lib_q_zkp/
api.rs

1//! High-level lib-Q API for zero-knowledge proofs
2//!
3//! This module provides easy-to-use functions for common ZKP operations,
4//! following lib-Q's design principles: simple functions for common problems,
5//! secure by default, and consistent naming while allowing advanced users
6//! to access lower-level APIs.
7
8extern 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/// Merkle path for inclusion proofs
22#[derive(Debug, Clone)]
23pub struct MerklePath {
24    /// Direction bits for each level (false = left, true = right)
25    pub path_bits: Vec<bool>,
26    /// Sibling hashes at each level (already computed hashes, not raw data)
27    pub siblings: Vec<crate::air::merkle_inclusion::MerkleHash>,
28}
29
30/// Build a Poseidon Merkle tree from leaf data.
31///
32/// The tree uses the same hashing as `MerkleInclusionAir`. Use `tree.path(index)`
33/// to get `(path_bits, siblings)` and construct `MerklePath { path_bits, siblings }`
34/// for `prove_membership`. Use `tree.root_bytes()` for `verify_membership`.
35///
36/// # Errors
37///
38/// Returns error if `leaves` is empty or padded size exceeds 2^64.
39pub fn build_merkle_tree(leaves: &[&[u8]]) -> Result<crate::merkle::PoseidonMerkleTree> {
40    crate::merkle::PoseidonMerkleTree::from_leaves(leaves)
41}
42
43/// Build a MerklePath for a leaf index from a built tree.
44///
45/// Convenience for `prove_membership(leaves[leaf_index], &path)`.
46pub 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
57/// Prove membership in a Merkle tree
58///
59/// This is a high-level function that proves a leaf value is included
60/// in a Merkle tree with a given root hash.
61///
62/// # Arguments
63///
64/// * `leaf` - The leaf value to prove membership of
65/// * `path` - The Merkle authentication path
66///
67/// # Returns
68///
69/// A zero-knowledge proof of membership with embedded tree depth metadata.
70/// The proof is self-describing and can be verified without knowing the
71/// tree depth in advance.
72///
73/// # Example
74///
75/// ```rust,ignore
76/// use lib_q_zkp::api::{prove_membership, MerklePath};
77///
78/// use lib_q_zkp::air::merkle_inclusion::MerkleHash;
79///
80/// let leaf = b"my leaf data";
81/// let path = MerklePath {
82///     path_bits: vec![false, true, false],
83///     siblings: vec![
84///         MerkleHash::from_bytes(&[0u8; 32])?,
85///         MerkleHash::from_bytes(&[0u8; 32])?,
86///         MerkleHash::from_bytes(&[0u8; 32])?,
87///     ],
88/// };
89///
90/// let proof = prove_membership(leaf, &path)?;
91/// // Tree depth (3) is stored in proof.metadata
92/// ```
93pub fn prove_membership(leaf: &[u8], path: &MerklePath) -> Result<ZkpProof> {
94    prove_membership_with_config(leaf, path, crate::stark::default_config())
95}
96
97/// Prove Merkle membership with an explicit STARK configuration.
98///
99/// [`prove_membership`] uses production [`crate::stark::default_config`]. For integration
100/// tests, use [`crate::stark::fast_proof_config`] so proving stays within a reasonable time.
101/// The verifier must be constructed with the **same** configuration type and FRI parameters.
102pub 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    // Store tree depth in proof metadata for self-describing verification
154    let metadata = ProofMetadata::MerkleInclusion {
155        tree_depth: tree_depth as u8,
156    };
157
158    ZkpProof::from_stark_proof(&proof, metadata)
159}
160
161/// Verify membership in a Merkle tree with explicit tree depth
162///
163/// This is the recommended verification function when the tree depth is known.
164/// It provides O(1) verification and prevents potential depth confusion attacks.
165///
166/// # Arguments
167///
168/// * `proof` - The proof to verify
169/// * `root` - The expected Merkle root hash (bytes)
170/// * `expected_tree_depth` - The expected tree depth (must match proof)
171///
172/// # Returns
173///
174/// `Ok(true)` if the proof is valid and tree depth matches, `Ok(false)` otherwise
175///
176/// # Security
177///
178/// Using explicit tree depth verification prevents depth confusion attacks where
179/// a malicious prover might craft proofs that validate at unexpected depths.
180///
181/// # Example
182///
183/// ```rust,ignore
184/// use lib_q_zkp::api::verify_membership_with_depth;
185///
186/// let root = b"expected root hash";
187/// let expected_depth = 4; // Known tree depth
188/// let is_valid = verify_membership_with_depth(&proof, root, expected_depth)?;
189/// ```
190pub 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
203/// Same as [`verify_membership_with_depth`], but uses the given STARK configuration (must
204/// match the prover configuration used to create `proof`).
205pub 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    // Validate tree depth against proof metadata if present
227    if let ProofMetadata::MerkleInclusion { tree_depth } = &proof.metadata &&
228        *tree_depth as usize != expected_tree_depth
229    {
230        // Tree depth mismatch - reject to prevent depth confusion
231        return Ok(false);
232    }
233
234    // Validate tree depth bounds
235    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    // Create AIR with the specified depth
243    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    // Deserialize root bytes to the single PoseidonField (no extra hash)
251    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    // Deserialize and verify the STARK proof
259    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
268/// Verify membership in a Merkle tree
269///
270/// This verifies a proof generated by `prove_membership`.
271///
272/// The proof must contain tree depth metadata. Proofs created with
273/// `prove_membership()` automatically include this metadata.
274///
275/// # Arguments
276///
277/// * `proof` - The proof to verify (must contain MerkleInclusion metadata)
278/// * `root` - The expected Merkle root hash (bytes)
279///
280/// # Returns
281///
282/// `Ok(true)` if the proof is valid, `Ok(false)` if invalid or missing metadata
283///
284/// # Security Note
285///
286/// For maximum security and explicit depth validation, prefer
287/// `verify_membership_with_depth` when you know the expected tree depth.
288///
289/// # Example
290///
291/// ```rust,ignore
292/// use lib_q_zkp::api::verify_membership;
293///
294/// let root = b"expected root hash";
295/// let is_valid = verify_membership(&proof, root)?;
296/// ```
297pub fn verify_membership(proof: &ZkpProof, root: &[u8]) -> Result<bool> {
298    verify_membership_with_config(proof, root, crate::stark::default_config())
299}
300
301/// Same as [`verify_membership`], but uses the given STARK configuration (must match the
302/// prover configuration used to create `proof`).
303pub 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    // Proof must contain tree depth metadata
324    let ProofMetadata::MerkleInclusion { tree_depth } = &proof.metadata else {
325        // Missing metadata - proof is invalid
326        return Ok(false);
327    };
328
329    let depth = *tree_depth as usize;
330
331    // Validate depth bounds
332    if depth == 0 || depth > 64 {
333        return Ok(false);
334    }
335
336    // Deserialize root bytes to the single PoseidonField (no extra hash)
337    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    // Create AIR with the metadata depth
344    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    // Deserialize and verify the STARK proof
350    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
359/// Prove knowledge of a preimage without revealing it
360///
361/// This generates a proof that the prover knows a secret value that hashes
362/// to a given output, without revealing the secret.
363///
364/// # Arguments
365///
366/// * `secret` - The secret preimage to prove knowledge of
367///
368/// # Returns
369///
370/// A zero-knowledge proof of preimage knowledge
371///
372/// # Example
373///
374/// ```rust,ignore
375/// use lib_q_zkp::api::prove_preimage;
376///
377/// let secret = b"my secret password";
378/// let proof = prove_preimage(secret)?;
379/// ```
380pub fn prove_preimage(secret: &[u8]) -> Result<ZkpProof> {
381    let mut prover = ZkpProver::new();
382    let public_statement = b""; // Empty for preimage proof
383    prover.prove_secret_value(secret, public_statement)
384}
385
386/// Verify a preimage proof
387///
388/// This verifies a proof generated by `prove_preimage`.
389///
390/// # Arguments
391///
392/// * `proof` - The proof to verify
393/// * `expected_hash` - The expected hash output (bytes)
394///
395/// # Returns
396///
397/// `Ok(true)` if the proof is valid, `Ok(false)` or `Err` otherwise
398pub 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
403/// Prove knowledge of a preimage using NIST cSHAKE256
404///
405/// Same as [`prove_preimage`] but uses cSHAKE256 (domain `b"HashPreimageNistAir"`).
406/// Use when 100% NIST compliance is required; prover cost is higher than Poseidon-based proofs.
407///
408/// # Arguments
409///
410/// * `secret` - The secret preimage
411///
412/// # Returns
413///
414/// A zero-knowledge proof (NIST variant)
415pub fn prove_preimage_nist(secret: &[u8]) -> Result<ZkpProof> {
416    let mut prover = ZkpProver::new();
417    prover.prove_secret_value_nist(secret, b"")
418}
419
420/// Verify a NIST (cSHAKE256) preimage proof
421///
422/// Verifies a proof from [`prove_preimage_nist`]. `expected_hash` is the 32-byte cSHAKE256 output.
423pub 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}