gorb_ct_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_program::{
7        account_info::{next_account_info, AccountInfo},
8        entrypoint::ProgramResult,
9        instruction::{AccountMeta, Instruction},
10        msg,
11        program_error::ProgramError,
12        pubkey::Pubkey,
13        sysvar::{self, instructions::get_instruction_relative},
14    },
15    solana_zk_sdk::zk_elgamal_proof_program::{
16        self,
17        instruction::ProofInstruction,
18        proof_data::{ProofType, ZkProofData},
19        state::ProofContextState,
20    },
21    gorb_pod::bytemuck::pod_from_bytes,
22    std::{num::NonZeroI8, slice::Iter},
23};
24
25/// Checks that the supplied program ID is correct for the ZK ElGamal proof
26/// program
27pub fn check_zk_elgamal_proof_program_account(
28    zk_elgamal_proof_program_id: &Pubkey,
29) -> ProgramResult {
30    if zk_elgamal_proof_program_id != &solana_zk_sdk::zk_elgamal_proof_program::id() {
31        return Err(ProgramError::IncorrectProgramId);
32    }
33    Ok(())
34}
35
36/// If a proof is to be read from a record account, the proof instruction data
37/// must be 5 bytes: 1 byte for the proof type and 4 bytes for the `u32` offset
38const INSTRUCTION_DATA_LENGTH_WITH_RECORD_ACCOUNT: usize = 5;
39
40/// Decodes the proof context data associated with a zero-knowledge proof
41/// instruction.
42pub fn decode_proof_instruction_context<T: Pod + ZkProofData<U>, U: Pod>(
43    account_info_iter: &mut Iter<'_, AccountInfo<'_>>,
44    expected: ProofInstruction,
45    instruction: &Instruction,
46) -> Result<U, ProgramError> {
47    if instruction.program_id != zk_elgamal_proof_program::id()
48        || ProofInstruction::instruction_type(&instruction.data) != Some(expected)
49    {
50        msg!("Unexpected proof instruction");
51        return Err(ProgramError::InvalidInstructionData);
52    }
53
54    // If the instruction data size is exactly 5 bytes, then interpret it as an
55    // offset byte for a record account. This behavior is identical to that of
56    // the ZK ElGamal proof program.
57    if instruction.data.len() == INSTRUCTION_DATA_LENGTH_WITH_RECORD_ACCOUNT {
58        let record_account = next_account_info(account_info_iter)?;
59
60        // first byte is the proof type
61        let start_offset = u32::from_le_bytes(instruction.data[1..].try_into().unwrap()) as usize;
62        let end_offset = start_offset
63            .checked_add(std::mem::size_of::<T>())
64            .ok_or(ProgramError::InvalidAccountData)?;
65
66        let record_account_data = record_account.data.borrow();
67        let raw_proof_data = record_account_data
68            .get(start_offset..end_offset)
69            .ok_or(ProgramError::AccountDataTooSmall)?;
70
71        bytemuck::try_from_bytes::<T>(raw_proof_data)
72            .map(|proof_data| *ZkProofData::context_data(proof_data))
73            .map_err(|_| ProgramError::InvalidAccountData)
74    } else {
75        ProofInstruction::proof_data::<T, U>(&instruction.data)
76            .map(|proof_data| *ZkProofData::context_data(proof_data))
77            .ok_or(ProgramError::InvalidInstructionData)
78    }
79}
80
81/// A proof location type meant to be used for arguments to instruction
82/// constructors.
83#[derive(Clone, Copy)]
84pub enum ProofLocation<'a, T> {
85    /// The proof is included in the same transaction of a corresponding
86    /// token-2022 instruction.
87    InstructionOffset(NonZeroI8, ProofData<'a, T>),
88    /// The proof is pre-verified into a context state account.
89    ContextStateAccount(&'a Pubkey),
90}
91
92impl<'a, T> ProofLocation<'a, T> {
93    /// Returns true if the proof location is an instruction offset
94    pub fn is_instruction_offset(&self) -> bool {
95        match self {
96            Self::InstructionOffset(_, _) => true,
97            Self::ContextStateAccount(_) => false,
98        }
99    }
100}
101
102/// A proof data type to distinguish between proof data included as part of
103/// zk-token proof instruction data and proof data stored in a record account.
104#[derive(Clone, Copy)]
105pub enum ProofData<'a, T> {
106    /// The proof data
107    InstructionData(&'a T),
108    /// The address of a record account containing the proof data and its byte
109    /// offset
110    RecordAccount(&'a Pubkey, u32),
111}
112
113/// Verify zero-knowledge proof and return the corresponding proof context.
114pub fn verify_and_extract_context<'a, T: Pod + ZkProofData<U>, U: Pod>(
115    account_info_iter: &mut Iter<'_, AccountInfo<'a>>,
116    proof_instruction_offset: i64,
117    sysvar_account_info: Option<&'_ AccountInfo<'a>>,
118) -> Result<U, ProgramError> {
119    if proof_instruction_offset == 0 {
120        // interpret `account_info` as a context state account
121        let context_state_account_info = next_account_info(account_info_iter)?;
122        check_zk_elgamal_proof_program_account(context_state_account_info.owner)?;
123        let context_state_account_data = context_state_account_info.data.borrow();
124        let context_state = pod_from_bytes::<ProofContextState<U>>(&context_state_account_data)?;
125
126        if context_state.proof_type != T::PROOF_TYPE.into() {
127            return Err(ProgramError::InvalidInstructionData);
128        }
129
130        Ok(context_state.proof_context)
131    } else {
132        // if sysvar account is not provided, then get the sysvar account
133        let sysvar_account_info = if let Some(sysvar_account_info) = sysvar_account_info {
134            sysvar_account_info
135        } else {
136            next_account_info(account_info_iter)?
137        };
138        let zkp_instruction =
139            get_instruction_relative(proof_instruction_offset, sysvar_account_info)?;
140        let expected_proof_type = zk_proof_type_to_instruction(T::PROOF_TYPE)?;
141        Ok(decode_proof_instruction_context::<T, U>(
142            account_info_iter,
143            expected_proof_type,
144            &zkp_instruction,
145        )?)
146    }
147}
148
149/// Processes a proof location for instruction creation. Adds relevant accounts
150/// to supplied account vector
151///
152/// If the proof location is an instruction offset the corresponding proof
153/// instruction is created and added to the `proof_instructions` vector.
154pub fn process_proof_location<T, U>(
155    accounts: &mut Vec<AccountMeta>,
156    expected_instruction_offset: &mut i8,
157    proof_instructions: &mut Vec<Instruction>,
158    proof_location: ProofLocation<T>,
159    push_sysvar_to_accounts: bool,
160    proof_instruction_type: ProofInstruction,
161) -> Result<i8, ProgramError>
162where
163    T: Pod + ZkProofData<U>,
164    U: Pod,
165{
166    match proof_location {
167        ProofLocation::InstructionOffset(proof_instruction_offset, proof_data) => {
168            let proof_instruction_offset: i8 = proof_instruction_offset.into();
169            if &proof_instruction_offset != expected_instruction_offset {
170                return Err(ProgramError::InvalidInstructionData);
171            }
172
173            if push_sysvar_to_accounts {
174                accounts.push(AccountMeta::new_readonly(sysvar::instructions::id(), false));
175            }
176            match proof_data {
177                ProofData::InstructionData(data) => proof_instructions
178                    .push(proof_instruction_type.encode_verify_proof::<T, U>(None, data)),
179                ProofData::RecordAccount(address, offset) => {
180                    accounts.push(AccountMeta::new_readonly(*address, false));
181                    proof_instructions.push(
182                        proof_instruction_type
183                            .encode_verify_proof_from_account(None, address, offset),
184                    )
185                }
186            };
187            *expected_instruction_offset = expected_instruction_offset
188                .checked_add(1)
189                .ok_or(ProgramError::InvalidInstructionData)?;
190            Ok(proof_instruction_offset)
191        }
192        ProofLocation::ContextStateAccount(context_state_account) => {
193            accounts.push(AccountMeta::new_readonly(*context_state_account, false));
194            Ok(0)
195        }
196    }
197}
198
199/// Converts a zk proof type to a corresponding ZK ElGamal proof program
200/// instruction that verifies the proof.
201pub fn zk_proof_type_to_instruction(
202    proof_type: ProofType,
203) -> Result<ProofInstruction, ProgramError> {
204    match proof_type {
205        ProofType::ZeroCiphertext => Ok(ProofInstruction::VerifyZeroCiphertext),
206        ProofType::CiphertextCiphertextEquality => {
207            Ok(ProofInstruction::VerifyCiphertextCiphertextEquality)
208        }
209        ProofType::PubkeyValidity => Ok(ProofInstruction::VerifyPubkeyValidity),
210        ProofType::BatchedRangeProofU64 => Ok(ProofInstruction::VerifyBatchedRangeProofU64),
211        ProofType::BatchedRangeProofU128 => Ok(ProofInstruction::VerifyBatchedRangeProofU128),
212        ProofType::BatchedRangeProofU256 => Ok(ProofInstruction::VerifyBatchedRangeProofU256),
213        ProofType::CiphertextCommitmentEquality => {
214            Ok(ProofInstruction::VerifyCiphertextCommitmentEquality)
215        }
216        ProofType::GroupedCiphertext2HandlesValidity => {
217            Ok(ProofInstruction::VerifyGroupedCiphertext2HandlesValidity)
218        }
219        ProofType::BatchedGroupedCiphertext2HandlesValidity => {
220            Ok(ProofInstruction::VerifyBatchedGroupedCiphertext2HandlesValidity)
221        }
222        ProofType::PercentageWithCap => Ok(ProofInstruction::VerifyPercentageWithCap),
223        ProofType::GroupedCiphertext3HandlesValidity => {
224            Ok(ProofInstruction::VerifyGroupedCiphertext3HandlesValidity)
225        }
226        ProofType::BatchedGroupedCiphertext3HandlesValidity => {
227            Ok(ProofInstruction::VerifyBatchedGroupedCiphertext3HandlesValidity)
228        }
229        ProofType::Uninitialized => Err(ProgramError::InvalidInstructionData),
230    }
231}
gorb_ct_proof_extraction/instruction.rs

gorb_ct_proof_extraction/
instruction.rs
