qfe 0.3.0

// src/zkp/mod.rs
//! This module contains experimental implementations related to Zero-Knowledge Proofs
//! using the QFE simulation framework, focusing on a simple validity proof scheme.
//!
//! It provides structures and methods for:
//! - Establishing a shared context (`Sqs`) based on public proof parameters.
//! - Generating challenges and validity proofs (`ZkpChallenge`, `ZkpValidityResponse`).
//! - Verifying the validity proofs within the shared context.
//!
//! **Note:** This implementation is a simulation for conceptual exploration and is
//! **not** cryptographically secure for production use without formal analysis.

// Import necessary items from the parent module (src/lib.rs) or crate root
use crate::{Frame, Sqs, QfeError, PatternType, Sha512Hash};
use crate::{PHI, RESONANCE_FREQ};
use std::hash::{Hash, Hasher};
use std::collections::hash_map::DefaultHasher;
use sha2::{Sha512, Digest};
use rand::RngCore;

// --- ZKP Struct Definitions ---

/// Represents the Verifier's challenge in the ZKP protocol.
///
/// Contains random data generated by the verifier to ensure freshness
/// and prevent replay attacks in the proof generation.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ZkpChallenge {
    /// The challenge data, typically random bytes. The length can vary.
    pub value: Vec<u8>,
}

/// Represents the Prover's response in the simple validity ZKP.
///
/// Contains a single hash derived from the challenge, the SQS context,
/// and the result of the Prover checking their witness against the public statement.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ZkpValidityResponse {
    /// A hash proving the witness validity check was performed correctly relative
    /// to the challenge and SQS context. Calculated as:
    /// `Sha512(DomainSep || Challenge || ValidityBool || SQSContext || Constants)`
    pub validity_proof_hash: Sha512Hash, // [u8; 64]
}


// --- ZKP SQS Establishment ---

/// Establishes a shared ZKP context (SQS_ZKP) based purely on public information.
///
/// Both Prover and Verifier compute this independently using the same inputs
/// to arrive at the identical shared state (`Sqs`) needed for the ZKP interaction rounds.
/// This process does not involve message exchange for the SQS itself and is transparent.
///
/// # Arguments
/// * `prover_id`: Identifier (`&str`) for the Prover Frame.
/// * `verifier_id`: Identifier (`&str`) for the Verifier Frame.
/// * `public_statement`: Byte representation (`&[u8]`) of the public statement being proven
///   (e.g., the target hash `H_public` for a hash preimage proof).
/// * `context_string`: A domain separation string (`&str`) unique to this specific proof
///   instance or protocol version to prevent cross-context attacks.
///
/// # Returns
/// * `Ok(Sqs)` containing the derived shared state (`Sqs` struct).
/// * `Err(QfeError::InternalError)` if the derived SQS components have an unexpected length.
pub fn establish_zkp_sqs(
    prover_id: &str,
    verifier_id: &str,
    public_statement: &[u8],
    context_string: &str,
) -> Result<Sqs, QfeError> {
    // 1. Derive SQS components using SHA-512
    let mut components_hasher = Sha512::new();
    components_hasher.update(b"QFE_ZKP_SQS_COMPONENTS_V1");
    components_hasher.update(prover_id.as_bytes());
    components_hasher.update(verifier_id.as_bytes());
    components_hasher.update(public_statement);
    components_hasher.update(context_string.as_bytes());
    components_hasher.update(PHI.to_le_bytes());
    components_hasher.update(RESONANCE_FREQ.to_le_bytes());
    let sqs_components: Vec<u8> = components_hasher.finalize().to_vec();

    // 2. Derive shared phase lock using DefaultHasher
    let mut phase_hasher = DefaultHasher::new();
    b"QFE_ZKP_SQS_PHASE_V1".hash(&mut phase_hasher);
    prover_id.as_bytes().hash(&mut phase_hasher);
    verifier_id.as_bytes().hash(&mut phase_hasher);
    public_statement.hash(&mut phase_hasher);
    context_string.as_bytes().hash(&mut phase_hasher);
    PHI.to_bits().hash(&mut phase_hasher);
    RESONANCE_FREQ.to_bits().hash(&mut phase_hasher);
    let phase_hash_output = phase_hasher.finish();
    let shared_phase_lock = (phase_hash_output as f64 / u64::MAX as f64) * 2.0 * std::f64::consts::PI;

    // 3. Construct the Sqs object
    let sqs = Sqs {
        pattern_type: PatternType::Sqs,
        components: sqs_components,
        shared_phase_lock,
        resonance_freq: RESONANCE_FREQ,
        validation: true, // Derived directly, assume valid structure
    };
    if sqs.components.len() != 64 {
         return Err(QfeError::InternalError(format!(
            "Derived ZKP SQS components have unexpected length: {}", sqs.components.len()
         )));
    }
    Ok(sqs)
}

