Zero-Knowledge Proofs in Rust & Circom: SnarkJS Verification & Production Guide

pragma circom 2.1.5;

template IdentityProof() {
    signal input x; // Private input
    signal input identity_hash; // Public input

    // Constraint: x must equal identity_hash
    x === identity_hash;
}

component main = IdentityProof();

pragma circom 2.1.5;

// Helper component to check if a number is within a range [0, N-1]
// by decomposing it into bits.
template Num2Bits(n) {
    signal input in;
    signal output out[n];

    var sum = 0;
    for (var i = 0; i < n; i++) {
        out[i] <-- (in >> i) & 1; // Extract i-th bit
        out[i] * (1 - out[i]) === 0; // Constraint: bit must be 0 or 1
        sum += out[i] * (1 << i);
    }
    sum === in; // Constraint: sum of bits must equal input
}

template RangeProof(n_bits, min_val, max_val) {
    signal input x; // Private input
    signal input public_min; // Public input (for flexibility, can be hardcoded)
    signal input public_max; // Public input (for flexibility, can be hardcoded)

    // Ensure x is non-negative (if range starts from 0)
    // For simplicity, we assume x >= 0. If negative numbers are possible,
    // a more complex range check is needed.

    // Check x >= public_min
    signal diff_min;
    diff_min <== x - public_min;
    component bits_diff_min = Num2Bits(n_bits); // n_bits should be sufficient for max_val - min_val
    bits_diff_min.in <== diff_min;

    // Check x <= public_max
    signal diff_max;
    diff_max <== public_max - x;
    component bits_diff_max = Num2Bits(n_bits); // n_bits should be sufficient for max_val - min_val
    bits_diff_max.in <== diff_max;

    // Public inputs must match the hardcoded range for this specific proof instance
    public_min === min_val;
    public_max === max_val;
}

// Example: Prove x is in range [10, 100] using 8 bits (max value 2^8 - 1 = 255)
component main = RangeProof(8, 10, 100);

# Install circom if not already installed
# npm install -g circom_tester # or yarn global add circom_tester

# Compile IdentityProof
circom identity.circom --r1cs --wasm --sym --json

# Compile RangeProof
circom range.circom --r1cs --wasm --sym --json

# Generate a new trusted setup (powers of tau)
snarkjs powersoftau new bn128 12 pot12_0000.ptau -v

# Contribute to the setup (simulated for dev)
snarkjs powersoftau contribute pot12_0000.ptau pot12_0001.ptau --name="First contribution" -v

# Apply the circuit-specific phase 2 setup for IdentityProof
snarkjs groth16 setup identity.r1cs pot12_0001.ptau identity_0000.zkey

# Contribute to the phase 2 setup (simulated for dev)
snarkjs zkey contribute identity_0000.zkey identity_final.zkey --name="Second contribution" -v

# Export verification key
snarkjs zkey export verificationkey identity_final.zkey verification_key_identity.json

# Repeat for RangeProof
snarkjs groth16 setup range.r1cs pot12_0001.ptau range_0000.zkey
snarkjs zkey contribute range_0000.zkey range_final.zkey --name="Range contribution" -v
snarkjs zkey export verificationkey range_final.zkey verification_key_range.json

// witness_generator.ts
import { WitnessCalculator } from 'circom_runtime';
import * as fs from 'fs';
import * as path from 'path';

async function generateWitness(circuitName: string, inputs: Record<string, any>): Promise<any> {
    const wasmPath = path.join(__dirname, `${circuitName}_js`, `${circuitName}.wasm`);
    const witnessCalculator = await WitnessCalculator(fs.readFileSync(wasmPath));

    const fullWitness = await witnessCalculator.calculateWitness(inputs, true);
    // The first element is the signal `1`, subsequent elements are public inputs, then private inputs.
    // We usually only need the full witness for proof generation.
    return fullWitness;
}

