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                .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                .consume_checked(VERIFY_ZERO_CIPHERTEXT_COMPUTE_UNITS)
207                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
208            ic_msg!(invoke_context, "VerifyZeroCiphertext");
209            process_verify_proof::<ZeroCiphertextProofData, ZeroCiphertextProofContext>(
210                invoke_context,
211            )
212        }
213        ProofInstruction::VerifyCiphertextCiphertextEquality => {
214            invoke_context
215                .consume_checked(VERIFY_CIPHERTEXT_CIPHERTEXT_EQUALITY_COMPUTE_UNITS)
216                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
217            ic_msg!(invoke_context, "VerifyCiphertextCiphertextEquality");
218            process_verify_proof::<
219                CiphertextCiphertextEqualityProofData,
220                CiphertextCiphertextEqualityProofContext,
221            >(invoke_context)
222        }
223        ProofInstruction::VerifyCiphertextCommitmentEquality => {
224            invoke_context
225                .consume_checked(VERIFY_CIPHERTEXT_COMMITMENT_EQUALITY_COMPUTE_UNITS)
226                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
227            ic_msg!(invoke_context, "VerifyCiphertextCommitmentEquality");
228            process_verify_proof::<
229                CiphertextCommitmentEqualityProofData,
230                CiphertextCommitmentEqualityProofContext,
231            >(invoke_context)
232        }
233        ProofInstruction::VerifyPubkeyValidity => {
234            invoke_context
235                .consume_checked(VERIFY_PUBKEY_VALIDITY_COMPUTE_UNITS)
236                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
237            ic_msg!(invoke_context, "VerifyPubkeyValidity");
238            process_verify_proof::<PubkeyValidityProofData, PubkeyValidityProofContext>(
239                invoke_context,
240            )
241        }
242        ProofInstruction::VerifyPercentageWithCap => {
243            invoke_context
244                .consume_checked(VERIFY_PERCENTAGE_WITH_CAP_COMPUTE_UNITS)
245                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
246            ic_msg!(invoke_context, "VerifyPercentageWithCap");
247            process_verify_proof::<PercentageWithCapProofData, PercentageWithCapProofContext>(
248                invoke_context,
249            )
250        }
251        ProofInstruction::VerifyBatchedRangeProofU64 => {
252            invoke_context
253                .consume_checked(VERIFY_BATCHED_RANGE_PROOF_U64_COMPUTE_UNITS)
254                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
255            ic_msg!(invoke_context, "VerifyBatchedRangeProofU64");
256            process_verify_proof::<BatchedRangeProofU64Data, BatchedRangeProofContext>(
257                invoke_context,
258            )
259        }
260        ProofInstruction::VerifyBatchedRangeProofU128 => {
261            invoke_context
262                .consume_checked(VERIFY_BATCHED_RANGE_PROOF_U128_COMPUTE_UNITS)
263                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
264            ic_msg!(invoke_context, "VerifyBatchedRangeProofU128");
265            process_verify_proof::<BatchedRangeProofU128Data, BatchedRangeProofContext>(
266                invoke_context,
267            )
268        }
269        ProofInstruction::VerifyBatchedRangeProofU256 => {
270            invoke_context
271                .consume_checked(VERIFY_BATCHED_RANGE_PROOF_U256_COMPUTE_UNITS)
272                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
273            ic_msg!(invoke_context, "VerifyBatchedRangeProofU256");
274            process_verify_proof::<BatchedRangeProofU256Data, BatchedRangeProofContext>(
275                invoke_context,
276            )
277        }
278        ProofInstruction::VerifyGroupedCiphertext2HandlesValidity => {
279            invoke_context
280                .consume_checked(VERIFY_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_COMPUTE_UNITS)
281                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
282            ic_msg!(invoke_context, "VerifyGroupedCiphertext2HandlesValidity");
283            process_verify_proof::<
284                GroupedCiphertext2HandlesValidityProofData,
285                GroupedCiphertext2HandlesValidityProofContext,
286            >(invoke_context)
287        }
288        ProofInstruction::VerifyBatchedGroupedCiphertext2HandlesValidity => {
289            invoke_context
290                .consume_checked(VERIFY_BATCHED_GROUPED_CIPHERTEXT_2_HANDLES_VALIDITY_COMPUTE_UNITS)
291                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
292            ic_msg!(
293                invoke_context,
294                "VerifyBatchedGroupedCiphertext2HandlesValidity"
295            );
296            process_verify_proof::<
297                BatchedGroupedCiphertext2HandlesValidityProofData,
298                BatchedGroupedCiphertext2HandlesValidityProofContext,
299            >(invoke_context)
300        }
301        ProofInstruction::VerifyGroupedCiphertext3HandlesValidity => {
302            invoke_context
303                .consume_checked(VERIFY_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_COMPUTE_UNITS)
304                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
305            ic_msg!(invoke_context, "VerifyGroupedCiphertext3HandlesValidity");
306            process_verify_proof::<
307                GroupedCiphertext3HandlesValidityProofData,
308                GroupedCiphertext3HandlesValidityProofContext,
309            >(invoke_context)
310        }
311        ProofInstruction::VerifyBatchedGroupedCiphertext3HandlesValidity => {
312            invoke_context
313                .consume_checked(VERIFY_BATCHED_GROUPED_CIPHERTEXT_3_HANDLES_VALIDITY_COMPUTE_UNITS)
314                .map_err(|_| InstructionError::ComputationalBudgetExceeded)?;
315            ic_msg!(
316                invoke_context,
317                "VerifyBatchedGroupedCiphertext3HandlesValidity"
318            );
319            process_verify_proof::<
320                BatchedGroupedCiphertext3HandlesValidityProofData,
321                BatchedGroupedCiphertext3HandlesValidityProofContext,
322            >(invoke_context)
323        }
324    }
325});
solana_zk_elgamal_proof_program/lib.rs

solana_zk_elgamal_proof_program/
lib.rs