/// Generates a random challenge for a ZKP round.
///
/// Uses a cryptographically secure random number generator (`rand::thread_rng`).
///
/// # Arguments
/// * `challenge_len`: The desired length of the challenge data in bytes (e.g., 32 or 64).
///
/// # Returns
/// * `ZkpChallenge` containing the randomly generated byte vector.
pub fn generate_zkp_challenge(challenge_len: usize) -> ZkpChallenge {
    let mut challenge_value = vec![0u8; challenge_len];
    rand::rng().fill_bytes(&mut challenge_value);
    ZkpChallenge { value: challenge_value }
}

// --- ZKP methods within Frame ---
impl Frame {

    /// Stores witness data within the Frame for ZKP operations.
    /// Overwrites any previously stored witness. This data is considered secret.
    ///
    /// # Arguments
    /// * `witness`: The secret witness data as a byte slice.
    ///
    /// # Errors
    /// * `QfeError::FrameInvalid` if the frame is already in an invalid state.
    pub fn store_zkp_witness(&mut self, witness: &[u8]) -> Result<(), QfeError> {
        if !self.is_valid() { return Err(QfeError::FrameInvalid); }
        self.zkp_witness = Some(witness.to_vec());
        Ok(())
    }

    // --- Simple Validity ZKP Prover Method ---

    /// Prover: Generates a simple validity proof hash based on witness, challenge, and SQS.
    ///
    /// This method implements the Prover's role in the simple validity ZKP scheme.
    /// It retrieves the stored witness, calculates `H(W)`, checks if it matches the
    /// provided `public_statement_h_public`, and then computes a hash incorporating
    /// this boolean validity result, the challenge, and the shared SQS context.
    ///
    /// # Arguments
    /// * `challenge`: The `ZkpChallenge` received from the verifier for this round.
    /// * `zkp_sqs`: The shared ZKP context established via `establish_zkp_sqs`.
    /// * `public_statement_h_public`: The public statement `H = Hash(W)` being proven.
    ///
    /// # Returns
    /// * `Ok(ZkpValidityResponse)` containing the resulting validity proof hash.
    /// * `Err(QfeError)` if:
    ///     - The frame is invalid (`QfeError::FrameInvalid`).
    ///     - The provided SQS is invalid (`QfeError::InternalError`).
    ///     - The witness has not been set (`QfeError::InternalError`).
    pub fn generate_validity_proof(
        &self,
        challenge: &ZkpChallenge,
        zkp_sqs: &Sqs,
        public_statement_h_public: &[u8],
    ) -> Result<ZkpValidityResponse, QfeError> {
        if !self.is_valid() { return Err(QfeError::FrameInvalid); }
        if !zkp_sqs.validation { return Err(QfeError::InternalError("Invalid ZKP SQS provided for proof".to_string())); }
        let witness_w = self.zkp_witness.as_ref().ok_or_else(|| QfeError::InternalError("ZKP witness not set for proof generation".to_string()))?;

        // 1. Calculate H(W)
        let calculated_hash_of_w: Sha512Hash = Sha512::digest(witness_w).into();

        // 2. Determine validity
        let is_valid_witness: bool = calculated_hash_of_w.as_slice() == public_statement_h_public;

        // 3. Compute the response hash: Hash(DomainSep || Challenge || Validity || SQS Context || Constants)
        let mut response_hasher = Sha512::new();
        response_hasher.update(b"QFE_ZKP_VALIDITY_PROOF_V1"); // Domain separation
        response_hasher.update(&challenge.value);
        response_hasher.update([is_valid_witness as u8]); // Hash the boolean result (as 1 or 0)
        response_hasher.update(&zkp_sqs.components);
        response_hasher.update(zkp_sqs.shared_phase_lock.to_le_bytes());
        response_hasher.update(PHI.to_le_bytes());
        response_hasher.update(RESONANCE_FREQ.to_le_bytes());

        let proof_hash: Sha512Hash = response_hasher.finalize().into();

        Ok(ZkpValidityResponse { validity_proof_hash: proof_hash })
    }

