|
| 1 | +//! Utility functions to simplify the handling of ZK ElGamal proof program |
| 2 | +//! instruction data in SPL crates |
| 3 | +
|
| 4 | +use { |
| 5 | + bytemuck::Pod, |
| 6 | + solana_program::{ |
| 7 | + account_info::{next_account_info, AccountInfo}, |
| 8 | + entrypoint::ProgramResult, |
| 9 | + instruction::Instruction, |
| 10 | + msg, |
| 11 | + program_error::ProgramError, |
| 12 | + pubkey::Pubkey, |
| 13 | + sysvar::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 | + spl_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 |
| 27 | +pub 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 |
| 38 | +const INSTRUCTION_DATA_LENGTH_WITH_RECORD_ACCOUNT: usize = 5; |
| 39 | + |
| 40 | +/// Decodes the proof context data associated with a zero-knowledge proof |
| 41 | +/// instruction. |
| 42 | +pub 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)] |
| 84 | +pub 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 | + |
| 92 | +impl<'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)] |
| 105 | +pub 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. |
| 114 | +pub 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 | +/// Converts a zk proof type to a corresponding ZK ElGamal proof program |
| 150 | +/// instruction that verifies the proof. |
| 151 | +pub fn zk_proof_type_to_instruction( |
| 152 | + proof_type: ProofType, |
| 153 | +) -> Result<ProofInstruction, ProgramError> { |
| 154 | + match proof_type { |
| 155 | + ProofType::ZeroCiphertext => Ok(ProofInstruction::VerifyZeroCiphertext), |
| 156 | + ProofType::CiphertextCiphertextEquality => { |
| 157 | + Ok(ProofInstruction::VerifyCiphertextCiphertextEquality) |
| 158 | + } |
| 159 | + ProofType::PubkeyValidity => Ok(ProofInstruction::VerifyPubkeyValidity), |
| 160 | + ProofType::BatchedRangeProofU64 => Ok(ProofInstruction::VerifyBatchedRangeProofU64), |
| 161 | + ProofType::BatchedRangeProofU128 => Ok(ProofInstruction::VerifyBatchedRangeProofU128), |
| 162 | + ProofType::BatchedRangeProofU256 => Ok(ProofInstruction::VerifyBatchedRangeProofU256), |
| 163 | + ProofType::CiphertextCommitmentEquality => { |
| 164 | + Ok(ProofInstruction::VerifyCiphertextCommitmentEquality) |
| 165 | + } |
| 166 | + ProofType::GroupedCiphertext2HandlesValidity => { |
| 167 | + Ok(ProofInstruction::VerifyGroupedCiphertext2HandlesValidity) |
| 168 | + } |
| 169 | + ProofType::BatchedGroupedCiphertext2HandlesValidity => { |
| 170 | + Ok(ProofInstruction::VerifyBatchedGroupedCiphertext2HandlesValidity) |
| 171 | + } |
| 172 | + ProofType::PercentageWithCap => Ok(ProofInstruction::VerifyPercentageWithCap), |
| 173 | + ProofType::GroupedCiphertext3HandlesValidity => { |
| 174 | + Ok(ProofInstruction::VerifyGroupedCiphertext3HandlesValidity) |
| 175 | + } |
| 176 | + ProofType::BatchedGroupedCiphertext3HandlesValidity => { |
| 177 | + Ok(ProofInstruction::VerifyBatchedGroupedCiphertext3HandlesValidity) |
| 178 | + } |
| 179 | + ProofType::Uninitialized => Err(ProgramError::InvalidInstructionData), |
| 180 | + } |
| 181 | +} |
0 commit comments