miraland_zk_token_proof_program/
lib.rs

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

miraland_zk_token_proof_program/
lib.rs
