fix: ctoken cpi context (#1870)

* fix: check cpi context account

* fix tests

Author: ananas-block

program-tests/system-cpi-test/tests/test.rs

@@ -1387,7 +1387,7 @@ async fn test_approve_revoke_burn_freeze_thaw_with_cpi_context() {
             .value
             .items[0]
             .clone();
-        perform_with_input_accounts(
+        let res = perform_with_input_accounts(
             &mut test_indexer,
             &mut rpc,
             &payer,
@@ -1397,20 +1397,16 @@ async fn test_approve_revoke_burn_freeze_thaw_with_cpi_context() {
             u32::MAX,
             WithInputAccountsMode::Burn,
         )
-        .await
+        .await;
+        assert_rpc_error(
+            res,
+            0,
+            light_compressed_token::ErrorCode::CpiContextSetNotUsable.into(),
+        )
         .unwrap();
-        let compressed_token_data = test_indexer
-            .get_compressed_token_accounts_by_owner(&payer.pubkey(), None, None)
-            .await
-            .unwrap()
-            .value
-            .items[0]
-            .clone();
-        let mut ref_data = ref_compressed_token_data.token.clone();
-        ref_data.amount = 1;
-        assert_eq!(compressed_token_data.token, ref_data);
     }
 }
+
 /// Test:
 /// 1. Cannot create an address in a program owned address Merkle tree owned by a different program (InvalidMerkleTreeOwner)
 /// 2. Cannot create a compressed account in a program owned state Merkle tree owned by a different program (InvalidMerkleTreeOwner)

programs/compressed-token/src/burn.rs

@@ -36,6 +36,7 @@ pub fn process_burn<'a, 'b, 'c, 'info: 'b + 'c>(
 ) -> Result<()> {
     let inputs: CompressedTokenInstructionDataBurn =
         CompressedTokenInstructionDataBurn::deserialize(&mut inputs.as_slice())?;
+    crate::check_cpi_context(&inputs.cpi_context)?;
     burn_spl_from_pool_pda(&ctx, &inputs)?;
     let mint = ctx.accounts.mint.key();
     let (compressed_input_accounts, output_compressed_accounts) =

programs/compressed-token/src/delegation.rs

@@ -49,6 +49,7 @@ pub fn process_approve<'a, 'b, 'c, 'info: 'b + 'c>(
 ) -> Result<()> {
     let inputs: CompressedTokenInstructionDataApprove =
         CompressedTokenInstructionDataApprove::deserialize(&mut inputs.as_slice())?;
+    // CPI context check not needed: delegation operations don't modify Solana account state
     let (compressed_input_accounts, output_compressed_accounts) =
         create_input_and_output_accounts_approve(
             &inputs,
@@ -183,6 +184,7 @@ pub fn process_revoke<'a, 'b, 'c, 'info: 'b + 'c>(
 ) -> Result<()> {
     let inputs: CompressedTokenInstructionDataRevoke =
         CompressedTokenInstructionDataRevoke::deserialize(&mut inputs.as_slice())?;
+    // CPI context check not needed: delegation operations don't modify Solana account state
     let (compressed_input_accounts, output_compressed_accounts) =
         create_input_and_output_accounts_revoke(
             &inputs,

programs/compressed-token/src/freeze.rs

@@ -42,6 +42,7 @@ pub fn process_freeze_or_thaw<
 ) -> Result<()> {
     let inputs: CompressedTokenInstructionDataFreeze =
         CompressedTokenInstructionDataFreeze::deserialize(&mut inputs.as_slice())?;
+    // CPI context check not needed: freeze/thaw operations don't modify Solana account state
     let (compressed_input_accounts, output_compressed_accounts) =
         create_input_and_output_accounts_freeze_or_thaw::<FROZEN_INPUTS, FROZEN_OUTPUTS>(
             &inputs,

programs/compressed-token/src/lib.rs

@@ -148,6 +148,10 @@ pub mod light_compressed_token {
         inputs.extend_from_slice(&[0u8; 1]);
         let inputs: CompressedTokenInstructionDataTransfer =
             CompressedTokenInstructionDataTransfer::deserialize(&mut inputs.as_slice())?;
+        // Only check CPI context if we're compressing or decompressing (modifying Solana account state)
+        if inputs.compress_or_decompress_amount.is_some() {
+            check_cpi_context(&inputs.cpi_context)?;
+        }
         process_transfer::process_transfer(ctx, inputs)
     }
 
@@ -276,4 +280,17 @@ pub enum ErrorCode {
     NoMatchingBumpFound,
     NoAmount,
     AmountsAndAmountProvided,
+    #[msg("Cpi context set and set first is not usable with burn, compression(transfer ix) or decompress(transfer).")]
+    CpiContextSetNotUsable,
+}
+
+/// Checks if CPI context usage is valid for the current instruction
+/// Throws an error if cpi_context is Some and (set_context OR first_set_context is true)
+fn check_cpi_context(cpi_context: &Option<CompressedCpiContext>) -> Result<()> {
+    if let Some(ctx) = cpi_context {
+        if ctx.set_context || ctx.first_set_context {
+            return Err(ErrorCode::CpiContextSetNotUsable.into());
+        }
+    }
+    Ok(())
 }

programs/compressed-token/src/process_compress_spl_token_account.rs

@@ -15,6 +15,7 @@ pub fn process_compress_spl_token_account<'info>(
     remaining_amount: Option<u64>,
     cpi_context: Option<CompressedCpiContext>,
 ) -> Result<()> {
+    crate::check_cpi_context(&cpi_context)?;
     let compression_token_account =
         if let Some(token_account) = ctx.accounts.compress_or_decompress_token_account.as_ref() {
             token_account