fix: clear cpi context account

Author: ananas-block

programs/system/src/cpi_context/process_cpi_context.rs

@@ -16,7 +16,10 @@ use pinocchio::{account_info::AccountInfo, program_error::ProgramError, pubkey::
 use solana_msg::msg;
 
 use super::state::{deserialize_cpi_context_account, ZCpiContextAccount};
-use crate::{context::WrappedInstructionData, errors::SystemProgramError, Result};
+use crate::{
+    context::WrappedInstructionData, cpi_context::state::deserialize_cpi_context_account_cleared,
+    errors::SystemProgramError, Result,
+};
 
 /// Diff:
 /// 1. return Cpi context instead of combined data.
@@ -107,18 +110,22 @@ pub fn set_cpi_context<'a, 'info, T: InstructionData<'a>>(
     // relay_fee
     // 2. Subsequent invocations check the proof and fee payer
 
-    let mut cpi_context_account =
-        deserialize_cpi_context_account(cpi_context_account_info).map_err(ProgramError::from)?;
-
     if instruction_data.cpi_context().unwrap().first_set_context {
-        cpi_context_account.clear();
+        let mut cpi_context_account =
+            deserialize_cpi_context_account_cleared(cpi_context_account_info)
+                .map_err(ProgramError::from)?;
         *cpi_context_account.fee_payer = fee_payer.into();
         cpi_context_account.store_data(&instruction_data)?;
-    } else if *cpi_context_account.fee_payer == fee_payer && !cpi_context_account.is_empty() {
-        cpi_context_account.store_data(&instruction_data)?;
     } else {
-        msg!(format!(" {:?} != {:?}", fee_payer, cpi_context_account.fee_payer).as_str());
-        return Err(SystemProgramError::CpiContextFeePayerMismatch.into());
+        let mut cpi_context_account = deserialize_cpi_context_account(cpi_context_account_info)
+            .map_err(ProgramError::from)?;
+
+        if *cpi_context_account.fee_payer == fee_payer && !cpi_context_account.is_empty() {
+            cpi_context_account.store_data(&instruction_data)?;
+        } else {
+            msg!(format!(" {:?} != {:?}", fee_payer, cpi_context_account.fee_payer).as_str());
+            return Err(SystemProgramError::CpiContextFeePayerMismatch.into());
+        }
     }
 
     Ok(())

programs/system/src/cpi_context/state.rs

@@ -56,17 +56,6 @@ impl<'a> ZCpiContextAccount<'a> {
             && self.output_data.is_empty()
     }
 
-    pub fn clear(&mut self) {
-        self.new_addresses.clear();
-        self.readonly_addresses.clear();
-        self.readonly_accounts.clear();
-        self.in_accounts.clear();
-        self.out_accounts.clear();
-        *self.output_data_len = 0.into();
-        self.output_data.clear();
-        self.remaining_data.fill(0);
-    }
-
     pub fn store_data<
         'b,
         T: light_compressed_account::instruction_data::traits::InstructionData<'b>,
@@ -185,6 +174,18 @@ impl Discriminator for ZCpiContextAccount<'_> {
 
 pub fn deserialize_cpi_context_account<'a>(
     account_info: &AccountInfo,
+) -> std::result::Result<ZCpiContextAccount<'a>, ZeroCopyError> {
+    deserialize_cpi_context_account_inner::<false>(account_info)
+}
+
+pub fn deserialize_cpi_context_account_cleared<'a>(
+    account_info: &AccountInfo,
+) -> std::result::Result<ZCpiContextAccount<'a>, ZeroCopyError> {
+    deserialize_cpi_context_account_inner::<true>(account_info)
+}
+
+fn deserialize_cpi_context_account_inner<'a, const CLEARED: bool>(
+    account_info: &AccountInfo,
 ) -> std::result::Result<ZCpiContextAccount<'a>, ZeroCopyError> {
     check_owner(&ID, account_info).map_err(|_| ZeroCopyError::IterFromOutOfBounds)?;
     let mut account_data = account_info
