solana_zk_token_proof_program/
lib.rs

1#![forbid(unsafe_code)]
2
3use {
4    bytemuck::Pod,
5    solana_feature_set as feature_set,
6    solana_log_collector::ic_msg,
7    solana_program_runtime::{declare_process_instruction, invoke_context::InvokeContext},
8    solana_sdk::{
9        instruction::{InstructionError, TRANSACTION_LEVEL_STACK_HEIGHT},
10        system_program,
11    },
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        if !invoke_context
55            .get_feature_set()
56            .is_active(&feature_set::enable_zk_proof_from_account::id())
57        {
58            return Err(InstructionError::InvalidInstructionData);
59        }
60
61        let proof_data_account = instruction_context
62            .try_borrow_instruction_account(transaction_context, 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            .try_borrow_instruction_account(
110                transaction_context,
111                accessed_accounts.checked_add(1).unwrap(),
112            )?
113            .get_key();
114
115        let mut proof_context_account = instruction_context
116            .try_borrow_instruction_account(transaction_context, accessed_accounts)?;
117
118        if *proof_context_account.get_owner() != id() {
119            return Err(InstructionError::InvalidAccountOwner);
120        }
121
122        let proof_context_state_meta =
123            ProofContextStateMeta::try_from_bytes(proof_context_account.get_data())?;
124
125        if proof_context_state_meta.proof_type != ProofType::Uninitialized.into() {
126            return Err(InstructionError::AccountAlreadyInitialized);
127        }
128
129        let context_state_data =
130            ProofContextState::encode(&context_state_authority, T::PROOF_TYPE, &context_data);
131
132        if proof_context_account.get_data().len() != context_state_data.len() {
133            return Err(InstructionError::InvalidAccountData);
134        }
135
136        proof_context_account.set_data_from_slice(&context_state_data)?;
137    }
138
139    Ok(())
140}
141
142fn process_close_proof_context(invoke_context: &mut InvokeContext) -> Result<(), InstructionError> {
143    let transaction_context = &invoke_context.transaction_context;
144    let instruction_context = transaction_context.get_current_instruction_context()?;
145
146    let owner_pubkey = {
147        let owner_account =
148            instruction_context.try_borrow_instruction_account(transaction_context, 2)?;
149
150        if !owner_account.is_signer() {
151            return Err(InstructionError::MissingRequiredSignature);
152        }
153        *owner_account.get_key()
154    }; // done with `owner_account`, so drop it to prevent a potential double borrow
155
156    let proof_context_account_pubkey = *instruction_context
157        .try_borrow_instruction_account(transaction_context, 0)?
158        .get_key();
159    let destination_account_pubkey = *instruction_context
160        .try_borrow_instruction_account(transaction_context, 1)?
161        .get_key();
162    if proof_context_account_pubkey == destination_account_pubkey {
163        return Err(InstructionError::InvalidInstructionData);
164    }
165
166    let mut proof_context_account =
167        instruction_context.try_borrow_instruction_account(transaction_context, 0)?;
168    let proof_context_state_meta =
169        ProofContextStateMeta::try_from_bytes(proof_context_account.get_data())?;
170    let expected_owner_pubkey = proof_context_state_meta.context_state_authority;
171
172    if owner_pubkey != expected_owner_pubkey {
173        return Err(InstructionError::InvalidAccountOwner);
174    }
175
176    let mut destination_account =
177        instruction_context.try_borrow_instruction_account(transaction_context, 1)?;
178    destination_account.checked_add_lamports(proof_context_account.get_lamports())?;
179    proof_context_account.set_lamports(0)?;
180    proof_context_account.set_data_length(0)?;
181    proof_context_account.set_owner(system_program::id().as_ref())?;
182
183    Ok(())
184}
185
186declare_process_instruction!(Entrypoint, 0, |invoke_context| {
187    let enable_zk_transfer_with_fee = invoke_context
188        .get_feature_set()
189        .is_active(&feature_set::enable_zk_transfer_with_fee::id());
190
191    let transaction_context = &invoke_context.transaction_context;
192    let instruction_context = transaction_context.get_current_instruction_context()?;
193    let instruction_data = instruction_context.get_instruction_data();
194    let instruction = ProofInstruction::instruction_type(instruction_data)
195        .ok_or(InstructionError::InvalidInstructionData)?;
196
197    if invoke_context.get_stack_height() != TRANSACTION_LEVEL_STACK_HEIGHT
198        && instruction != ProofInstruction::CloseContextState
199    {
200        // Proof verification instructions are not supported as an inner instruction
201        return Err(InstructionError::UnsupportedProgramId);
202    }
203
204    match instruction {
205        ProofInstruction::CloseContextState => {
206            invoke_context
207                .consume_checked(CLOSE_CONTEXT_STATE_COMPUTE_UNITS)
208                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
209            ic_msg!(invoke_context, "CloseContextState");
210            process_close_proof_context(invoke_context)
211        }
212        ProofInstruction::VerifyZeroBalance => {
213            invoke_context
214                .consume_checked(VERIFY_ZERO_BALANCE_COMPUTE_UNITS)
215                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
216            ic_msg!(invoke_context, "VerifyZeroBalance");
217            process_verify_proof::<ZeroBalanceProofData, ZeroBalanceProofContext>(invoke_context)
218        }
219        ProofInstruction::VerifyWithdraw => {
220            invoke_context
221                .consume_checked(VERIFY_WITHDRAW_COMPUTE_UNITS)
222                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
223            ic_msg!(invoke_context, "VerifyWithdraw");
224            process_verify_proof::<WithdrawData, WithdrawProofContext>(invoke_context)
225        }
226        ProofInstruction::VerifyCiphertextCiphertextEquality => {
227            invoke_context
228                .consume_checked(VERIFY_CIPHERTEXT_CIPHERTEXT_EQUALITY_COMPUTE_UNITS)
229                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
230            ic_msg!(invoke_context, "VerifyCiphertextCiphertextEquality");
231            process_verify_proof::<
232                CiphertextCiphertextEqualityProofData,
233                CiphertextCiphertextEqualityProofContext,
234            >(invoke_context)
235        }
236        ProofInstruction::VerifyTransfer => {
237            invoke_context
238                .consume_checked(VERIFY_TRANSFER_COMPUTE_UNITS)
239                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
240            ic_msg!(invoke_context, "VerifyTransfer");
241            process_verify_proof::<TransferData, TransferProofContext>(invoke_context)
242        }
243        ProofInstruction::VerifyTransferWithFee => {
244            // transfer with fee related proofs are not enabled
245            if !enable_zk_transfer_with_fee {
246                return Err(InstructionError::InvalidInstructionData);
247            }
248
249            invoke_context
250                .consume_checked(VERIFY_TRANSFER_WITH_FEE_COMPUTE_UNITS)
251                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
252            ic_msg!(invoke_context, "VerifyTransferWithFee");
253            process_verify_proof::<TransferWithFeeData, TransferWithFeeProofContext>(invoke_context)
254        }
255        ProofInstruction::VerifyPubkeyValidity => {
256            invoke_context
257                .consume_checked(VERIFY_PUBKEY_VALIDITY_COMPUTE_UNITS)
258                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
259            ic_msg!(invoke_context, "VerifyPubkeyValidity");
260            process_verify_proof::<PubkeyValidityData, PubkeyValidityProofContext>(invoke_context)
261        }
262        ProofInstruction::VerifyRangeProofU64 => {
263            invoke_context
264                .consume_checked(VERIFY_RANGE_PROOF_U64_COMPUTE_UNITS)
265                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
266            ic_msg!(invoke_context, "VerifyRangeProof");
267            process_verify_proof::<RangeProofU64Data, RangeProofContext>(invoke_context)
268        }
269        ProofInstruction::VerifyBatchedRangeProofU64 => {
270            invoke_context
271                .consume_checked(VERIFY_BATCHED_RANGE_PROOF_U64_COMPUTE_UNITS)
272                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
273            ic_msg!(invoke_context, "VerifyBatchedRangeProof64");
274            process_verify_proof::<BatchedRangeProofU64Data, BatchedRangeProofContext>(
275                invoke_context,
276            )
277        }
278        ProofInstruction::VerifyBatchedRangeProofU128 => {
279            invoke_context
280                .consume_checked(VERIFY_BATCHED_RANGE_PROOF_U128_COMPUTE_UNITS)
281                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
282            ic_msg!(invoke_context, "VerifyBatchedRangeProof128");
283            process_verify_proof::<BatchedRangeProofU128Data, BatchedRangeProofContext>(
284                invoke_context,
285            )
286        }
287        ProofInstruction::VerifyBatchedRangeProofU256 => {
288            // transfer with fee related proofs are not enabled
289            if !enable_zk_transfer_with_fee {
290                return Err(InstructionError::InvalidInstructionData);
291            }
292
293            invoke_context
294                .consume_checked(VERIFY_BATCHED_RANGE_PROOF_U256_COMPUTE_UNITS)
295                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
296            ic_msg!(invoke_context, "VerifyBatchedRangeProof256");
297            process_verify_proof::<BatchedRangeProofU256Data, BatchedRangeProofContext>(
298                invoke_context,
299            )
300        }
301        ProofInstruction::VerifyCiphertextCommitmentEquality => {
302            invoke_context
303                .consume_checked(VERIFY_CIPHERTEXT_COMMITMENT_EQUALITY_COMPUTE_UNITS)
304                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
305            ic_msg!(invoke_context, "VerifyCiphertextCommitmentEquality");
306            process_verify_proof::<
307                CiphertextCommitmentEqualityProofData,
308                CiphertextCommitmentEqualityProofContext,
309            >(invoke_context)
310        }
311        ProofInstruction::VerifyGroupedCiphertext2HandlesValidity => {
312            invoke_context
313                .consume_checked(VERIFY_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_COMPUTE_UNITS)
314                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
315            ic_msg!(invoke_context, "VerifyGroupedCiphertext2HandlesValidity");
316            process_verify_proof::<
317                GroupedCiphertext2HandlesValidityProofData,
318                GroupedCiphertext2HandlesValidityProofContext,
319            >(invoke_context)
320        }
321        ProofInstruction::VerifyBatchedGroupedCiphertext2HandlesValidity => {
322            invoke_context
323                .consume_checked(VERIFY_BATCHED_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_COMPUTE_UNITS)
324                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
325            ic_msg!(
326                invoke_context,
327                "VerifyBatchedGroupedCiphertext2HandlesValidity"
328            );
329            process_verify_proof::<
330                BatchedGroupedCiphertext2HandlesValidityProofData,
331                BatchedGroupedCiphertext2HandlesValidityProofContext,
332            >(invoke_context)
333        }
334        ProofInstruction::VerifyFeeSigma => {
335            // transfer with fee related proofs are not enabled
336            if !enable_zk_transfer_with_fee {
337                return Err(InstructionError::InvalidInstructionData);
338            }
339
340            invoke_context
341                .consume_checked(VERIFY_FEE_SIGMA_COMPUTE_UNITS)
342                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
343            ic_msg!(invoke_context, "VerifyFeeSigma");
344            process_verify_proof::<FeeSigmaProofData, FeeSigmaProofContext>(invoke_context)
345        }
346        ProofInstruction::VerifyGroupedCiphertext3HandlesValidity => {
347            invoke_context
348                .consume_checked(VERIFY_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_COMPUTE_UNITS)
349                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
350            ic_msg!(invoke_context, "VerifyGroupedCiphertext3HandlesValidity");
351            process_verify_proof::<
352                GroupedCiphertext3HandlesValidityProofData,
353                GroupedCiphertext3HandlesValidityProofContext,
354            >(invoke_context)
355        }
356        ProofInstruction::VerifyBatchedGroupedCiphertext3HandlesValidity => {
357            invoke_context
358                .consume_checked(VERIFY_BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_COMPUTE_UNITS)
359                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
360            ic_msg!(
361                invoke_context,
362                "VerifyBatchedGroupedCiphertext3HandlesValidity"
363            );
364            process_verify_proof::<
365                BatchedGroupedCiphertext3HandlesValidityProofData,
366                BatchedGroupedCiphertext3HandlesValidityProofContext,
367            >(invoke_context)
368        }
369    }
370});
solana_zk_token_proof_program/lib.rs

solana_zk_token_proof_program/
lib.rs