// Example usage for IdentityProof
async function generateIdentityWitness() {
    const privateX = 12345;
    const publicIdentityHash = 12345; // Must match privateX for a valid proof

    const inputs = {
        x: privateX,
        identity_hash: publicIdentityHash,
    };

    console.log(`Generating witness for IdentityProof with inputs: ${JSON.stringify(inputs)}`);
    const witness = await generateWitness('identity', inputs);
    console.log('IdentityProof Witness generated.');
    // In a real scenario, you might serialize this witness or pass it directly.
    // For arkworks, we need the public inputs and the full witness.
    // The arkworks prover expects a specific format, typically a vector of field elements.
    // The witness array from snarkjs is already field elements.
    return { witness, publicInputs: [publicIdentityHash] };
}

// Example usage for RangeProof
async function generateRangeWitness() {
    const privateX = 50;
    const publicMin = 10;
    const publicMax = 100;

    const inputs = {
        x: privateX,
        public_min: publicMin,
        public_max: publicMax,
    };

    console.log(`Generating witness for RangeProof with inputs: ${JSON.stringify(inputs)}`);
    const witness = await generateWitness('range', inputs);
    console.log('RangeProof Witness generated.');
    return { witness, publicInputs: [publicMin, publicMax] };
}

// To run:
// ts-node witness_generator.ts
// (Ensure circom_runtime is installed: npm install circom_runtime)

cargo new --bin zkp_prover
cd zkp_prover

[package]
name = "zkp_prover"
version = "0.1.0"
edition = "2021"

