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
36pub fn verify_locally(proof: ZKProof) -> Result<(), SVMVerificationError> {
38 let (groth_proof, vk) = extract_proof_and_verification_key(proof)?;
40
41 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
51pub fn verify_on_chain(
53 solana_rpc_client: RpcClient,
54 program_id: Pubkey,
55 payer: Keypair,
56 proof: ZKProof,
57) -> Result<Signature, SVMVerificationError> {
58 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 let set_cu_limit_instruction = ComputeBudgetInstruction::set_compute_unit_limit(300_000);
76
77 let instruction = Instruction {
79 program_id,
80 accounts: vec![],
81 data,
82 };
83
84 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 let latest_blockhash = solana_rpc_client.get_latest_blockhash()?;
93 transaction.sign(&[&payer], latest_blockhash);
94 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 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 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 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 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 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 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 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 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}