solana_zk_elgamal_proof_program/
lib.rs

1#![forbid(unsafe_code)]
2
3use {
4    bytemuck::Pod,
5    solana_log_collector::ic_msg,
6    solana_program_runtime::{declare_process_instruction, invoke_context::InvokeContext},
7    solana_sdk::{instruction::InstructionError, system_program},
8    solana_zk_sdk::zk_elgamal_proof_program::{
9        id,
10        instruction::ProofInstruction,
11        proof_data::*,
12        state::{ProofContextState, ProofContextStateMeta},
13    },
14    std::result::Result,
15};
16
17pub const CLOSE_CONTEXT_STATE_COMPUTE_UNITS: u64 = 3_300;
18pub const VERIFY_ZERO_CIPHERTEXT_COMPUTE_UNITS: u64 = 6_000;
19pub const VERIFY_CIPHERTEXT_CIPHERTEXT_EQUALITY_COMPUTE_UNITS: u64 = 8_000;
20pub const VERIFY_CIPHERTEXT_COMMITMENT_EQUALITY_COMPUTE_UNITS: u64 = 6_400;
21pub const VERIFY_PUBKEY_VALIDITY_COMPUTE_UNITS: u64 = 2_600;
22pub const VERIFY_PERCENTAGE_WITH_CAP_COMPUTE_UNITS: u64 = 6_500;
23pub const VERIFY_BATCHED_RANGE_PROOF_U64_COMPUTE_UNITS: u64 = 111_000;
24pub const VERIFY_BATCHED_RANGE_PROOF_U128_COMPUTE_UNITS: u64 = 200_000;
25pub const VERIFY_BATCHED_RANGE_PROOF_U256_COMPUTE_UNITS: u64 = 368_000;
26pub const VERIFY_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_COMPUTE_UNITS: u64 = 6_400;
27pub const VERIFY_BATCHED_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_COMPUTE_UNITS: u64 = 13_000;
28pub const VERIFY_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_COMPUTE_UNITS: u64 = 8_100;
29pub const VERIFY_BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_COMPUTE_UNITS: u64 = 16_400;
30
31const INSTRUCTION_DATA_LENGTH_WITH_PROOF_ACCOUNT: usize = 5;
32
33fn process_verify_proof<T, U>(invoke_context: &mut InvokeContext) -> Result<(), InstructionError>
34where
35    T: Pod + ZkProofData<U>,
36    U: Pod,
37{
38    let transaction_context = &invoke_context.transaction_context;
39    let instruction_context = transaction_context.get_current_instruction_context()?;
40    let instruction_data = instruction_context.get_instruction_data();
41
42    // number of accessed accounts so far
43    let mut accessed_accounts = 0_u16;
44
45    // if instruction data is exactly 5 bytes, then read proof from an account
46    let context_data = if instruction_data.len() == INSTRUCTION_DATA_LENGTH_WITH_PROOF_ACCOUNT {
47        let proof_data_account = instruction_context
48            .try_borrow_instruction_account(transaction_context, accessed_accounts)?;
49        accessed_accounts = accessed_accounts.checked_add(1).unwrap();
50
51        let proof_data_offset = u32::from_le_bytes(
52            // the first byte is the instruction discriminator
53            instruction_data[1..INSTRUCTION_DATA_LENGTH_WITH_PROOF_ACCOUNT]
54                .try_into()
55                .map_err(|_| InstructionError::InvalidInstructionData)?,
56        );
57        let proof_data_start: usize = proof_data_offset
58            .try_into()
59            .map_err(|_| InstructionError::InvalidInstructionData)?;
60        let proof_data_end = proof_data_start
61            .checked_add(std::mem::size_of::<T>())
62            .ok_or(InstructionError::InvalidInstructionData)?;
63        let proof_data_raw = proof_data_account
64            .get_data()
65            .get(proof_data_start..proof_data_end)
66            .ok_or(InstructionError::InvalidAccountData)?;
67
68        let proof_data = bytemuck::try_from_bytes::<T>(proof_data_raw).map_err(|_| {
69            ic_msg!(invoke_context, "invalid proof data");
70            InstructionError::InvalidInstructionData
71        })?;
72        proof_data.verify_proof().map_err(|err| {
73            ic_msg!(invoke_context, "proof verification failed: {:?}", err);
74            InstructionError::InvalidInstructionData
75        })?;
76
77        *proof_data.context_data()
78    } else {
79        let proof_data =
80            ProofInstruction::proof_data::<T, U>(instruction_data).ok_or_else(|| {
81                ic_msg!(invoke_context, "invalid proof data");
82                InstructionError::InvalidInstructionData
83            })?;
84        proof_data.verify_proof().map_err(|err| {
85            ic_msg!(invoke_context, "proof_verification failed: {:?}", err);
86            InstructionError::InvalidInstructionData
87        })?;
88
89        *proof_data.context_data()
90    };
91
92    // create context state if additional accounts are provided with the instruction
93    if instruction_context.get_number_of_instruction_accounts() > accessed_accounts {
94        let context_state_authority = *instruction_context
95            .try_borrow_instruction_account(
96                transaction_context,
97                accessed_accounts.checked_add(1).unwrap(),
98            )?
99            .get_key();
100
101        let mut proof_context_account = instruction_context
102            .try_borrow_instruction_account(transaction_context, accessed_accounts)?;
103
104        if *proof_context_account.get_owner() != id() {
105            return Err(InstructionError::InvalidAccountOwner);
106        }
107
108        let proof_context_state_meta =
109            ProofContextStateMeta::try_from_bytes(proof_context_account.get_data())?;
110
111        if proof_context_state_meta.proof_type != ProofType::Uninitialized.into() {
112            return Err(InstructionError::AccountAlreadyInitialized);
113        }
114
115        let context_state_data =
116            ProofContextState::encode(&context_state_authority, T::PROOF_TYPE, &context_data);
117
118        if proof_context_account.get_data().len() != context_state_data.len() {
119            return Err(InstructionError::InvalidAccountData);
120        }
121
122        proof_context_account.set_data_from_slice(&context_state_data)?;
123    }
124
125    Ok(())
126}
127
128fn process_close_proof_context(invoke_context: &mut InvokeContext) -> Result<(), InstructionError> {
129    let transaction_context = &invoke_context.transaction_context;
130    let instruction_context = transaction_context.get_current_instruction_context()?;
131
132    let owner_pubkey = {
133        let owner_account =
134            instruction_context.try_borrow_instruction_account(transaction_context, 2)?;
135
136        if !owner_account.is_signer() {
137            return Err(InstructionError::MissingRequiredSignature);
138        }
139        *owner_account.get_key()
140    }; // done with `owner_account`, so drop it to prevent a potential double borrow
141
142    let proof_context_account_pubkey = *instruction_context
143        .try_borrow_instruction_account(transaction_context, 0)?
144        .get_key();
145    let destination_account_pubkey = *instruction_context
146        .try_borrow_instruction_account(transaction_context, 1)?
147        .get_key();
148    if proof_context_account_pubkey == destination_account_pubkey {
149        return Err(InstructionError::InvalidInstructionData);
150    }
151
152    let mut proof_context_account =
153        instruction_context.try_borrow_instruction_account(transaction_context, 0)?;
154    let proof_context_state_meta =
155        ProofContextStateMeta::try_from_bytes(proof_context_account.get_data())?;
156    let expected_owner_pubkey = proof_context_state_meta.context_state_authority;
157
158    if owner_pubkey != expected_owner_pubkey {
159        return Err(InstructionError::InvalidAccountOwner);
160    }
161
162    let mut destination_account =
163        instruction_context.try_borrow_instruction_account(transaction_context, 1)?;
164    destination_account.checked_add_lamports(proof_context_account.get_lamports())?;
165    proof_context_account.set_lamports(0)?;
166    proof_context_account.set_data_length(0)?;
167    proof_context_account.set_owner(system_program::id().as_ref())?;
168
169    Ok(())
170}
171
172declare_process_instruction!(Entrypoint, 0, |invoke_context| {
173    let transaction_context = &invoke_context.transaction_context;
174    let instruction_context = transaction_context.get_current_instruction_context()?;
175    let instruction_data = instruction_context.get_instruction_data();
176    let instruction = ProofInstruction::instruction_type(instruction_data)
177        .ok_or(InstructionError::InvalidInstructionData)?;
178
179    match instruction {
180        ProofInstruction::CloseContextState => {
181            invoke_context
182                .consume_checked(CLOSE_CONTEXT_STATE_COMPUTE_UNITS)
183                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
184            ic_msg!(invoke_context, "CloseContextState");
185            process_close_proof_context(invoke_context)
186        }
187        ProofInstruction::VerifyZeroCiphertext => {
188            invoke_context
189                .consume_checked(VERIFY_ZERO_CIPHERTEXT_COMPUTE_UNITS)
190                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
191            ic_msg!(invoke_context, "VerifyZeroCiphertext");
192            process_verify_proof::<ZeroCiphertextProofData, ZeroCiphertextProofContext>(
193                invoke_context,
194            )
195        }
196        ProofInstruction::VerifyCiphertextCiphertextEquality => {
197            invoke_context
198                .consume_checked(VERIFY_CIPHERTEXT_CIPHERTEXT_EQUALITY_COMPUTE_UNITS)
199                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
200            ic_msg!(invoke_context, "VerifyCiphertextCiphertextEquality");
201            process_verify_proof::<
202                CiphertextCiphertextEqualityProofData,
203                CiphertextCiphertextEqualityProofContext,
204            >(invoke_context)
205        }
206        ProofInstruction::VerifyCiphertextCommitmentEquality => {
207            invoke_context
208                .consume_checked(VERIFY_CIPHERTEXT_COMMITMENT_EQUALITY_COMPUTE_UNITS)
209                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
210            ic_msg!(invoke_context, "VerifyCiphertextCommitmentEquality");
211            process_verify_proof::<
212                CiphertextCommitmentEqualityProofData,
213                CiphertextCommitmentEqualityProofContext,
214            >(invoke_context)
215        }
216        ProofInstruction::VerifyPubkeyValidity => {
217            invoke_context
218                .consume_checked(VERIFY_PUBKEY_VALIDITY_COMPUTE_UNITS)
219                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
220            ic_msg!(invoke_context, "VerifyPubkeyValidity");
221            process_verify_proof::<PubkeyValidityProofData, PubkeyValidityProofContext>(
222                invoke_context,
223            )
224        }
225        ProofInstruction::VerifyPercentageWithCap => {
226            invoke_context
227                .consume_checked(VERIFY_PERCENTAGE_WITH_CAP_COMPUTE_UNITS)
228                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
229            ic_msg!(invoke_context, "VerifyPercentageWithCap");
230            process_verify_proof::<PercentageWithCapProofData, PercentageWithCapProofContext>(
231                invoke_context,
232            )
233        }
234        ProofInstruction::VerifyBatchedRangeProofU64 => {
235            invoke_context
236                .consume_checked(VERIFY_BATCHED_RANGE_PROOF_U64_COMPUTE_UNITS)
237                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
238            ic_msg!(invoke_context, "VerifyBatchedRangeProofU64");
239            process_verify_proof::<BatchedRangeProofU64Data, BatchedRangeProofContext>(
240                invoke_context,
241            )
242        }
243        ProofInstruction::VerifyBatchedRangeProofU128 => {
244            invoke_context
245                .consume_checked(VERIFY_BATCHED_RANGE_PROOF_U128_COMPUTE_UNITS)
246                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
247            ic_msg!(invoke_context, "VerifyBatchedRangeProofU128");
248            process_verify_proof::<BatchedRangeProofU128Data, BatchedRangeProofContext>(
249                invoke_context,
250            )
251        }
252        ProofInstruction::VerifyBatchedRangeProofU256 => {
253            invoke_context
254                .consume_checked(VERIFY_BATCHED_RANGE_PROOF_U256_COMPUTE_UNITS)
255                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
256            ic_msg!(invoke_context, "VerifyBatchedRangeProofU256");
257            process_verify_proof::<BatchedRangeProofU256Data, BatchedRangeProofContext>(
258                invoke_context,
259            )
260        }
261        ProofInstruction::VerifyGroupedCiphertext2HandlesValidity => {
262            invoke_context
263                .consume_checked(VERIFY_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_COMPUTE_UNITS)
264                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
265            ic_msg!(invoke_context, "VerifyGroupedCiphertext2HandlesValidity");
266            process_verify_proof::<
267                GroupedCiphertext2HandlesValidityProofData,
268                GroupedCiphertext2HandlesValidityProofContext,
269            >(invoke_context)
270        }
271        ProofInstruction::VerifyBatchedGroupedCiphertext2HandlesValidity => {
272            invoke_context
273                .consume_checked(VERIFY_BATCHED_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_COMPUTE_UNITS)
274                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
275            ic_msg!(
276                invoke_context,
277                "VerifyBatchedGroupedCiphertext2HandlesValidity"
278            );
279            process_verify_proof::<
280                BatchedGroupedCiphertext2HandlesValidityProofData,
281                BatchedGroupedCiphertext2HandlesValidityProofContext,
282            >(invoke_context)
283        }
284        ProofInstruction::VerifyGroupedCiphertext3HandlesValidity => {
285            invoke_context
286                .consume_checked(VERIFY_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_COMPUTE_UNITS)
287                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
288            ic_msg!(invoke_context, "VerifyGroupedCiphertext3HandlesValidity");
289            process_verify_proof::<
290                GroupedCiphertext3HandlesValidityProofData,
291                GroupedCiphertext3HandlesValidityProofContext,
292            >(invoke_context)
293        }
294        ProofInstruction::VerifyBatchedGroupedCiphertext3HandlesValidity => {
295            invoke_context
296                .consume_checked(VERIFY_BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_COMPUTE_UNITS)
297                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
298            ic_msg!(
299                invoke_context,
300                "VerifyBatchedGroupedCiphertext3HandlesValidity"
301            );
302            process_verify_proof::<
303                BatchedGroupedCiphertext3HandlesValidityProofData,
304                BatchedGroupedCiphertext3HandlesValidityProofContext,
305            >(invoke_context)
306        }
307    }
308});
solana_zk_elgamal_proof_program/lib.rs

solana_zk_elgamal_proof_program/
lib.rs