[dependencies]
ark-bn254 = { version = "0.4.0", features = ["curve"] }
ark-circom = "0.4.0"
ark-ff = "0.4.0"
ark-groth16 = "0.4.0"
ark-relations = "0.4.0"
ark-std = { version = "0.4.0", features = ["print-trace"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
hex = "0.4"

// src/main.rs
use ark_bn254::{Bn254, Fr};
use ark_circom::{CircomBuilder, CircomCircuit, R1CS};
use ark_ff::{BigInt, PrimeField};
use ark_groth16::{
    create_random_proof, generate_random_parameters, prepare_verifying_key, verify_proof, Proof,
};
use ark_std::{rand::thread_rng, UniformRand};
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, fs::File, io::BufReader, path::PathBuf};

// Structure to hold the proof data for serialization
#[derive(Serialize, Deserialize, Debug)]
pub struct Groth16Proof {
    pub a: Vec<String>, // G1 point (x, y)
    pub b: Vec<Vec<String>>, // G2 point (x[0], x[1], y[0], y[1])
    pub c: Vec<String>, // G1 point (x, y)
}

impl From<Proof<Bn254>> for Groth16Proof {
    fn from(proof: Proof<Bn254>) -> Self {
        Self {
            a: vec![
                format!("{:?}", proof.a.x),
                format!("{:?}", proof.a.y),
            ],
            b: vec![
                vec![
                    format!("{:?}", proof.b.x.c0),
                    format!("{:?}", proof.b.x.c1),
                ],
                vec![
                    format!("{:?}", proof.b.y.c0),
                    format!("{:?}", proof.b.y.c1),
                ],
            ],
            c: vec![
                format!("{:?}", proof.c.x),
                format!("{:?}", proof.c.y),
            ],
        }
    }
}

// Helper to convert field elements to BigInt for CircomBuilder
fn field_to_bigint<F: PrimeField>(f: F) -> BigInt<4> {
    f.into_bigint()
}

async fn generate_identity_proof(
    r1cs_path: PathBuf,
    zkey_path: PathBuf,
    private_x: u64,
    public_identity_hash: u64,
) -> Result<(Groth16Proof, Vec<Fr>), Box<dyn std::error::Error>> {
    let mut rng = thread_rng();

    // 1. Load R1CS
    let r1cs = R1CS::from_file(r1cs_path)?;

    // 2. Create CircomBuilder and assign inputs
    let mut builder = CircomBuilder::new(r1cs);
    builder.push_input("x", private_x);
    builder.push_input("identity_hash", public_identity_hash);

    // 3. Build the circuit and generate the witness
    let circuit = builder.setup();
    let full_assignments = circuit.generate_witness()?;

    // Extract public inputs from the full witness
    // The first element is always 1, then public inputs, then private inputs.
    // For IdentityProof, public_identity_hash is the only public input.
    let public_inputs = vec![full_assignments[1]]; // Assuming identity_hash is the first public signal after 1

    // 4. Load proving key from zkey file
    // In a real scenario, you'd load the proving key directly, not generate it.
    // For simplicity, we'll generate parameters here using the R1CS.
    // For production, the `zkey_path` would contain the pre-computed proving key.
    // ark-circom's `CircomBuilder` can directly load the proving key from a zkey.
    // However, `ark-groth16` expects `ProvingKey`.
    // A more direct way to load from zkey is complex due to `ark-circom`'s current API.
    // For this example, we'll generate parameters from R1CS, which is NOT production-ready.
    // Production: Load PK from `zkey_path` (e.g., using `snarkjs zkey export json` and parsing).
    let params = generate_random_parameters::<Bn254, _, _>(circuit.clone(), &mut rng)?;

    // 5. Create the proof
    let proof = create_random_proof(circuit, params, &mut rng)?;

    Ok((proof.into(), public_inputs))
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Paths to your compiled Circom artifacts and zkey
    let r1cs_path = PathBuf::from("./identity.r1cs");
    let zkey_path = PathBuf::from("./identity_final.zkey"); // This is the proving key

    let private_x = 12345;
    let public_identity_hash = 12345;

    println!("Generating IdentityProof...");
    let (proof, public_inputs) = generate_identity_proof(
        r1cs_path,
        zkey_path,
        private_x,
        public_identity_hash,
    ).await?;

    println!("Proof generated: {:?}", proof);
    println!("Public inputs: {:?}", public_inputs);

    // For verification, you would typically send `proof` and `public_inputs` to a verifier.
    // Here, we'll demonstrate local verification using ark-groth16.
    // In production, the verification key would be loaded from `verification_key_identity.json`.

    // Load verification key (for local verification demonstration)
    let vk_file = File::open("./verification_key_identity.json")?;
    let vk_json: serde_json::Value = serde_json::from_reader(BufReader::new(vk_file))?;

    // Parse the verification key from JSON into ark-groth16's `VerifyingKey` struct.
    // This parsing is non-trivial and requires careful mapping of JSON fields to struct fields.
    // For simplicity, we'll re-generate VK from params for local verification.
    // Production: Parse `verification_key_identity.json` into `ark_groth16::VerifyingKey<Bn254>`.
    // This is a complex step due to the structure of `ark-groth16`'s `VerifyingKey` and `snarkjs`'s JSON output.
    // A common approach is to use `snarkjs zkey export solidityverifier` and then use the generated Solidity contract.
    // Or, write a custom parser for the `verification_key_identity.json` into `ark_groth16::VerifyingKey`.
    // For this example, we'll use the `params` generated earlier to get the VK.
    // This means the local verification uses the same setup as proof generation, which is not how production works.
    // In production, the VK is fixed after trusted setup.
    let r1cs_for_vk = R1CS::from_file(PathBuf::from("./identity.r1cs"))?;
    let circuit_for_vk = CircomBuilder::new(r1cs_for_vk).setup();
    let params_for_vk = generate_random_parameters::<Bn254, _, _>(circuit_for_vk, &mut rng)?;
    let pvk = prepare_verifying_key(&params_for_vk.vk);

    // Convert the generated proof back to ark-groth16's Proof struct for local verification
    let ark_proof = Proof {
        a: ark_bn254::G1Affine::new(
            Fr::from_str_radix(&proof.a[0].trim_start_matches("0x"), 16)?,
            Fr::from_str_radix(&proof.a[1].trim_start_matches("0x"), 16)?,
        ),
        b: ark_bn254::G2Affine::new(
            ark_bn254::Fq2::new(
                Fr::from_str_radix(&proof.b[0][0].trim_start_matches("0x"), 16)?,
                Fr::from_str_radix(&proof.b[0][1].trim_start_matches("0x"), 16)?,
            ),
            ark_bn254::Fq2::new(
                Fr::from_str_radix(&proof.b[1][0].trim_start_matches("0x"), 16)?,
                Fr::from_str_radix(&proof.b[1][1].trim_start_matches("0x"), 16)?,
            ),
        ),
        c: ark_bn254::G1Affine::new(
            Fr::from_str_radix(&proof.c[0].trim_start_matches("0x"), 16)?,
            Fr::from_str_radix(&proof.c[1].trim_start_matches("0x"), 16)?,
        ),
    };

    let is_valid = verify_proof(&pvk, &ark_proof, &public_inputs)?;
    println!("Local verification result: {}", is_valid);

    Ok(())
}

// verifier.ts
import * as snarkjs from 'snarkjs';
import * as fs from 'fs';
import * as path from 'path';

async function verifyProof(
    circuitName: string,
    publicInputs: any[],
    proof: any
): Promise<boolean> {
    const vKeyPath = path.join(__dirname, `verification_key_${circuitName}.json`);
    const vKey = JSON.parse(fs.readFileSync(vKeyPath, 'utf-8'));

    const res = await snarkjs.groth16.verify(vKey, publicInputs, proof);
    return res;
}

// Example usage (assuming you have a proof and public inputs from the Rust prover)
async function runVerification() {
    // These would typically come from the Rust prover service
    const identityProof = {
        pi_a: ["0x...", "0x...", "0x..."], // G1 point
        pi_b: [["0x...", "0x..."], ["0x...", "0x..."]], // G2 point
        pi_c: ["0x...", "0x...", "0x..."] // G1 point
    };
    const identityPublicInputs = ["12345"];

    console.log('Verifying IdentityProof...');
    const isValidIdentity = await verifyProof('identity', identityPublicInputs, identityProof);
    console.log(`IdentityProof verification result: ${isValidIdentity}`);

    // Example for RangeProof
    const rangeProof = { /* ... */ };
    const rangePublicInputs = ["10", "100"];
    console.log('Verifying RangeProof...');
    const isValidRange = await verifyProof('range', rangePublicInputs, rangeProof);
    console.log(`RangeProof verification result: ${isValidRange}`);
}

// To run:
// ts-node verifier.ts
// (Ensure snarkjs is installed: npm install snarkjs)

# Export Solidity verifier for IdentityProof
snarkjs zkey export solidityverifier identity_final.zkey verifier_identity.sol

# Export Solidity verifier for RangeProof
snarkjs zkey export solidityverifier range_final.zkey verifier_range.sol

// verifier_identity.sol (simplified)
pragma solidity ^0.8.0;

contract Verifier {
    function verifyProof(
        uint[2] memory _pA,
        uint[2][2] memory _pB,
        uint[2] memory _pC,
        uint[1] memory _pubSignals
    ) public view returns (bool) {
        // ... cryptographic verification logic ...
        // This function will return true if the proof is valid for the given public signals.
    }
}

// web3_verifier.ts (example using ethers.js)
import { ethers } from 'ethers';
import * as fs from 'fs';
import * as path from 'path';

// Assuming you have a deployed Verifier contract
const VERIFIER_CONTRACT_ADDRESS = "0x..."; // Replace with your deployed contract address
const VERIFIER_ABI = JSON.parse(fs.readFileSync(path.join(__dirname, 'Verifier_abi.json'), 'utf-8'));

async function verifyOnChain(
    provider: ethers.Provider,
    proof: any, // SnarkJS proof format
    publicInputs: string[]
): Promise<boolean> {
    const verifier = new ethers.Contract(VERIFIER_CONTRACT_ADDRESS, VERIFIER_ABI, provider);

    // Convert SnarkJS proof format to Solidity-compatible format
    const pA: [string, string] = [proof.pi_a[0], proof.pi_a[1]];
    const pB: [[string, string], [string, string]] = [
        [proof.pi_b[0][0], proof.pi_b[0][1]],
        [proof.pi_b[1][0], proof.pi_b[1][1]]
    ];
    const pC: [string, string] = [proof.pi_c[0], proof.pi_c[1]];

    // Public signals need to be `uint` in Solidity, so convert from string
    const pubSignals = publicInputs.map(s => ethers.BigNumber.from(s));

    try {
        const isValid = await verifier.verifyProof(pA, pB, pC, pubSignals);
        return isValid;
    } catch (error) {
        console.error("On-chain verification failed:", error);
        return false;
    }
}

// Example usage:
// const provider = new ethers.JsonRpcProvider("YOUR_RPC_URL");
// const proof = { /* ... from Rust prover ... */ };
// const publicInputs = ["12345"];
// const isValid = await verifyOnChain(provider, proof, publicInputs);
// console.log(`On-chain verification result: ${isValid}`);