solana_zk_elgamal_proof_program/
lib.rs

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

solana_zk_elgamal_proof_program/
lib.rs
