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

solana_zk_elgamal_proof_program/
lib.rs
