zksvm_client/
verify.rs

1use std::collections::{BTreeMap, HashSet};
2
3use rs_merkle::{algorithms::Sha256, Hasher, MerkleProof, MerkleTree};
4use solana_client::rpc_client::RpcClient;
5use solana_sdk::{
6    account::Account,
7    compute_budget::ComputeBudgetInstruction,
8    instruction::Instruction,
9    message::Message,
10    pubkey::Pubkey,
11    signature::{Keypair, Signature},
12    signer::Signer,
13    transaction::Transaction,
14};
15use sp1_sdk::{HashableKey, SP1ProofWithPublicValues, SP1VerifyingKey};
16use thiserror::Error;
17use zksvm_api_types::common::ZKProof;
18use zksvm_lib::merkle::{merkelize_accounts, merkelize_transactions};
19
20use crate::types::{InstructionData, SP1Groth16Proof};
21
22#[derive(Debug, Error)]
23pub enum SVMVerificationError {
24    #[error("Failed Solana RPC request: {0}")]
25    SolanaRPCError(#[from] solana_client::client_error::ClientError),
26    #[error("Failed Merkle root verification: {0}")]
27    MerkleRootError(String),
28    #[error("Failed to deserialize bytes: {0}")]
29    SerializationError(#[from] bincode::Error),
30    #[error("Failed to convert to slice: {0}")]
31    TryFromError(#[from] std::array::TryFromSliceError),
32    #[error("Failed proof verification: {0}")]
33    ProofVerificationError(String),
34}
35
36/// Verifies a zero-knowledge proof locally using Groth16 verification.
37pub fn verify_locally(proof: ZKProof) -> Result<(), SVMVerificationError> {
38    // Prepare proof and verification key
39    let (groth_proof, vk) = extract_proof_and_verification_key(proof)?;
40
41    // try to verify the proof
42    sp1_verifier::Groth16Verifier::verify(
43        groth_proof.bytes().as_slice(),
44        groth_proof.public_values.as_slice(),
45        &vk.bytes32(),
46        *sp1_verifier::GROTH16_VK_BYTES,
47    )
48    .map_err(|e| SVMVerificationError::ProofVerificationError(e.to_string()))
49}
50
51/// Verifies a zero-knowledge proof on-chain by sending a transaction to the Solana network.
52pub fn verify_on_chain(
53    solana_rpc_client: RpcClient,
54    program_id: Pubkey,
55    payer: Keypair,
56    proof: ZKProof,
57) -> Result<Signature, SVMVerificationError> {
58    // Prepare proof and verification key
59    let (groth_proof, vk) = extract_proof_and_verification_key(proof)?;
60
61    let data = InstructionData {
62        groth16_proof: SP1Groth16Proof {
63            proof: groth_proof.bytes(),
64            sp1_public_inputs: groth_proof.public_values.to_vec(),
65        },
66        vkey_hash: vk.bytes32(),
67    };
68    let data = borsh::to_vec(&data).map_err(|e| {
69        SVMVerificationError::ProofVerificationError(format!(
70            "Error translating data to raw format: {e}"
71        ))
72    })?;
73
74    // Verification takes ~280K CUs
75    let set_cu_limit_instruction = ComputeBudgetInstruction::set_compute_unit_limit(300_000);
76
77    // Create the instruction for the program
78    let instruction = Instruction {
79        program_id,
80        accounts: vec![],
81        data,
82    };
83
84    // Create a transaction with the instructions
85    let message = Message::new(
86        &[set_cu_limit_instruction, instruction],
87        Some(&payer.pubkey()),
88    );
89    let mut transaction = Transaction::new_unsigned(message);
90
91    // Sign the transaction with the payer's keypair
92    let latest_blockhash = solana_rpc_client.get_latest_blockhash()?;
93    transaction.sign(&[&payer], latest_blockhash);
94    // Send the transaction and confirm
95    Ok(solana_rpc_client.send_and_confirm_transaction(&transaction)?)
96}
97
98fn extract_proof_and_verification_key(
99    proof: ZKProof,
100) -> Result<(SP1ProofWithPublicValues, SP1VerifyingKey), SVMVerificationError> {
101    // Prepare proof
102    let bytes = proof.proof.ok_or_else(|| {
103        SVMVerificationError::ProofVerificationError("Could not unwrap proof bytes".to_string())
104    })?;
105    let groth_proof = bincode::deserialize::<SP1ProofWithPublicValues>(bytes.as_slice())?;
106
107    // Prepare verification key
108    let bytes = proof.verification_key.ok_or_else(|| {
109        SVMVerificationError::ProofVerificationError(
110            "Could not deserialize unwrap verification key bytes".to_string(),
111        )
112    })?;
113    let vk = bincode::deserialize::<SP1VerifyingKey>(&bytes)?;
114
115    Ok((groth_proof, vk))
116}
117
118#[allow(dead_code)]
119pub struct ZKProverMerkleRoots {
120    initial_state_merkle: [u8; 32],
121    final_state_merkle: [u8; 32],
122    transactions: [u8; 32],
123}
124
125impl ZKProverMerkleRoots {
126    pub fn new(proof: ZKProof) -> Result<Self, SVMVerificationError> {
127        let bytes = proof.proof.clone().ok_or_else(|| {
128            SVMVerificationError::ProofVerificationError("Could not unwrap proof bytes".to_string())
129        })?;
130        let groth_proof = bincode::deserialize::<SP1ProofWithPublicValues>(bytes.as_slice())?;
131        let merkle_roots = groth_proof.public_values.to_vec();
132
133        if merkle_roots.len() != 96 {
134            return Err(SVMVerificationError::MerkleRootError(
135                "The ZK proof should commit to 3 Merkle root for a total of 96 bytes".to_string(),
136            ));
137        }
138        let initial_state_merkle = merkle_roots[..33].try_into()?;
139        let transactions = merkle_roots[33..65].try_into()?;
140        let final_state_merkle = merkle_roots[65..].try_into()?;
141
142        Ok(Self {
143            initial_state_merkle,
144            transactions,
145            final_state_merkle,
146        })
147    }
148}
149
150pub struct AccountsMerkle {
151    accounts: BTreeMap<Pubkey, Account>,
152    tree: MerkleTree<Sha256>,
153}
154
155impl AccountsMerkle {
156    pub fn new(accounts: BTreeMap<Pubkey, Account>) -> Self {
157        Self {
158            tree: merkelize_accounts(&accounts),
159            accounts,
160        }
161    }
162
163    /// Verify that the Merkle root matches the expected hash.
164    pub fn verify_root(&self, expected_root: &[u8; 32]) -> bool {
165        match self.tree.root() {
166            Some(actual_root) => actual_root == *expected_root,
167            None => false,
168        }
169    }
170
171    /// Create a Merkle proof that specific accounts are committed to by the Merkle root.
172    pub fn generate_merkle_proof(
173        &self,
174        pubkeys: &[Pubkey],
175    ) -> Result<MerkleProof<Sha256>, SVMVerificationError> {
176        let mut leaf_indices = Vec::new();
177        let pubkeys: HashSet<_> = pubkeys.iter().cloned().collect();
178        for (index, (pubkey, _)) in self.accounts.iter().enumerate() {
179            if pubkeys.contains(pubkey) {
180                leaf_indices.push(index);
181            }
182        }
183
184        let bytes = self.tree.proof(leaf_indices.as_slice()).to_bytes();
185        MerkleProof::<Sha256>::try_from(bytes)
186            .map_err(|e| SVMVerificationError::MerkleRootError(e.to_string()))
187    }
188
189    /// Verifies the Merkle proof for the given accounts.
190    pub fn verify_merkle_proof(
191        &self,
192        proof: &MerkleProof<Sha256>,
193        pubkeys: &[Pubkey],
194    ) -> Result<bool, SVMVerificationError> {
195        let mut indices = Vec::new();
196        let mut leaves = Vec::new();
197        let pubkeys: HashSet<_> = pubkeys.iter().cloned().collect();
198        for (index, (pubkey, account)) in self.accounts.iter().enumerate() {
199            if pubkeys.contains(pubkey) {
200                let bytes = [pubkey.as_ref(), bincode::serialize(account)?.as_slice()].concat();
201                let hash = Sha256::hash(bytes.as_slice());
202                indices.push(index);
203                leaves.push(hash);
204            }
205        }
206
207        let root = self.tree.root().ok_or_else(|| {
208            SVMVerificationError::ProofVerificationError("Could not retrieve tree root".to_string())
209        })?;
210        Ok(proof.verify(
211            root,
212            indices.as_slice(),
213            leaves.as_slice(),
214            self.accounts.len(),
215        ))
216    }
217}
218
219pub struct TransactionsMerkle {
220    transactions: Vec<Transaction>,
221    tree: MerkleTree<Sha256>,
222}
223
224impl TransactionsMerkle {
225    pub fn new(transactions: Vec<Transaction>) -> Self {
226        Self {
227            tree: merkelize_transactions(&transactions),
228            transactions,
229        }
230    }
231
232    /// Verify that the Merkle root matches the expected hash.
233    pub fn verify_root(&self, expected_root: &[u8; 32]) -> bool {
234        match self.tree.root() {
235            Some(actual_root) => actual_root == *expected_root,
236            None => false,
237        }
238    }
239
240    /// Generates a Merkle proof for the specified transaction indices.
241    ///
242    /// This function creates a Merkle proof that proves the inclusion of specific
243    /// transactions within the Merkle tree. The proof can be used to verify that
244    /// these transactions are committed to by the Merkle root.
245    ///
246    /// # Arguments
247    ///
248    /// * `indices` - A slice of indices representing the transactions for which
249    ///   the proof should be generated.
250    ///
251    /// # Returns
252    ///
253    /// * `Ok(MerkleProof<Sha256>)` - The generated Merkle proof.
254    /// * `Err(SVMVerificationError)` - An error if proof generation fails.
255    ///
256    /// # Errors
257    ///
258    /// This function returns an error if the proof cannot be generated or if there
259    /// is an issue converting the proof bytes into a `MerkleProof<Sha256>`.
260    pub fn generate_merkle_proof(
261        &self,
262        indices: &[usize],
263    ) -> Result<MerkleProof<Sha256>, SVMVerificationError> {
264        let bytes = self.tree.proof(indices).to_bytes();
265        MerkleProof::<Sha256>::try_from(bytes)
266            .map_err(|e| SVMVerificationError::MerkleRootError(e.to_string()))
267    }
268
269    /// Verifies a Merkle proof for the specified transactions.
270    ///
271    /// This function checks whether the provided Merkle proof correctly proves
272    /// the inclusion of transactions at the given indices in the Merkle tree.
273    ///
274    /// # Arguments
275    ///
276    /// * `proof` - A reference to the `MerkleProof<Sha256>` to be verified.
277    /// * `indices` - A slice of indices corresponding to the transactions
278    ///   that the proof should validate.
279    ///
280    /// # Returns
281    ///
282    /// * `true` if the proof is valid and the transactions are correctly
283    ///   committed to by the Merkle root.
284    /// * `false` if the proof is invalid.
285    ///
286    /// # Panics
287    ///
288    /// This function may panic if:
289    /// * An index is out of bounds and the corresponding transaction cannot be retrieved.
290    /// * Serialization of a transaction fails.
291    /// * The Merkle tree root is unavailable.
292    pub fn verify_merkle_proof(
293        &self,
294        proof: &MerkleProof<Sha256>,
295        indices: &[usize],
296    ) -> Result<bool, SVMVerificationError> {
297        let mut leaves = Vec::new();
298
299        for i in indices {
300            let tx = self.transactions.get(*i).ok_or_else(|| {
301                SVMVerificationError::ProofVerificationError(
302                    "Could not retrieve transaction".to_string(),
303                )
304            })?;
305            let bytes = bincode::serialize(tx)?;
306            leaves.push(Sha256::hash(bytes.as_slice()));
307        }
308
309        let root = self.tree.root().ok_or_else(|| {
310            SVMVerificationError::ProofVerificationError("Could not retrieve tree root".to_string())
311        })?;
312        Ok(proof.verify(root, indices, leaves.as_slice(), self.transactions.len()))
313    }
314}