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