use {
crate::check_zk_elgamal_proof_program_account,
bytemuck::Pod,
solana_program::{
account_info::{next_account_info, AccountInfo},
instruction::Instruction,
msg,
program_error::ProgramError,
pubkey::Pubkey,
sysvar::instructions::get_instruction_relative,
},
solana_zk_sdk::zk_elgamal_proof_program::{
self,
instruction::ProofInstruction,
proof_data::{ProofType, ZkProofData},
state::ProofContextState,
},
spl_pod::bytemuck::pod_from_bytes,
std::{num::NonZeroI8, slice::Iter},
};
const INSTRUCTION_DATA_LENGTH_WITH_RECORD_ACCOUNT: usize = 5;
pub fn decode_proof_instruction_context<T: Pod + ZkProofData<U>, U: Pod>(
account_info_iter: &mut Iter<'_, AccountInfo<'_>>,
expected: ProofInstruction,
instruction: &Instruction,
) -> Result<U, ProgramError> {
if instruction.program_id != zk_elgamal_proof_program::id()
|| ProofInstruction::instruction_type(&instruction.data) != Some(expected)
{
msg!("Unexpected proof instruction");
return Err(ProgramError::InvalidInstructionData);
}
if instruction.data.len() == INSTRUCTION_DATA_LENGTH_WITH_RECORD_ACCOUNT {
let record_account = next_account_info(account_info_iter)?;
let start_offset = u32::from_le_bytes(instruction.data[1..].try_into().unwrap()) as usize;
let end_offset = start_offset
.checked_add(std::mem::size_of::<T>())
.ok_or(ProgramError::InvalidAccountData)?;
let record_account_data = record_account.data.borrow();
let raw_proof_data = record_account_data
.get(start_offset..end_offset)
.ok_or(ProgramError::AccountDataTooSmall)?;
bytemuck::try_from_bytes::<T>(raw_proof_data)
.map(|proof_data| *ZkProofData::context_data(proof_data))
.map_err(|_| ProgramError::InvalidAccountData)
} else {
ProofInstruction::proof_data::<T, U>(&instruction.data)
.map(|proof_data| *ZkProofData::context_data(proof_data))
.ok_or(ProgramError::InvalidInstructionData)
}
}
#[derive(Clone, Copy)]
pub enum ProofLocation<'a, T> {
InstructionOffset(NonZeroI8, ProofData<'a, T>),
ContextStateAccount(&'a Pubkey),
}
impl<'a, T> ProofLocation<'a, T> {
pub fn is_instruction_offset(&self) -> bool {
match self {
Self::InstructionOffset(_, _) => true,
Self::ContextStateAccount(_) => false,
}
}
}
#[derive(Clone, Copy)]
pub enum ProofData<'a, T> {
InstructionData(&'a T),
RecordAccount(&'a Pubkey, u32),
}
pub fn verify_and_extract_context<'a, T: Pod + ZkProofData<U>, U: Pod>(
account_info_iter: &mut Iter<'_, AccountInfo<'a>>,
proof_instruction_offset: i64,
sysvar_account_info: Option<&'_ AccountInfo<'a>>,
) -> Result<U, ProgramError> {
if proof_instruction_offset == 0 {
let context_state_account_info = next_account_info(account_info_iter)?;
check_zk_elgamal_proof_program_account(context_state_account_info.owner)?;
let context_state_account_data = context_state_account_info.data.borrow();
let context_state = pod_from_bytes::<ProofContextState<U>>(&context_state_account_data)?;
if context_state.proof_type != T::PROOF_TYPE.into() {
return Err(ProgramError::InvalidInstructionData);
}
Ok(context_state.proof_context)
} else {
let sysvar_account_info = if let Some(sysvar_account_info) = sysvar_account_info {
sysvar_account_info
} else {
next_account_info(account_info_iter)?
};
let zkp_instruction =
get_instruction_relative(proof_instruction_offset, sysvar_account_info)?;
let expected_proof_type = zk_proof_type_to_instruction(T::PROOF_TYPE)?;
Ok(decode_proof_instruction_context::<T, U>(
account_info_iter,
expected_proof_type,
&zkp_instruction,
)?)
}
}
pub fn zk_proof_type_to_instruction(
proof_type: ProofType,
) -> Result<ProofInstruction, ProgramError> {
match proof_type {
ProofType::ZeroCiphertext => Ok(ProofInstruction::VerifyZeroCiphertext),
ProofType::CiphertextCiphertextEquality => {
Ok(ProofInstruction::VerifyCiphertextCiphertextEquality)
}
ProofType::PubkeyValidity => Ok(ProofInstruction::VerifyPubkeyValidity),
ProofType::BatchedRangeProofU64 => Ok(ProofInstruction::VerifyBatchedRangeProofU64),
ProofType::BatchedRangeProofU128 => Ok(ProofInstruction::VerifyBatchedRangeProofU128),
ProofType::BatchedRangeProofU256 => Ok(ProofInstruction::VerifyBatchedRangeProofU256),
ProofType::CiphertextCommitmentEquality => {
Ok(ProofInstruction::VerifyCiphertextCommitmentEquality)
}
ProofType::GroupedCiphertext2HandlesValidity => {
Ok(ProofInstruction::VerifyGroupedCiphertext2HandlesValidity)
}
ProofType::BatchedGroupedCiphertext2HandlesValidity => {
Ok(ProofInstruction::VerifyBatchedGroupedCiphertext2HandlesValidity)
}
ProofType::PercentageWithCap => Ok(ProofInstruction::VerifyPercentageWithCap),
ProofType::GroupedCiphertext3HandlesValidity => {
Ok(ProofInstruction::VerifyGroupedCiphertext3HandlesValidity)
}
ProofType::BatchedGroupedCiphertext3HandlesValidity => {
Ok(ProofInstruction::VerifyBatchedGroupedCiphertext3HandlesValidity)
}
ProofType::Uninitialized => Err(ProgramError::InvalidInstructionData),
}
}
#[derive(Clone, Copy)]
pub struct SplitContextStateAccountsConfig {
pub no_op_on_uninitialized_split_context_state: bool,
pub close_split_context_state_on_execution: bool,
}