    // --- Simple Validity ZKP Verifier Method ---

     /// Verifier: Verifies the simple validity proof hash.
     ///
     /// This method implements the Verifier's role in the simple validity ZKP scheme.
     /// It recomputes the expected proof hash by performing the same hashing operation
     /// as the prover, but crucially *assuming* the witness was valid (`is_valid = true`).
     /// It then compares this expected hash to the hash received in the `response`.
     /// If the hashes do not match, the frame's `validation_status` is set to `false`.
     ///
     /// # Arguments
     /// * `challenge`: The `ZkpChallenge` sent by this Verifier for this round.
     /// * `response`: The `ZkpValidityResponse` received from the Prover.
     /// * `zkp_sqs`: The shared ZKP context established via `establish_zkp_sqs`.
     /// * `public_statement_h_public`: The public statement `H = Hash(W)`. Although not
     ///   directly used in the hash calculation here (it's implicitly part of the `zkp_sqs` context),
     ///   it's kept for clarity and potential future scheme refinements.
     ///
     /// # Returns
     /// * `Ok(())` if the verification check passes.
     /// * `Err(QfeError::DecodingFailed)` if the validity proof check fails.
     /// * `Err(QfeError::FrameInvalid)` if the Verifier frame is already invalid.
     /// * `Err(QfeError::InternalError)` if the provided SQS is invalid.
     pub fn verify_validity_proof(
         &mut self, // Mutable to update validation status
         challenge: &ZkpChallenge,
         response: &ZkpValidityResponse,
         zkp_sqs: &Sqs,
         // Keep H_public for context, even if not directly hashed inside this function
         _public_statement_h_public: &[u8],
     ) -> Result<(), QfeError> {
        if !self.is_valid() { return Err(QfeError::FrameInvalid); }
        if !zkp_sqs.validation { return Err(QfeError::InternalError("Invalid ZKP SQS provided for verification".to_string())); }

        // Calculate the hash Verifier expects if Prover's witness was valid
        let expected_hash = {
            let mut response_hasher = Sha512::new();
            response_hasher.update(b"QFE_ZKP_VALIDITY_PROOF_V1"); // Same domain separation
            response_hasher.update(&challenge.value);
            response_hasher.update([true as u8]); // Verifier *assumes* validity (true -> 1 byte)
            response_hasher.update(&zkp_sqs.components);
            response_hasher.update(zkp_sqs.shared_phase_lock.to_le_bytes());
            response_hasher.update(PHI.to_le_bytes());
            response_hasher.update(RESONANCE_FREQ.to_le_bytes());
            let hash: Sha512Hash = response_hasher.finalize().into();
            hash
        };

        // Compare expected hash with the one received from Prover
        if response.validity_proof_hash != expected_hash {
            self.validation_status = false; // Mark invalid on failure
            return Err(QfeError::DecodingFailed(
                "ZKP Validity Proof Check Failed".to_string()
            ));
        }

        // If hash matches
        Ok(())
     }

} // end impl Frame

// --- Unit Tests for Simple Validity ZKP ---
#[cfg(test)]
mod tests {
    use super::*; // Import items from outer scope (zkp module)
    use crate::{Frame}; // Import Frame from crate root
    use sha2::{Sha512, Digest}; // Import Sha512 for calculating H_public in tests

    // --- Test Setup Helper ---

    #[allow(dead_code)]
    struct SimpleZkpTestData {
        prover: Frame,
        verifier: Frame,
        zkp_sqs: Sqs,
        witness: Vec<u8>,
        h_public: Sha512Hash, // Public statement H(W)
        context: String,
    }

