safe_zk_token_proof_program/
lib.rs

1#![forbid(unsafe_code)]
2
3use {
4    bytemuck::Pod,
5    solana_program_runtime::{ic_msg, invoke_context::InvokeContext},
6    solana_sdk::instruction::{InstructionError, TRANSACTION_LEVEL_STACK_HEIGHT},
7    safe_zk_token_sdk::zk_token_proof_instruction::*,
8    std::result::Result,
9};
10
11fn verify<T: Pod + Verifiable>(invoke_context: &mut InvokeContext) -> Result<(), InstructionError> {
12    let transaction_context = &invoke_context.transaction_context;
13    let instruction_context = transaction_context.get_current_instruction_context()?;
14    let instruction_data = instruction_context.get_instruction_data();
15    let instruction = ProofInstruction::decode_data::<T>(instruction_data);
16
17    let proof = instruction.ok_or_else(|| {
18        ic_msg!(invoke_context, "invalid proof data");
19        InstructionError::InvalidInstructionData
20    })?;
21
22    proof.verify().map_err(|err| {
23        ic_msg!(invoke_context, "proof verification failed: {:?}", err);
24        InstructionError::InvalidInstructionData
25    })
26}
27
28pub fn process_instruction(
29    _first_instruction_account: usize,
30    invoke_context: &mut InvokeContext,
31) -> Result<(), InstructionError> {
32    if invoke_context.get_stack_height() != TRANSACTION_LEVEL_STACK_HEIGHT {
33        // Not supported as an inner instruction
34        return Err(InstructionError::UnsupportedProgramId);
35    }
36
37    // Consume compute units since proof verification is an expensive operation
38    {
39        let compute_meter = invoke_context.get_compute_meter();
40        // TODO: Tune the number of units consumed.  The current value is just a rough estimate
41        compute_meter.borrow_mut().consume(100_000)?;
42    }
43
44    let transaction_context = &invoke_context.transaction_context;
45    let instruction_context = transaction_context.get_current_instruction_context()?;
46    let instruction_data = instruction_context.get_instruction_data();
47    let instruction = ProofInstruction::decode_type(instruction_data);
48
49    match instruction.ok_or(InstructionError::InvalidInstructionData)? {
50        ProofInstruction::VerifyCloseAccount => {
51            ic_msg!(invoke_context, "VerifyCloseAccount");
52            verify::<CloseAccountData>(invoke_context)
53        }
54        ProofInstruction::VerifyWithdraw => {
55            ic_msg!(invoke_context, "VerifyWithdraw");
56            verify::<WithdrawData>(invoke_context)
57        }
58        ProofInstruction::VerifyWithdrawWithheldTokens => {
59            ic_msg!(invoke_context, "VerifyWithdrawWithheldTokens");
60            verify::<WithdrawWithheldTokensData>(invoke_context)
61        }
62        ProofInstruction::VerifyTransfer => {
63            ic_msg!(invoke_context, "VerifyTransfer");
64            verify::<TransferData>(invoke_context)
65        }
66        ProofInstruction::VerifyTransferWithFee => {
67            ic_msg!(invoke_context, "VerifyTransferWithFee");
68            verify::<TransferWithFeeData>(invoke_context)
69        }
70        ProofInstruction::VerifyPubkeyValidity => {
71            ic_msg!(invoke_context, "VerifyPubkeyValidity");
72            verify::<PubkeyValidityData>(invoke_context)
73        }
74    }
75}