solana_zk_elgamal_proof_program/
lib.rs

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

solana_zk_elgamal_proof_program/
lib.rs