    /// Sets up Prover, Verifier, calculates H(W), stores W, establishes ZKP SQS for Validity Proof.
    fn setup_simple_zkp_test() -> SimpleZkpTestData {
        let mut prover = Frame::initialize("ValidityProver".to_string(), 20250404);
        let verifier = Frame::initialize("ValidityVerifier".to_string(), 20250404);
        let witness = b"a_valid_witness_for_simple_zkp".to_vec();
        let h_public: Sha512Hash = Sha512::digest(&witness).into(); // Calculate H(W)
        prover.store_zkp_witness(&witness).expect("Failed to store witness");
        let context = "simple_validity_test_v1".to_string();
        let zkp_sqs = establish_zkp_sqs(
            prover.id(),
            verifier.id(),
            &h_public, // Use the calculated public hash H(W)
            &context,
        ).expect("Failed to establish ZKP SQS");
        let zkp_sqs_clone = zkp_sqs.clone();
        SimpleZkpTestData {
            prover,
            verifier,
            zkp_sqs: zkp_sqs_clone,
            witness,
            h_public,
            context,
        }
    }

    // --- Simple Validity ZKP Tests ---

    #[test]
    fn test_simple_zkp_successful_proof() {
        let test_data = setup_simple_zkp_test();
        let prover = test_data.prover;
        let mut verifier = test_data.verifier; // Verifier needs mut for verify call
        let zkp_sqs = test_data.zkp_sqs;
        let h_public = test_data.h_public;

        // 1. Verifier generates challenge
        let challenge = generate_zkp_challenge(32); // 32-byte challenge

        // 2. Prover generates validity proof response
        let response = prover.generate_validity_proof(
            &challenge,
            &zkp_sqs,
            &h_public,
        ).expect("Prover failed to generate validity proof");

        // 3. Verifier verifies the proof
        let verification_result = verifier.verify_validity_proof(
            &challenge,
            &response,
            &zkp_sqs,
            &h_public, // Pass H_public for context/potential future use
        );

        // Assert verification success
        assert!(verification_result.is_ok(), "Verification failed unexpectedly: {:?}", verification_result.err());
        assert!(verifier.is_valid(), "Verifier should remain valid after successful verification");
    }

    #[test]
    fn test_simple_zkp_invalid_witness() {
        let test_data = setup_simple_zkp_test();
        let mut prover = test_data.prover; // Need mut to store wrong witness
        let mut verifier = test_data.verifier;
        let zkp_sqs = test_data.zkp_sqs;
        let h_public = test_data.h_public; // Correct H(W)

        // Store WRONG witness
        let wrong_witness = b"this_is_the_wrong_witness".to_vec();
        prover.store_zkp_witness(&wrong_witness).expect("Storing wrong witness failed");

        // Verifier generates challenge
        let challenge = generate_zkp_challenge(32);

        // Prover generates response using the wrong witness
        // This means the `is_valid_witness` flag inside generate_validity_proof will be false.
        let response = prover.generate_validity_proof(
            &challenge,
            &zkp_sqs,
            &h_public, // P uses correct public H for comparison context
        ).expect("Prover failed proof generation (using wrong witness)");

        // Verifier verifies response.
        // V calculates expected hash assuming `is_valid=true`.
        // P calculated hash using `is_valid=false`. Hashes won't match.
        let verification_result = verifier.verify_validity_proof(
            &challenge,
            &response,
            &zkp_sqs,
            &h_public,
        );

        // Assert failure
        assert!(verification_result.is_err(), "Verification should fail for invalid witness");
        let err = verification_result.unwrap_err();
        assert!(matches!(err, QfeError::DecodingFailed(_)), "Expected DecodingFailed, got {:?}", err);
        if let QfeError::DecodingFailed(msg) = err {
             assert!(msg.contains("Validity Proof Check Failed"), "Expected Validity Proof failure message, got: {}", msg);
        }
        assert!(!verifier.is_valid(), "Verifier should become invalid after failed verification");
    }

    #[test]
    fn test_simple_zkp_tampered_response_hash() {
        let test_data = setup_simple_zkp_test();
        let prover = test_data.prover;
        let mut verifier = test_data.verifier;
        let zkp_sqs = test_data.zkp_sqs;
        let h_public = test_data.h_public;

        let challenge = generate_zkp_challenge(32);

        // P generates a valid response first
        let mut response = prover.generate_validity_proof(&challenge, &zkp_sqs, &h_public)
            .expect("Prover failed proof generation");

        // Tamper with the validity proof hash
        response.validity_proof_hash[0] ^= 0xAA; // Flip some bits

        // V verifies tampered response
        let verification_result = verifier.verify_validity_proof(
            &challenge,
            &response, // Pass tampered response
            &zkp_sqs,
            &h_public,
        );

        // Assert failure
        assert!(verification_result.is_err(), "Verification should fail for tampered validity proof hash");
        let err = verification_result.unwrap_err();
        assert!(matches!(err, QfeError::DecodingFailed(_)), "Expected DecodingFailed, got {:?}", err);
        if let QfeError::DecodingFailed(msg) = err {
             assert!(msg.contains("Validity Proof Check Failed"), "Expected Validity Proof failure message, got: {}", msg);
        }
        assert!(!verifier.is_valid());
    }

