solana_zk_elgamal_proof_program/
lib.rs

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

solana_zk_elgamal_proof_program/
lib.rs
