spl_token_confidential_transfer_proof_extraction/
instruction.rs

1//! Utility functions to simplify the handling of ZK ElGamal proof program
2//! instruction data in SPL crates
3
4use {
5    bytemuck::Pod,
6    solana_account_info::{next_account_info, AccountInfo},
7    solana_instruction::{AccountMeta, Instruction},
8    solana_instructions_sysvar::get_instruction_relative,
9    solana_msg::msg,
10    solana_program_error::{ProgramError, ProgramResult},
11    solana_pubkey::Pubkey,
12    solana_zk_sdk::zk_elgamal_proof_program::{
13        self,
14        instruction::ProofInstruction,
15        proof_data::{ProofType, ZkProofData},
16        state::ProofContextState,
17    },
18    spl_pod::bytemuck::pod_from_bytes,
19    std::{num::NonZeroI8, slice::Iter},
20};
21
22/// Checks that the supplied program ID is correct for the ZK ElGamal proof
23/// program
24pub fn check_zk_elgamal_proof_program_account(
25    zk_elgamal_proof_program_id: &Pubkey,
26) -> ProgramResult {
27    if zk_elgamal_proof_program_id != &solana_zk_sdk::zk_elgamal_proof_program::id() {
28        return Err(ProgramError::IncorrectProgramId);
29    }
30    Ok(())
31}
32
33/// Decodes the proof context data associated with a zero-knowledge proof
34/// instruction.
35pub fn decode_proof_instruction_context<T: Pod + ZkProofData<U>, U: Pod>(
36    expected: ProofInstruction,
37    instruction: &Instruction,
38) -> Result<U, ProgramError> {
39    if instruction.program_id != zk_elgamal_proof_program::id()
40        || ProofInstruction::instruction_type(&instruction.data) != Some(expected)
41    {
42        msg!("Unexpected proof instruction");
43        return Err(ProgramError::InvalidInstructionData);
44    }
45    ProofInstruction::proof_data::<T, U>(&instruction.data)
46        .map(|proof_data| *ZkProofData::context_data(proof_data))
47        .ok_or(ProgramError::InvalidInstructionData)
48}
49
50/// A proof location type meant to be used for arguments to instruction
51/// constructors.
52#[derive(Clone, Copy)]
53pub enum ProofLocation<'a, T> {
54    /// The proof is included in the same transaction of a corresponding
55    /// token-2022 instruction.
56    InstructionOffset(NonZeroI8, &'a T),
57    /// The proof is pre-verified into a context state account.
58    ContextStateAccount(&'a Pubkey),
59}
60
61impl<T> ProofLocation<'_, T> {
62    /// Returns true if the proof location is an instruction offset
63    pub fn is_instruction_offset(&self) -> bool {
64        match self {
65            Self::InstructionOffset(_, _) => true,
66            Self::ContextStateAccount(_) => false,
67        }
68    }
69}
70
71/// Verify zero-knowledge proof and return the corresponding proof context.
72pub fn verify_and_extract_context<'a, T: Pod + ZkProofData<U>, U: Pod>(
73    account_info_iter: &mut Iter<'_, AccountInfo<'a>>,
74    proof_instruction_offset: i64,
75    sysvar_account_info: Option<&'_ AccountInfo<'a>>,
76) -> Result<U, ProgramError> {
77    if proof_instruction_offset == 0 {
78        // interpret `account_info` as a context state account
79        let context_state_account_info = next_account_info(account_info_iter)?;
80        check_zk_elgamal_proof_program_account(context_state_account_info.owner)?;
81        let context_state_account_data = context_state_account_info.data.borrow();
82        let context_state = pod_from_bytes::<ProofContextState<U>>(&context_state_account_data)?;
83
84        if context_state.proof_type != T::PROOF_TYPE.into() {
85            return Err(ProgramError::InvalidInstructionData);
86        }
87
88        Ok(context_state.proof_context)
89    } else {
90        // if sysvar account is not provided, then get the sysvar account
91        let sysvar_account_info = if let Some(sysvar_account_info) = sysvar_account_info {
92            sysvar_account_info
93        } else {
94            next_account_info(account_info_iter)?
95        };
96        let zkp_instruction =
97            get_instruction_relative(proof_instruction_offset, sysvar_account_info)?;
98        let expected_proof_type = zk_proof_type_to_instruction(T::PROOF_TYPE)?;
99        Ok(decode_proof_instruction_context::<T, U>(
100            expected_proof_type,
101            &zkp_instruction,
102        )?)
103    }
104}
105
106/// Processes a proof location for instruction creation. Adds relevant accounts
107/// to supplied account vector
108///
109/// If the proof location is an instruction offset the corresponding proof
110/// instruction is created and added to the `proof_instructions` vector.
111pub fn process_proof_location<T, U>(
112    accounts: &mut Vec<AccountMeta>,
113    expected_instruction_offset: &mut i8,
114    proof_instructions: &mut Vec<Instruction>,
115    proof_location: ProofLocation<T>,
116    push_sysvar_to_accounts: bool,
117    proof_instruction_type: ProofInstruction,
118) -> Result<i8, ProgramError>
119where
120    T: Pod + ZkProofData<U>,
121    U: Pod,
122{
123    match proof_location {
124        ProofLocation::InstructionOffset(proof_instruction_offset, proof_data) => {
125            let proof_instruction_offset: i8 = proof_instruction_offset.into();
126            if &proof_instruction_offset != expected_instruction_offset {
127                return Err(ProgramError::InvalidInstructionData);
128            }
129
130            if push_sysvar_to_accounts {
131                accounts.push(AccountMeta::new_readonly(
132                    solana_sdk_ids::sysvar::instructions::id(),
133                    false,
134                ));
135            }
136            proof_instructions
137                .push(proof_instruction_type.encode_verify_proof::<T, U>(None, proof_data));
138            *expected_instruction_offset = expected_instruction_offset
139                .checked_add(1)
140                .ok_or(ProgramError::InvalidInstructionData)?;
141            Ok(proof_instruction_offset)
142        }
143        ProofLocation::ContextStateAccount(context_state_account) => {
144            accounts.push(AccountMeta::new_readonly(*context_state_account, false));
145            Ok(0)
146        }
147    }
148}
149
150/// Converts a zk proof type to a corresponding ZK ElGamal proof program
151/// instruction that verifies the proof.
152pub fn zk_proof_type_to_instruction(
153    proof_type: ProofType,
154) -> Result<ProofInstruction, ProgramError> {
155    match proof_type {
156        ProofType::ZeroCiphertext => Ok(ProofInstruction::VerifyZeroCiphertext),
157        ProofType::CiphertextCiphertextEquality => {
158            Ok(ProofInstruction::VerifyCiphertextCiphertextEquality)
159        }
160        ProofType::PubkeyValidity => Ok(ProofInstruction::VerifyPubkeyValidity),
161        ProofType::BatchedRangeProofU64 => Ok(ProofInstruction::VerifyBatchedRangeProofU64),
162        ProofType::BatchedRangeProofU128 => Ok(ProofInstruction::VerifyBatchedRangeProofU128),
163        ProofType::BatchedRangeProofU256 => Ok(ProofInstruction::VerifyBatchedRangeProofU256),
164        ProofType::CiphertextCommitmentEquality => {
165            Ok(ProofInstruction::VerifyCiphertextCommitmentEquality)
166        }
167        ProofType::GroupedCiphertext2HandlesValidity => {
168            Ok(ProofInstruction::VerifyGroupedCiphertext2HandlesValidity)
169        }
170        ProofType::BatchedGroupedCiphertext2HandlesValidity => {
171            Ok(ProofInstruction::VerifyBatchedGroupedCiphertext2HandlesValidity)
172        }
173        ProofType::PercentageWithCap => Ok(ProofInstruction::VerifyPercentageWithCap),
174        ProofType::GroupedCiphertext3HandlesValidity => {
175            Ok(ProofInstruction::VerifyGroupedCiphertext3HandlesValidity)
176        }
177        ProofType::BatchedGroupedCiphertext3HandlesValidity => {
178            Ok(ProofInstruction::VerifyBatchedGroupedCiphertext3HandlesValidity)
179        }
180        ProofType::Uninitialized => Err(ProgramError::InvalidInstructionData),
181    }
182}
spl_token_confidential_transfer_proof_extraction/instruction.rs

spl_token_confidential_transfer_proof_extraction/
instruction.rs