@@ -202,21 +203,34 @@ pub fn deserialize_cpi_context_account<'a>(
     let (associated_merkle_tree, data) =
         Ref::<&'a mut [u8], light_compressed_account::pubkey::Pubkey>::from_prefix(data)?;
 
-    let (new_addresses, data) =
+    let (mut new_addresses, data) =
         ZeroCopyVecU8::<'a, CpiContextNewAddressParamsAssignedPacked>::from_bytes_at(data)?;
-    let (readonly_addresses, data) =
+    let (mut readonly_addresses, data) =
         ZeroCopyVecU8::<'a, ZPackedReadOnlyAddress>::from_bytes_at(data)?;
-    let (readonly_accounts, data) =
+    let (mut readonly_accounts, data) =
         ZeroCopyVecU8::<'a, ZPackedReadOnlyCompressedAccount>::from_bytes_at(data)?;
-    let (in_accounts, data) = ZeroCopyVecU8::<'a, CpiContextInAccount>::from_bytes_at(data)?;
-    let (out_accounts, data) = ZeroCopyVecU8::<'a, CpiContextOutAccount>::from_bytes_at(data)?;
-    let (output_data_len, mut data) = Ref::<&'a mut [u8], U16>::from_prefix(data)?;
-    let mut output_data = Vec::with_capacity(output_data_len.get() as usize);
-    for _ in 0..output_data_len.get() {
-        let (output_data_slice, inner_data) = ZeroCopySliceMut::<U16, u8>::from_bytes_at(data)?;
-        output_data.push(output_data_slice);
-        data = inner_data;
-    }
+    let (mut in_accounts, data) = ZeroCopyVecU8::<'a, CpiContextInAccount>::from_bytes_at(data)?;
+    let (mut out_accounts, data) = ZeroCopyVecU8::<'a, CpiContextOutAccount>::from_bytes_at(data)?;
+    let (mut output_data_len, mut data) = Ref::<&'a mut [u8], U16>::from_prefix(data)?;
+    let output_data = if CLEARED {
+        new_addresses.clear();
+        readonly_addresses.clear();
+        readonly_accounts.clear();
+        in_accounts.clear();
+        out_accounts.clear();
+        output_data_len.set(0);
+        // 65 CU
+        data.fill(0);
+        Vec::new()
+    } else {
+        let mut output_data = Vec::with_capacity(output_data_len.get() as usize);
+        for _ in 0..output_data_len.get() {
+            let (output_data_slice, inner_data) = ZeroCopySliceMut::<U16, u8>::from_bytes_at(data)?;
+            output_data.push(output_data_slice);
+            data = inner_data;
+        }
+        output_data
+    };
 
     Ok(ZCpiContextAccount {
         fee_payer,

programs/system/src/invoke_cpi/processor.rs

@@ -5,7 +5,9 @@ pub use crate::Result;
 use crate::{
     accounts::account_traits::{CpiContextAccountTrait, InvokeAccounts, SignerAccounts},
     context::WrappedInstructionData,
-    cpi_context::process_cpi_context::process_cpi_context,
+    cpi_context::{
+        process_cpi_context::process_cpi_context, state::deserialize_cpi_context_account_cleared,
+    },
     invoke_cpi::verify_signer::cpi_signer_checks,
     processor::process::process,
 };
@@ -59,16 +61,7 @@ pub fn process_invoke_cpi<
 
     // 4. clear cpi context account
     if cpi_context_inputs_len > 0 {
-        // TODO: reimplement this doesn't work anymore
-        // clear_cpi_context_account(accounts.get_cpi_context_account())?;
+        deserialize_cpi_context_account_cleared(accounts.get_cpi_context_account().unwrap())?;
     }
     Ok(())
 }
-
-/// Clear the CPI context account by setting the length to 0.
-pub fn clear_cpi_context_account(account_info: Option<&AccountInfo>) -> Result<()> {
-    let mut data = account_info.unwrap().try_borrow_mut_data()?;
-    let start_offset = 8 + 32 + 32;
-    data[start_offset..start_offset + 4].copy_from_slice(&[0u8, 0u8, 0u8, 0u8]);
-    Ok(())
-}