solana_zk_token_proof_program/
lib.rs

1#![forbid(unsafe_code)]
2// Allow deprecated warnings since this crate will be removed along with
3// `solana-zk-token-sdk` will be removed
4#![allow(deprecated)]
5
6use {
7    bytemuck::Pod,
8    solana_instruction::{error::InstructionError, TRANSACTION_LEVEL_STACK_HEIGHT},
9    solana_program_runtime::{declare_process_instruction, invoke_context::InvokeContext},
10    solana_sdk_ids::system_program,
11    solana_svm_log_collector::ic_msg,
12    solana_zk_token_sdk::{
13        zk_token_proof_instruction::*,
14        zk_token_proof_program::id,
15        zk_token_proof_state::{ProofContextState, ProofContextStateMeta},
16    },
17    std::result::Result,
18};
19
20pub const CLOSE_CONTEXT_STATE_COMPUTE_UNITS: u64 = 3_300;
21pub const VERIFY_ZERO_BALANCE_COMPUTE_UNITS: u64 = 6_000;
22pub const VERIFY_WITHDRAW_COMPUTE_UNITS: u64 = 110_000;
23pub const VERIFY_CIPHERTEXT_CIPHERTEXT_EQUALITY_COMPUTE_UNITS: u64 = 8_000;
24pub const VERIFY_TRANSFER_COMPUTE_UNITS: u64 = 219_000;
25pub const VERIFY_TRANSFER_WITH_FEE_COMPUTE_UNITS: u64 = 407_000;
26pub const VERIFY_PUBKEY_VALIDITY_COMPUTE_UNITS: u64 = 2_600;
27pub const VERIFY_RANGE_PROOF_U64_COMPUTE_UNITS: u64 = 105_000;
28pub const VERIFY_BATCHED_RANGE_PROOF_U64_COMPUTE_UNITS: u64 = 111_000;
29pub const VERIFY_BATCHED_RANGE_PROOF_U128_COMPUTE_UNITS: u64 = 200_000;
30pub const VERIFY_BATCHED_RANGE_PROOF_U256_COMPUTE_UNITS: u64 = 368_000;
31pub const VERIFY_CIPHERTEXT_COMMITMENT_EQUALITY_COMPUTE_UNITS: u64 = 6_400;
32pub const VERIFY_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_COMPUTE_UNITS: u64 = 6_400;
33pub const VERIFY_BATCHED_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_COMPUTE_UNITS: u64 = 13_000;
34pub const VERIFY_FEE_SIGMA_COMPUTE_UNITS: u64 = 6_500;
35pub const VERIFY_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_COMPUTE_UNITS: u64 = 8_100;
36pub const VERIFY_BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_COMPUTE_UNITS: u64 = 16_400;
37
38const INSTRUCTION_DATA_LENGTH_WITH_PROOF_ACCOUNT: usize = 5;
39
40fn process_verify_proof<T, U>(invoke_context: &mut InvokeContext) -> Result<(), InstructionError>
41where
42    T: Pod + ZkProofData<U>,
43    U: Pod,
44{
45    let transaction_context = &invoke_context.transaction_context;
46    let instruction_context = transaction_context.get_current_instruction_context()?;
47    let instruction_data = instruction_context.get_instruction_data();
48
49    // number of accessed accounts so far
50    let mut accessed_accounts = 0_u16;
51
52    // if instruction data is exactly 5 bytes, then read proof from an account
53    let context_data = if instruction_data.len() == INSTRUCTION_DATA_LENGTH_WITH_PROOF_ACCOUNT {
54        let enable_zk_proof_from_account = false;
55        // This code is disabled on purpose. If the feature is required to be enabled in future,
56        // a better way to lookup feature_set should be implemented/used.
57        if !enable_zk_proof_from_account {
58            return Err(InstructionError::InvalidInstructionData);
59        }
60
61        let proof_data_account =
62            instruction_context.try_borrow_instruction_account(accessed_accounts)?;
63        accessed_accounts = accessed_accounts.checked_add(1).unwrap();
64
65        let proof_data_offset = u32::from_le_bytes(
66            // the first byte is the instruction discriminator
67            instruction_data[1..INSTRUCTION_DATA_LENGTH_WITH_PROOF_ACCOUNT]
68                .try_into()
69                .map_err(|_| InstructionError::InvalidInstructionData)?,
70        );
71        let proof_data_start: usize = proof_data_offset
72            .try_into()
73            .map_err(|_| InstructionError::InvalidInstructionData)?;
74        let proof_data_end = proof_data_start
75            .checked_add(std::mem::size_of::<T>())
76            .ok_or(InstructionError::InvalidInstructionData)?;
77        let proof_data_raw = proof_data_account
78            .get_data()
79            .get(proof_data_start..proof_data_end)
80            .ok_or(InstructionError::InvalidAccountData)?;
81
82        let proof_data = bytemuck::try_from_bytes::<T>(proof_data_raw).map_err(|_| {
83            ic_msg!(invoke_context, "invalid proof data");
84            InstructionError::InvalidInstructionData
85        })?;
86        proof_data.verify_proof().map_err(|err| {
87            ic_msg!(invoke_context, "proof verification failed: {:?}", err);
88            InstructionError::InvalidInstructionData
89        })?;
90
91        *proof_data.context_data()
92    } else {
93        let proof_data =
94            ProofInstruction::proof_data::<T, U>(instruction_data).ok_or_else(|| {
95                ic_msg!(invoke_context, "invalid proof data");
96                InstructionError::InvalidInstructionData
97            })?;
98        proof_data.verify_proof().map_err(|err| {
99            ic_msg!(invoke_context, "proof_verification failed: {:?}", err);
100            InstructionError::InvalidInstructionData
101        })?;
102
103        *proof_data.context_data()
104    };
105
106    // create context state if additional accounts are provided with the instruction
107    if instruction_context.get_number_of_instruction_accounts() > accessed_accounts {
108        let context_state_authority = *instruction_context
109            .get_key_of_instruction_account(accessed_accounts.checked_add(1).unwrap())?;
110        let mut proof_context_account =
111            instruction_context.try_borrow_instruction_account(accessed_accounts)?;
112
113        if *proof_context_account.get_owner() != id() {
114            return Err(InstructionError::InvalidAccountOwner);
115        }
116
117        let proof_context_state_meta =
118            ProofContextStateMeta::try_from_bytes(proof_context_account.get_data())?;
119
120        if proof_context_state_meta.proof_type != ProofType::Uninitialized.into() {
121            return Err(InstructionError::AccountAlreadyInitialized);
122        }
123
124        let context_state_data =
125            ProofContextState::encode(&context_state_authority, T::PROOF_TYPE, &context_data);
126
127        if proof_context_account.get_data().len() != context_state_data.len() {
128            return Err(InstructionError::InvalidAccountData);
129        }
130
131        proof_context_account.set_data_from_slice(&context_state_data)?;
132    }
133
134    Ok(())
135}
136
137fn process_close_proof_context(invoke_context: &mut InvokeContext) -> Result<(), InstructionError> {
138    let transaction_context = &invoke_context.transaction_context;
139    let instruction_context = transaction_context.get_current_instruction_context()?;
140
141    let owner_pubkey = {
142        if !instruction_context.is_instruction_account_signer(2)? {
143            return Err(InstructionError::MissingRequiredSignature);
144        }
145
146        *instruction_context.get_key_of_instruction_account(2)?
147    };
148
149    let proof_context_account_pubkey = *instruction_context.get_key_of_instruction_account(0)?;
150    let destination_account_pubkey = *instruction_context.get_key_of_instruction_account(1)?;
151    if proof_context_account_pubkey == destination_account_pubkey {
152        return Err(InstructionError::InvalidInstructionData);
153    }
154
155    let mut proof_context_account = instruction_context.try_borrow_instruction_account(0)?;
156    let proof_context_state_meta =
157        ProofContextStateMeta::try_from_bytes(proof_context_account.get_data())?;
158    let expected_owner_pubkey = proof_context_state_meta.context_state_authority;
159
160    if owner_pubkey != expected_owner_pubkey {
161        return Err(InstructionError::InvalidAccountOwner);
162    }
163
164    let mut destination_account = instruction_context.try_borrow_instruction_account(1)?;
165    destination_account.checked_add_lamports(proof_context_account.get_lamports())?;
166    proof_context_account.set_lamports(0)?;
167    proof_context_account.set_data_length(0)?;
168    proof_context_account.set_owner(system_program::id().as_ref())?;
169
170    Ok(())
171}
172
173declare_process_instruction!(Entrypoint, 0, |invoke_context| {
174    let enable_zk_transfer_with_fee = false;
175
176    let transaction_context = &invoke_context.transaction_context;
177    let instruction_context = transaction_context.get_current_instruction_context()?;
178    let instruction_data = instruction_context.get_instruction_data();
179    let instruction = ProofInstruction::instruction_type(instruction_data)
180        .ok_or(InstructionError::InvalidInstructionData)?;
181
182    if invoke_context.get_stack_height() != TRANSACTION_LEVEL_STACK_HEIGHT
183        && instruction != ProofInstruction::CloseContextState
184    {
185        // Proof verification instructions are not supported as an inner instruction
186        return Err(InstructionError::UnsupportedProgramId);
187    }
188
189    match instruction {
190        ProofInstruction::CloseContextState => {
191            invoke_context
192                .consume_checked(CLOSE_CONTEXT_STATE_COMPUTE_UNITS)
193                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
194            ic_msg!(invoke_context, "CloseContextState");
195            process_close_proof_context(invoke_context)
196        }
197        ProofInstruction::VerifyZeroBalance => {
198            invoke_context
199                .consume_checked(VERIFY_ZERO_BALANCE_COMPUTE_UNITS)
200                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
201            ic_msg!(invoke_context, "VerifyZeroBalance");
202            process_verify_proof::<ZeroBalanceProofData, ZeroBalanceProofContext>(invoke_context)
203        }
204        ProofInstruction::VerifyWithdraw => {
205            invoke_context
206                .consume_checked(VERIFY_WITHDRAW_COMPUTE_UNITS)
207                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
208            ic_msg!(invoke_context, "VerifyWithdraw");
209            process_verify_proof::<WithdrawData, WithdrawProofContext>(invoke_context)
210        }
211        ProofInstruction::VerifyCiphertextCiphertextEquality => {
212            invoke_context
213                .consume_checked(VERIFY_CIPHERTEXT_CIPHERTEXT_EQUALITY_COMPUTE_UNITS)
214                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
215            ic_msg!(invoke_context, "VerifyCiphertextCiphertextEquality");
216            process_verify_proof::<
217                CiphertextCiphertextEqualityProofData,
218                CiphertextCiphertextEqualityProofContext,
219            >(invoke_context)
220        }
221        ProofInstruction::VerifyTransfer => {
222            invoke_context
223                .consume_checked(VERIFY_TRANSFER_COMPUTE_UNITS)
224                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
225            ic_msg!(invoke_context, "VerifyTransfer");
226            process_verify_proof::<TransferData, TransferProofContext>(invoke_context)
227        }
228        ProofInstruction::VerifyTransferWithFee => {
229            // transfer with fee related proofs are not enabled
230            if !enable_zk_transfer_with_fee {
231                return Err(InstructionError::InvalidInstructionData);
232            }
233
234            invoke_context
235                .consume_checked(VERIFY_TRANSFER_WITH_FEE_COMPUTE_UNITS)
236                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
237            ic_msg!(invoke_context, "VerifyTransferWithFee");
238            process_verify_proof::<TransferWithFeeData, TransferWithFeeProofContext>(invoke_context)
239        }
240        ProofInstruction::VerifyPubkeyValidity => {
241            invoke_context
242                .consume_checked(VERIFY_PUBKEY_VALIDITY_COMPUTE_UNITS)
243                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
244            ic_msg!(invoke_context, "VerifyPubkeyValidity");
245            process_verify_proof::<PubkeyValidityData, PubkeyValidityProofContext>(invoke_context)
246        }
247        ProofInstruction::VerifyRangeProofU64 => {
248            invoke_context
249                .consume_checked(VERIFY_RANGE_PROOF_U64_COMPUTE_UNITS)
250                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
251            ic_msg!(invoke_context, "VerifyRangeProof");
252            process_verify_proof::<RangeProofU64Data, RangeProofContext>(invoke_context)
253        }
254        ProofInstruction::VerifyBatchedRangeProofU64 => {
255            invoke_context
256                .consume_checked(VERIFY_BATCHED_RANGE_PROOF_U64_COMPUTE_UNITS)
257                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
258            ic_msg!(invoke_context, "VerifyBatchedRangeProof64");
259            process_verify_proof::<BatchedRangeProofU64Data, BatchedRangeProofContext>(
260                invoke_context,
261            )
262        }
263        ProofInstruction::VerifyBatchedRangeProofU128 => {
264            invoke_context
265                .consume_checked(VERIFY_BATCHED_RANGE_PROOF_U128_COMPUTE_UNITS)
266                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
267            ic_msg!(invoke_context, "VerifyBatchedRangeProof128");
268            process_verify_proof::<BatchedRangeProofU128Data, BatchedRangeProofContext>(
269                invoke_context,
270            )
271        }
272        ProofInstruction::VerifyBatchedRangeProofU256 => {
273            // transfer with fee related proofs are not enabled
274            if !enable_zk_transfer_with_fee {
275                return Err(InstructionError::InvalidInstructionData);
276            }
277
278            invoke_context
279                .consume_checked(VERIFY_BATCHED_RANGE_PROOF_U256_COMPUTE_UNITS)
280                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
281            ic_msg!(invoke_context, "VerifyBatchedRangeProof256");
282            process_verify_proof::<BatchedRangeProofU256Data, BatchedRangeProofContext>(
283                invoke_context,
284            )
285        }
286        ProofInstruction::VerifyCiphertextCommitmentEquality => {
287            invoke_context
288                .consume_checked(VERIFY_CIPHERTEXT_COMMITMENT_EQUALITY_COMPUTE_UNITS)
289                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
290            ic_msg!(invoke_context, "VerifyCiphertextCommitmentEquality");
291            process_verify_proof::<
292                CiphertextCommitmentEqualityProofData,
293                CiphertextCommitmentEqualityProofContext,
294            >(invoke_context)
295        }
296        ProofInstruction::VerifyGroupedCiphertext2HandlesValidity => {
297            invoke_context
298                .consume_checked(VERIFY_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_COMPUTE_UNITS)
299                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
300            ic_msg!(invoke_context, "VerifyGroupedCiphertext2HandlesValidity");
301            process_verify_proof::<
302                GroupedCiphertext2HandlesValidityProofData,
303                GroupedCiphertext2HandlesValidityProofContext,
304            >(invoke_context)
305        }
306        ProofInstruction::VerifyBatchedGroupedCiphertext2HandlesValidity => {
307            invoke_context
308                .consume_checked(VERIFY_BATCHED_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_COMPUTE_UNITS)
309                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
310            ic_msg!(
311                invoke_context,
312                "VerifyBatchedGroupedCiphertext2HandlesValidity"
313            );
314            process_verify_proof::<
315                BatchedGroupedCiphertext2HandlesValidityProofData,
316                BatchedGroupedCiphertext2HandlesValidityProofContext,
317            >(invoke_context)
318        }
319        ProofInstruction::VerifyFeeSigma => {
320            // transfer with fee related proofs are not enabled
321            if !enable_zk_transfer_with_fee {
322                return Err(InstructionError::InvalidInstructionData);
323            }
324
325            invoke_context
326                .consume_checked(VERIFY_FEE_SIGMA_COMPUTE_UNITS)
327                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
328            ic_msg!(invoke_context, "VerifyFeeSigma");
329            process_verify_proof::<FeeSigmaProofData, FeeSigmaProofContext>(invoke_context)
330        }
331        ProofInstruction::VerifyGroupedCiphertext3HandlesValidity => {
332            invoke_context
333                .consume_checked(VERIFY_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_COMPUTE_UNITS)
334                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
335            ic_msg!(invoke_context, "VerifyGroupedCiphertext3HandlesValidity");
336            process_verify_proof::<
337                GroupedCiphertext3HandlesValidityProofData,
338                GroupedCiphertext3HandlesValidityProofContext,
339            >(invoke_context)
340        }
341        ProofInstruction::VerifyBatchedGroupedCiphertext3HandlesValidity => {
342            invoke_context
343                .consume_checked(VERIFY_BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_COMPUTE_UNITS)
344                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
345            ic_msg!(
346                invoke_context,
347                "VerifyBatchedGroupedCiphertext3HandlesValidity"
348            );
349            process_verify_proof::<
350                BatchedGroupedCiphertext3HandlesValidityProofData,
351                BatchedGroupedCiphertext3HandlesValidityProofContext,
352            >(invoke_context)
353        }
354    }
355});