    #[test]
    fn test_simple_zkp_wrong_challenge() {
        // Verifier uses a different challenge during verification than Prover used
        let test_data = setup_simple_zkp_test();
        let prover = test_data.prover;
        let mut verifier = test_data.verifier;
        let zkp_sqs = test_data.zkp_sqs;
        let h_public = test_data.h_public;

        // V generates C1, P uses C1
        let challenge_c1 = generate_zkp_challenge(32);
        let response = prover.generate_validity_proof(&challenge_c1, &zkp_sqs, &h_public)
            .expect("Prover failed proof generation");

        // V generates C2 and uses it for verification
        let challenge_c2 = generate_zkp_challenge(32);
        // Ensure challenges are different (highly likely, but check anyway)
        if challenge_c1 == challenge_c2 {
            println!("Warning: Collision in random challenge generation, test may be less effective.");
        }
        assert_ne!(challenge_c1, challenge_c2);


        let verification_result = verifier.verify_validity_proof(
            &challenge_c2, // Use wrong challenge C2
            &response,
            &zkp_sqs,
            &h_public,
        );

        // Assert failure
        assert!(verification_result.is_err(), "Verification should fail when Verifier uses wrong challenge");
        let err = verification_result.unwrap_err();
        // Failure happens because V's expected hash uses C2, P's used C1
        assert!(matches!(err, QfeError::DecodingFailed(_)), "Expected DecodingFailed, got {:?}", err);
         if let QfeError::DecodingFailed(msg) = err {
             assert!(msg.contains("Validity Proof Check Failed"), "Expected Validity Proof failure message, got: {}", msg);
         }
        assert!(!verifier.is_valid());
    }

    #[test]
    fn test_simple_zkp_wrong_sqs() {
        // Verifier uses different SQS context for verification
        let test_data1 = setup_simple_zkp_test(); // P, V1, SQS1, H
        let prover = test_data1.prover;
        let zkp_sqs1 = test_data1.zkp_sqs;
        let h_public = test_data1.h_public;

        // Create V2 and SQS2
        let mut verifier2 = Frame::initialize("Verifier2_WrongSQS".to_string(), 909090);
        let zkp_sqs2 = establish_zkp_sqs(
            prover.id(),
            verifier2.id(), // Different V ID
            &h_public,
            "a_completely_different_context", // Different context string
        ).expect("Failed to establish ZKP SQS2");
        assert_ne!(zkp_sqs1.components, zkp_sqs2.components); // Ensure SQS differs

        // V2 sends challenge
        let challenge = generate_zkp_challenge(32);

        // P generates response using SQS1
        let response = prover.generate_validity_proof(&challenge, &zkp_sqs1, &h_public)
            .expect("Prover failed proof generation");

        // V2 verifies using SQS2
        let verification_result = verifier2.verify_validity_proof(
            &challenge,
            &response,
            &zkp_sqs2, // Use wrong SQS2
            &h_public,
        );

        // Assert failure
        assert!(verification_result.is_err(), "Verification should fail when using wrong SQS");
        let err = verification_result.unwrap_err();
        // Failure happens because SQS context affects hash calculation
         assert!(matches!(err, QfeError::DecodingFailed(_)), "Expected DecodingFailed, got {:?}", err);
         if let QfeError::DecodingFailed(msg) = err {
             assert!(msg.contains("Validity Proof Check Failed"), "Expected Validity Proof failure message, got: {}", msg);
         }
        assert!(!verifier2.is_valid());
    }

    // Note: The test for "Wrong Public Statement" is implicitly covered by the "Wrong SQS"
    // test IF establish_zkp_sqs correctly incorporates the statement into the SQS components,
    // making the SQS context statement-specific. If SQS was *not* statement-specific,
    // we would need to add H_public back into the hash calculations inside the proof/verify methods
    // and add a specific test for verifying with the wrong H_public.

} // end tests module