light_system_program/invoke_cpi/
process_cpi_context.rs

1use super::{account::CpiContextAccount, InstructionDataInvokeCpi};
2use crate::errors::SystemProgramError;
3use anchor_lang::prelude::*;
4
5/// Cpi context enables the use of input compressed accounts owned by different
6/// programs.
7///
8/// Example:
9/// - a transaction calling a pda program needs to transfer tokens and modify a
10///   compressed pda
11/// - the pda is owned by pda program while the tokens are owned by the compressed
12///   token program
13///
14/// without cpi context:
15/// - naively invoking each compressed token via cpi and modifying the pda
16///   requires two proofs 128 bytes and ~100,000 CU each
17///
18/// with cpi context:
19/// - only one proof is required -> less instruction data and CU cost
20/// 1. first invocation (token program) performs signer checks of the compressed
21///    token accounts, caches these in the cpi context and returns. The state
22///    transition is not executed yet.
23/// 2. second invocation (pda program) performs signer checks of the pda
24///    compressed account, reads cpi context and combines the instruction inputs
25///    with verified inputs from the cpi context. The proof is verified and
26///    other state transition is executed with the combined inputs.
27pub fn process_cpi_context<'info>(
28    mut inputs: InstructionDataInvokeCpi,
29    cpi_context_account: &mut Option<Account<'info, CpiContextAccount>>,
30    fee_payer: Pubkey,
31    remaining_accounts: &[AccountInfo<'info>],
32) -> Result<Option<InstructionDataInvokeCpi>> {
33    let cpi_context = &inputs.cpi_context;
34    if cpi_context_account.is_some() && cpi_context.is_none() {
35        msg!("cpi context account is some but cpi context is none");
36        return err!(SystemProgramError::CpiContextMissing);
37    }
38
39    if let Some(cpi_context) = cpi_context {
40        let cpi_context_account = match cpi_context_account {
41            Some(cpi_context_account) => cpi_context_account,
42            None => return err!(SystemProgramError::CpiContextAccountUndefined),
43        };
44        let index = if !inputs
45            .input_compressed_accounts_with_merkle_context
46            .is_empty()
47        {
48            inputs.input_compressed_accounts_with_merkle_context[0]
49                .merkle_context
50                .merkle_tree_pubkey_index
51        } else if !inputs.output_compressed_accounts.is_empty() {
52            inputs.output_compressed_accounts[0].merkle_tree_index
53        } else {
54            return err!(SystemProgramError::NoInputs);
55        };
56        let first_merkle_tree_pubkey = remaining_accounts[index as usize].key();
57        if first_merkle_tree_pubkey != cpi_context_account.associated_merkle_tree {
58            msg!(
59                "first_merkle_tree_pubkey {:?} != associated_merkle_tree {:?}",
60                first_merkle_tree_pubkey,
61                cpi_context_account.associated_merkle_tree
62            );
63            return err!(SystemProgramError::CpiContextAssociatedMerkleTreeMismatch);
64        }
65        if cpi_context.set_context {
66            set_cpi_context(fee_payer, cpi_context_account, inputs)?;
67            return Ok(None);
68        } else {
69            if cpi_context_account.context.is_empty() {
70                msg!("cpi context account : {:?}", cpi_context_account);
71                msg!("fee payer : {:?}", fee_payer);
72                msg!("cpi context  : {:?}", cpi_context);
73                return err!(SystemProgramError::CpiContextEmpty);
74            } else if cpi_context_account.fee_payer != fee_payer || cpi_context.first_set_context {
75                msg!("cpi context account : {:?}", cpi_context_account);
76                msg!("fee payer : {:?}", fee_payer);
77                msg!("cpi context  : {:?}", cpi_context);
78                return err!(SystemProgramError::CpiContextFeePayerMismatch);
79            }
80            inputs.combine(&cpi_context_account.context);
81            // Reset cpi context account
82            cpi_context_account.context = Vec::new();
83            cpi_context_account.fee_payer = Pubkey::default();
84        }
85    }
86    Ok(Some(inputs))
87}
88
89pub fn set_cpi_context(
90    fee_payer: Pubkey,
91    cpi_context_account: &mut CpiContextAccount,
92    mut inputs: InstructionDataInvokeCpi,
93) -> Result<()> {
94    // SAFETY Assumptions:
95    // -  previous data in cpi_context_account
96    //   -> we require the account to be wiped in the beginning of a
97    //   transaction
98    // - leaf over data: There cannot be any leftover data in the
99    //   account since if the transaction fails the account doesn't change.
100
101    // Expected usage:
102    // 1. The first invocation is marked with
103    // No need to store the proof (except in first invokation),
104    // cpi context, compress_or_decompress_lamports,
105    // relay_fee
106    // 2. Subsequent invocations check the proof and fee payer
107    if inputs.cpi_context.unwrap().first_set_context {
108        clean_input_data(&mut inputs);
109        cpi_context_account.context = vec![inputs];
110        cpi_context_account.fee_payer = fee_payer;
111    } else if fee_payer == cpi_context_account.fee_payer && !cpi_context_account.context.is_empty()
112    {
113        clean_input_data(&mut inputs);
114        cpi_context_account.context.push(inputs);
115    } else {
116        msg!(" {} != {}", fee_payer, cpi_context_account.fee_payer);
117        return err!(SystemProgramError::CpiContextFeePayerMismatch);
118    }
119    Ok(())
120}
121
122fn clean_input_data(inputs: &mut InstructionDataInvokeCpi) {
123    inputs.cpi_context = None;
124    inputs.compress_or_decompress_lamports = None;
125    inputs.relay_fee = None;
126    inputs.proof = None;
127}
128
129/// Set cpi context tests:
130/// 1. Functional: Set cpi context first invocation
131/// 2. Functional: Set cpi context subsequent invocation
132/// 3. Failing: Set cpi context fee payer mismatch
133/// 4. Failing: Set cpi context without first context
134///
135/// process cpi context:
136/// 1. CpiContextMissing
137/// 2. CpiContextAccountUndefined
138/// 3. NoInputs
139/// 4. CpiContextAssociatedMerkleTreeMismatch
140/// 5. CpiContextEmpty
141/// 6. CpiContextFeePayerMismatch
142///
143/// Functional process cpi context:
144/// 1. Set context
145/// 2. Combine (with malicious input in cpi context account)
146#[cfg(test)]
147mod tests {
148    use std::cell::RefCell;
149
150    use crate::{
151        sdk::{
152            compressed_account::{
153                CompressedAccount, PackedCompressedAccountWithMerkleContext, PackedMerkleContext,
154            },
155            CompressedCpiContext,
156        },
157        NewAddressParamsPacked, OutputCompressedAccountWithPackedContext,
158    };
159
160    use super::*;
161    use anchor_lang::solana_program::pubkey::Pubkey;
162
163    fn create_test_cpi_context_account() -> CpiContextAccount {
164        CpiContextAccount {
165            fee_payer: Pubkey::new_unique(),
166            associated_merkle_tree: Pubkey::new_unique(),
167            context: vec![],
168        }
169    }
170
171    fn create_test_instruction_data(
172        first_set_context: bool,
173        set_context: bool,
174        iter: u8,
175    ) -> InstructionDataInvokeCpi {
176        InstructionDataInvokeCpi {
177            proof: None,
178            new_address_params: vec![NewAddressParamsPacked {
179                seed: vec![iter; 32].try_into().unwrap(),
180                address_merkle_tree_account_index: iter,
181                address_merkle_tree_root_index: iter.into(),
182                address_queue_account_index: iter,
183            }],
184            input_compressed_accounts_with_merkle_context: vec![
185                PackedCompressedAccountWithMerkleContext {
186                    compressed_account: CompressedAccount {
187                        owner: Pubkey::new_unique(),
188                        lamports: iter.into(),
189                        address: None,
190                        data: None,
191                    },
192                    merkle_context: PackedMerkleContext {
193                        merkle_tree_pubkey_index: 0,
194                        nullifier_queue_pubkey_index: iter,
195                        leaf_index: 0,
196                        queue_index: None,
197                    },
198                    root_index: iter.into(),
199                    read_only: false,
200                },
201            ],
202            output_compressed_accounts: vec![OutputCompressedAccountWithPackedContext {
203                compressed_account: CompressedAccount {
204                    owner: Pubkey::new_unique(),
205                    lamports: iter.into(),
206                    address: None,
207                    data: None,
208                },
209                merkle_tree_index: iter,
210            }],
211            relay_fee: None,
212            compress_or_decompress_lamports: None,
213            is_compress: false,
214            cpi_context: Some(CompressedCpiContext {
215                first_set_context,
216                set_context,
217                cpi_context_account_index: 0,
218            }),
219        }
220    }
221
222    #[test]
223    fn test_set_cpi_context_first_invocation() {
224        let fee_payer = Pubkey::new_unique();
225        let mut cpi_context_account = create_test_cpi_context_account();
226        let mut inputs = create_test_instruction_data(true, true, 1);
227
228        let result = set_cpi_context(fee_payer, &mut cpi_context_account, inputs.clone());
229        assert!(result.is_ok());
230        assert_eq!(cpi_context_account.fee_payer, fee_payer);
231        assert_eq!(cpi_context_account.context.len(), 1);
232        assert_ne!(cpi_context_account.context[0], inputs);
233        clean_input_data(&mut inputs);
234        assert_eq!(cpi_context_account.context[0], inputs);
235    }
236
237    #[test]
238    fn test_set_cpi_context_subsequent_invocation() {
239        let fee_payer = Pubkey::new_unique();
240        let mut cpi_context_account = create_test_cpi_context_account();
241        let inputs_first = create_test_instruction_data(true, true, 1);
242        set_cpi_context(fee_payer, &mut cpi_context_account, inputs_first.clone()).unwrap();
243
244        let mut inputs_subsequent = create_test_instruction_data(false, true, 2);
245        let result = set_cpi_context(
246            fee_payer,
247            &mut cpi_context_account,
248            inputs_subsequent.clone(),
249        );
250        assert!(result.is_ok());
251        assert_eq!(cpi_context_account.context.len(), 2);
252        clean_input_data(&mut inputs_subsequent);
253        assert_eq!(cpi_context_account.context[1], inputs_subsequent);
254    }
255
256    #[test]
257    fn test_set_cpi_context_fee_payer_mismatch() {
258        let fee_payer = Pubkey::new_unique();
259        let mut cpi_context_account = create_test_cpi_context_account();
260        let inputs_first = create_test_instruction_data(true, true, 1);
261        set_cpi_context(fee_payer, &mut cpi_context_account, inputs_first.clone()).unwrap();
262
263        let different_fee_payer = Pubkey::new_unique();
264        let inputs_subsequent = create_test_instruction_data(false, true, 2);
265        let result = set_cpi_context(
266            different_fee_payer,
267            &mut cpi_context_account,
268            inputs_subsequent,
269        );
270        assert!(result.is_err());
271    }
272
273    #[test]
274    fn test_set_cpi_context_without_first_context() {
275        let fee_payer = Pubkey::new_unique();
276        let mut cpi_context_account = create_test_cpi_context_account();
277        let inputs_first = create_test_instruction_data(false, true, 1);
278        let result = set_cpi_context(fee_payer, &mut cpi_context_account, inputs_first.clone());
279        assert_eq!(
280            result,
281            Err(SystemProgramError::CpiContextFeePayerMismatch.into())
282        );
283    }
284
285    /// Check: process cpi 1
286    #[test]
287    fn test_process_cpi_context_both_none() {
288        let fee_payer = Pubkey::new_unique();
289        let inputs = create_test_instruction_data(false, true, 1);
290        let mut cpi_context_account: Option<Account<CpiContextAccount>> = None;
291
292        let result = process_cpi_context(inputs.clone(), &mut cpi_context_account, fee_payer, &[]);
293        assert_eq!(
294            result,
295            Err(SystemProgramError::CpiContextAccountUndefined.into())
296        );
297    }
298
299    /// Check: process cpi 1
300    #[test]
301    fn test_process_cpi_context_account_none_context_some() {
302        let fee_payer = Pubkey::new_unique();
303        let inputs = create_test_instruction_data(false, true, 1);
304        let mut cpi_context_account: Option<Account<CpiContextAccount>> = None;
305
306        let result = process_cpi_context(inputs, &mut cpi_context_account, fee_payer, &[]);
307        assert_eq!(
308            result,
309            Err(SystemProgramError::CpiContextAccountUndefined.into())
310        );
311    }
312
313    /// Check: process cpi 2
314    #[test]
315    fn test_process_cpi_context_account_some_context_none() {
316        let fee_payer = Pubkey::new_unique();
317        let inputs = InstructionDataInvokeCpi {
318            cpi_context: None,
319            ..create_test_instruction_data(false, true, 1)
320        };
321        let mut lamports = 0;
322        let cpi_context_content = CpiContextAccount {
323            fee_payer: Pubkey::default(),
324            associated_merkle_tree: Pubkey::new_unique(),
325            context: vec![],
326        };
327        let mut data = vec![22, 20, 149, 218, 74, 204, 128, 166];
328        data.extend_from_slice(&cpi_context_content.try_to_vec().unwrap());
329        let account_info = AccountInfo {
330            key: &Pubkey::new_unique(),
331            is_signer: false,
332            is_writable: false,
333            lamports: RefCell::new(&mut lamports).into(),
334            data: RefCell::new(data.as_mut_slice()).into(),
335            owner: &crate::ID,
336            rent_epoch: 0,
337            executable: false,
338        };
339        let mut cpi_context_account = Some(Account::try_from(account_info.as_ref()).unwrap());
340        let result = process_cpi_context(inputs, &mut cpi_context_account, fee_payer, &[]);
341        assert_eq!(result, Err(SystemProgramError::CpiContextMissing.into()));
342    }
343
344    /// Check: process cpi 3
345    #[test]
346    fn test_process_cpi_no_inputs() {
347        let fee_payer = Pubkey::new_unique();
348        let mut inputs = create_test_instruction_data(false, true, 1);
349        inputs.input_compressed_accounts_with_merkle_context = vec![];
350        inputs.output_compressed_accounts = vec![];
351
352        let mut lamports = 0;
353        let cpi_context_content = CpiContextAccount {
354            fee_payer: Pubkey::default(),
355            associated_merkle_tree: Pubkey::new_unique(),
356            context: vec![],
357        };
358        let mut data = vec![22, 20, 149, 218, 74, 204, 128, 166];
359        data.extend_from_slice(&cpi_context_content.try_to_vec().unwrap());
360        let account_info = AccountInfo {
361            key: &Pubkey::new_unique(),
362            is_signer: false,
363            is_writable: false,
364            lamports: RefCell::new(&mut lamports).into(),
365            data: RefCell::new(data.as_mut_slice()).into(),
366            owner: &crate::ID,
367            rent_epoch: 0,
368            executable: false,
369        };
370        let mut cpi_context_account = Some(Account::try_from(account_info.as_ref()).unwrap());
371        let result = process_cpi_context(inputs, &mut cpi_context_account, fee_payer, &[]);
372        assert_eq!(result, Err(SystemProgramError::NoInputs.into()));
373    }
374
375    /// Check: process cpi 4
376    #[test]
377    fn test_process_cpi_context_associated_tree_mismatch() {
378        let fee_payer = Pubkey::new_unique();
379        let inputs = create_test_instruction_data(true, true, 1);
380        let mut lamports = 0;
381        let merkle_tree_pubkey = Pubkey::new_unique();
382        let cpi_context_content = CpiContextAccount {
383            fee_payer: Pubkey::default(),
384            associated_merkle_tree: merkle_tree_pubkey,
385            context: vec![],
386        };
387        let mut data = vec![22, 20, 149, 218, 74, 204, 128, 166];
388        data.extend_from_slice(&cpi_context_content.try_to_vec().unwrap());
389        let account_info = AccountInfo {
390            key: &Pubkey::new_unique(),
391            is_signer: false,
392            is_writable: false,
393            lamports: RefCell::new(&mut lamports).into(),
394            data: RefCell::new(data.as_mut_slice()).into(),
395            owner: &crate::ID,
396            rent_epoch: 0,
397            executable: false,
398        };
399        let mut cpi_context_account = Some(Account::try_from(account_info.as_ref()).unwrap());
400        let mut mt_lamports = 0;
401        let mut data = vec![172, 43, 172, 186, 29, 73, 219, 84];
402        let invalid_merkle_tree_pubkey = Pubkey::new_unique();
403        let merkle_tree_account_info = AccountInfo {
404            key: &invalid_merkle_tree_pubkey,
405            is_signer: false,
406            is_writable: false,
407            lamports: RefCell::new(&mut mt_lamports).into(),
408            data: RefCell::new(data.as_mut_slice()).into(),
409            owner: &crate::ID,
410            rent_epoch: 0,
411            executable: false,
412        };
413        let remaining_accounts = &[merkle_tree_account_info];
414        let result = process_cpi_context(
415            inputs,
416            &mut cpi_context_account,
417            fee_payer,
418            remaining_accounts,
419        );
420        assert_eq!(
421            result,
422            Err(SystemProgramError::CpiContextAssociatedMerkleTreeMismatch.into())
423        );
424    }
425
426    /// Check: process cpi 5
427    #[test]
428    fn test_process_cpi_context_no_set_context() {
429        let fee_payer = Pubkey::new_unique();
430        let inputs = create_test_instruction_data(false, false, 1);
431        let mut lamports = 0;
432        let merkle_tree_pubkey = Pubkey::new_unique();
433        let cpi_context_content = CpiContextAccount {
434            fee_payer: Pubkey::default(),
435            associated_merkle_tree: merkle_tree_pubkey,
436            context: vec![],
437        };
438        let mut data = vec![22, 20, 149, 218, 74, 204, 128, 166];
439        data.extend_from_slice(&cpi_context_content.try_to_vec().unwrap());
440        let account_info = AccountInfo {
441            key: &Pubkey::new_unique(),
442            is_signer: false,
443            is_writable: false,
444            lamports: RefCell::new(&mut lamports).into(),
445            data: RefCell::new(data.as_mut_slice()).into(),
446            owner: &crate::ID,
447            rent_epoch: 0,
448            executable: false,
449        };
450        let mut cpi_context_account = Some(Account::try_from(account_info.as_ref()).unwrap());
451        let mut mt_lamports = 0;
452        let mut data = vec![172, 43, 172, 186, 29, 73, 219, 84];
453        let merkle_tree_account_info = AccountInfo {
454            key: &merkle_tree_pubkey,
455            is_signer: false,
456            is_writable: false,
457            lamports: RefCell::new(&mut mt_lamports).into(),
458            data: RefCell::new(data.as_mut_slice()).into(),
459            owner: &crate::ID,
460            rent_epoch: 0,
461            executable: false,
462        };
463        let remaining_accounts = &[merkle_tree_account_info];
464        let result = process_cpi_context(
465            inputs.clone(),
466            &mut cpi_context_account,
467            fee_payer,
468            remaining_accounts,
469        );
470        assert_eq!(result, Err(SystemProgramError::CpiContextEmpty.into()));
471    }
472
473    /// Check: process cpi 6
474    #[test]
475    fn test_process_cpi_context_empty_context_error() {
476        let fee_payer = Pubkey::default();
477        let inputs = create_test_instruction_data(false, true, 1);
478        let mut lamports = 0;
479        let merkle_tree_pubkey = Pubkey::new_unique();
480        let cpi_context_content = CpiContextAccount {
481            fee_payer: Pubkey::default(),
482            associated_merkle_tree: merkle_tree_pubkey,
483            context: vec![],
484        };
485        let mut data = vec![22, 20, 149, 218, 74, 204, 128, 166];
486        data.extend_from_slice(&cpi_context_content.try_to_vec().unwrap());
487        let account_info = AccountInfo {
488            key: &Pubkey::new_unique(),
489            is_signer: false,
490            is_writable: false,
491            lamports: RefCell::new(&mut lamports).into(),
492            data: RefCell::new(data.as_mut_slice()).into(),
493            owner: &crate::ID,
494            rent_epoch: 0,
495            executable: false,
496        };
497        let mut cpi_context_account = Some(Account::try_from(account_info.as_ref()).unwrap());
498        let mut mt_lamports = 0;
499        let mut data = vec![172, 43, 172, 186, 29, 73, 219, 84];
500        let merkle_tree_account_info = AccountInfo {
501            key: &merkle_tree_pubkey,
502            is_signer: false,
503            is_writable: false,
504            lamports: RefCell::new(&mut mt_lamports).into(),
505            data: RefCell::new(data.as_mut_slice()).into(),
506            owner: &crate::ID,
507            rent_epoch: 0,
508            executable: false,
509        };
510        let remaining_accounts = &[merkle_tree_account_info];
511        let result = process_cpi_context(
512            inputs,
513            &mut cpi_context_account,
514            fee_payer,
515            remaining_accounts,
516        );
517        assert_eq!(
518            result,
519            Err(SystemProgramError::CpiContextFeePayerMismatch.into())
520        );
521    }
522
523    /// Check: process cpi 6
524    #[test]
525    fn test_process_cpi_context_fee_payer_mismatch_error() {
526        let fee_payer = Pubkey::new_unique();
527        let inputs = create_test_instruction_data(true, true, 1);
528        let mut lamports = 0;
529        let merkle_tree_pubkey = Pubkey::new_unique();
530        let cpi_context_content = CpiContextAccount {
531            fee_payer: Pubkey::default(),
532            associated_merkle_tree: merkle_tree_pubkey,
533            context: vec![],
534        };
535        let mut data = vec![22, 20, 149, 218, 74, 204, 128, 166];
536        data.extend_from_slice(&cpi_context_content.try_to_vec().unwrap());
537        let account_info = AccountInfo {
538            key: &Pubkey::new_unique(),
539            is_signer: false,
540            is_writable: false,
541            lamports: RefCell::new(&mut lamports).into(),
542            data: RefCell::new(data.as_mut_slice()).into(),
543            owner: &crate::ID,
544            rent_epoch: 0,
545            executable: false,
546        };
547        let mut cpi_context_account = Some(Account::try_from(account_info.as_ref()).unwrap());
548        let mut mt_lamports = 0;
549        let mut data = vec![172, 43, 172, 186, 29, 73, 219, 84];
550        let merkle_tree_account_info = AccountInfo {
551            key: &merkle_tree_pubkey,
552            is_signer: false,
553            is_writable: false,
554            lamports: RefCell::new(&mut mt_lamports).into(),
555            data: RefCell::new(data.as_mut_slice()).into(),
556            owner: &crate::ID,
557            rent_epoch: 0,
558            executable: false,
559        };
560        let remaining_accounts = &[merkle_tree_account_info];
561        let result = process_cpi_context(
562            inputs.clone(),
563            &mut cpi_context_account,
564            fee_payer,
565            remaining_accounts,
566        );
567        assert!(result.is_ok());
568        let invalid_fee_payer = Pubkey::new_unique();
569        let inputs = create_test_instruction_data(false, true, 1);
570        let result = process_cpi_context(
571            inputs,
572            &mut cpi_context_account,
573            invalid_fee_payer,
574            remaining_accounts,
575        );
576        assert_eq!(
577            result,
578            Err(SystemProgramError::CpiContextFeePayerMismatch.into())
579        );
580    }
581
582    #[test]
583    fn test_process_cpi_context_set_context() {
584        let fee_payer = Pubkey::new_unique();
585        let mut inputs = create_test_instruction_data(true, true, 1);
586        let mut lamports = 0;
587        let merkle_tree_pubkey = Pubkey::new_unique();
588        let cpi_context_content = CpiContextAccount {
589            fee_payer: Pubkey::default(),
590            associated_merkle_tree: merkle_tree_pubkey,
591            context: vec![],
592        };
593        let mut data = vec![22, 20, 149, 218, 74, 204, 128, 166];
594        data.extend_from_slice(&cpi_context_content.try_to_vec().unwrap());
595        let account_info = AccountInfo {
596            key: &Pubkey::new_unique(),
597            is_signer: false,
598            is_writable: false,
599            lamports: RefCell::new(&mut lamports).into(),
600            data: RefCell::new(data.as_mut_slice()).into(),
601            owner: &crate::ID,
602            rent_epoch: 0,
603            executable: false,
604        };
605        let mut cpi_context_account = Some(Account::try_from(account_info.as_ref()).unwrap());
606        let mut mt_lamports = 0;
607        let mut data = vec![172, 43, 172, 186, 29, 73, 219, 84];
608        let merkle_tree_account_info = AccountInfo {
609            key: &merkle_tree_pubkey,
610            is_signer: false,
611            is_writable: false,
612            lamports: RefCell::new(&mut mt_lamports).into(),
613            data: RefCell::new(data.as_mut_slice()).into(),
614            owner: &crate::ID,
615            rent_epoch: 0,
616            executable: false,
617        };
618        let remaining_accounts = &[merkle_tree_account_info];
619        let result = process_cpi_context(
620            inputs.clone(),
621            &mut cpi_context_account,
622            fee_payer,
623            remaining_accounts,
624        );
625        assert!(result.is_ok());
626        assert_eq!(cpi_context_account.as_ref().unwrap().context.len(), 1);
627        assert_eq!(cpi_context_account.as_ref().unwrap().fee_payer, fee_payer);
628        clean_input_data(&mut inputs);
629        assert_eq!(cpi_context_account.as_ref().unwrap().context[0], inputs);
630        assert_eq!(result.unwrap(), None);
631    }
632
633    #[test]
634    fn test_process_cpi_context_combine() {
635        let fee_payer = Pubkey::new_unique();
636        let mut inputs = create_test_instruction_data(true, true, 1);
637        let malicious_inputs = create_test_instruction_data(true, true, 100);
638
639        let mut lamports = 0;
640        let merkle_tree_pubkey = Pubkey::new_unique();
641        let cpi_context_content = CpiContextAccount {
642            fee_payer: Pubkey::default(),
643            associated_merkle_tree: merkle_tree_pubkey,
644            context: vec![malicious_inputs],
645        };
646        let mut data = vec![22, 20, 149, 218, 74, 204, 128, 166];
647        data.extend_from_slice(&cpi_context_content.try_to_vec().unwrap());
648        let account_info = AccountInfo {
649            key: &Pubkey::new_unique(),
650            is_signer: false,
651            is_writable: false,
652            lamports: RefCell::new(&mut lamports).into(),
653            data: RefCell::new(data.as_mut_slice()).into(),
654            owner: &crate::ID,
655            rent_epoch: 0,
656            executable: false,
657        };
658        let mut cpi_context_account = Some(Account::try_from(account_info.as_ref()).unwrap());
659        let mut mt_lamports = 0;
660        let mut data = vec![172, 43, 172, 186, 29, 73, 219, 84];
661        let merkle_tree_account_info = AccountInfo {
662            key: &merkle_tree_pubkey,
663            is_signer: false,
664            is_writable: false,
665            lamports: RefCell::new(&mut mt_lamports).into(),
666            data: RefCell::new(data.as_mut_slice()).into(),
667            owner: &crate::ID,
668            rent_epoch: 0,
669            executable: false,
670        };
671        let remaining_accounts = &[merkle_tree_account_info];
672        let result = process_cpi_context(
673            inputs.clone(),
674            &mut cpi_context_account,
675            fee_payer,
676            remaining_accounts,
677        );
678        assert!(result.is_ok());
679        assert_eq!(cpi_context_account.as_ref().unwrap().context.len(), 1);
680        assert_eq!(cpi_context_account.as_ref().unwrap().fee_payer, fee_payer);
681        clean_input_data(&mut inputs);
682
683        assert_eq!(cpi_context_account.as_ref().unwrap().context[0], inputs);
684        assert_eq!(result.unwrap(), None);
685        for i in 2..10 {
686            let mut inputs = create_test_instruction_data(false, true, i);
687            let result = process_cpi_context(
688                inputs.clone(),
689                &mut cpi_context_account,
690                fee_payer,
691                remaining_accounts,
692            );
693            assert!(result.is_ok());
694            assert_eq!(
695                cpi_context_account.as_ref().unwrap().context.len(),
696                i as usize
697            );
698            assert_eq!(cpi_context_account.as_ref().unwrap().fee_payer, fee_payer);
699            clean_input_data(&mut inputs);
700            assert_eq!(
701                cpi_context_account.as_ref().unwrap().context[(i - 1) as usize],
702                inputs
703            );
704            assert_eq!(result.unwrap(), None);
705        }
706
707        let inputs = create_test_instruction_data(false, false, 10);
708        let result = process_cpi_context(
709            inputs.clone(),
710            &mut cpi_context_account,
711            fee_payer,
712            remaining_accounts,
713        );
714        assert!(result.is_ok());
715        let result = result.unwrap().unwrap();
716        for i in 1..10 {
717            assert_eq!(
718                result.output_compressed_accounts[i]
719                    .compressed_account
720                    .lamports,
721                i as u64
722            );
723            assert_eq!(
724                result.input_compressed_accounts_with_merkle_context[i]
725                    .compressed_account
726                    .lamports,
727                i as u64
728            );
729            assert_eq!(
730                result.new_address_params[i].seed,
731                <[u8; 32]>::try_from(vec![i as u8; 32]).unwrap()
732            );
733        }
734        assert_eq!(
735            cpi_context_account.as_ref().unwrap().associated_merkle_tree,
736            merkle_tree_pubkey
737        );
738        assert_eq!(
739            cpi_context_account.as_ref().unwrap().fee_payer,
740            Pubkey::default()
741        );
742        assert_eq!(cpi_context_account.as_ref().unwrap().context.len(), 0);
743    }
744}
light_system_program/invoke_cpi/process_cpi_context.rs

light_system_program/invoke_cpi/
process_cpi_context.rs
