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

solana_zk_token_proof_program/
lib.rs
