[program] Add cpi guard logic in confidential extensions (#819)

add cpi guard logic in confidential extensions

Author: samkim-crypto

program/src/extension/confidential_mint_burn/processor.rs

@@ -1,5 +1,13 @@
+#[cfg(feature = "zk-ops")]
+use {
+    crate::{
+        check_auditor_ciphertext,
+        extension::confidential_mint_burn::verify_proof::{verify_burn_proof, verify_mint_proof},
+    },
+    spl_token_confidential_transfer_ciphertext_arithmetic as ciphertext_arithmetic,
+};
 use {
-    crate::processor::Processor,
+    crate::{extension::cpi_guard::in_cpi, processor::Processor},
     solana_account_info::{next_account_info, AccountInfo},
     solana_msg::msg,
     solana_program_error::{ProgramError, ProgramResult},
@@ -26,6 +34,7 @@ use {
                 ConfidentialMintBurn,
             },
             confidential_transfer::{ConfidentialTransferAccount, ConfidentialTransferMint},
+            cpi_guard::CpiGuard,
             pausable::PausableConfig,
             BaseStateWithExtensions, BaseStateWithExtensionsMut, PodStateWithExtensionsMut,
         },
@@ -34,14 +43,6 @@ use {
     },
     spl_token_confidential_transfer_proof_extraction::instruction::verify_and_extract_context,
 };
-#[cfg(feature = "zk-ops")]
-use {
-    crate::{
-        check_auditor_ciphertext,
-        extension::confidential_mint_burn::verify_proof::{verify_burn_proof, verify_mint_proof},
-    },
-    spl_token_confidential_transfer_ciphertext_arithmetic as ciphertext_arithmetic,
-};
 
 /// Processes an [`InitializeMint`] instruction.
 fn process_initialize_mint(accounts: &[AccountInfo], data: &InitializeMintData) -> ProgramResult {
@@ -337,6 +338,17 @@ fn process_confidential_burn(
         account_info_iter.as_slice(),
     )?;
 
+    if let Ok(cpi_guard) = token_account.get_extension::<CpiGuard>() {
+        // Blocks all cases where the authority has signed if CPI Guard is
+        // enabled, including:
+        // * the account is delegated to the owner
+        // * the account owner is the permanent delegate
+        if *authority_info.key == token_account.base.owner && cpi_guard.lock_cpi.into() && in_cpi()
+        {
+            return Err(TokenError::CpiGuardBurnBlocked.into());
+        }
+    }
+
     if token_account.base.is_frozen() {
         return Err(TokenError::AccountFrozen.into());
     }

program/src/extension/confidential_transfer/processor.rs

@@ -11,7 +11,7 @@ use {
     crate::{
         check_elgamal_registry_program_account,
         extension::{
-            confidential_transfer::verify_proof::*,
+            confidential_transfer::verify_proof::*, cpi_guard::in_cpi,
             memo_transfer::check_previous_sibling_instruction_is_memo,
         },
         processor::Processor,
@@ -44,6 +44,7 @@ use {
                 ConfidentialTransferFeeAmount, ConfidentialTransferFeeConfig,
                 EncryptedWithheldAmount,
             },
+            cpi_guard::CpiGuard,
             memo_transfer::memo_required,
             pausable::PausableConfig,
             set_account_type,
@@ -872,6 +873,17 @@ fn process_source_for_transfer(
         signers,
     )?;
 
+    if let Ok(cpi_guard) = token_account.get_extension::<CpiGuard>() {
+        // Blocks all cases where the authority has signed if CPI Guard is
+        // enabled, including:
+        // * the account is delegated to the owner
+        // * the account owner is the permanent delegate
+        if *authority_info.key == token_account.base.owner && cpi_guard.lock_cpi.into() && in_cpi()
+        {
+            return Err(TokenError::CpiGuardTransferBlocked.into());
+        }
+    }
+
     if token_account.base.is_frozen() {
         return Err(TokenError::AccountFrozen.into());
     }
@@ -1008,6 +1020,17 @@ fn process_source_for_transfer_with_fee(
         signers,
     )?;
 
+    if let Ok(cpi_guard) = token_account.get_extension::<CpiGuard>() {
+        // Blocks all cases where the authority has signed if CPI Guard is
+        // enabled, including:
+        // * the account is delegated to the owner
+        // * the account owner is the permanent delegate
+        if *authority_info.key == token_account.base.owner && cpi_guard.lock_cpi.into() && in_cpi()
+        {
+            return Err(TokenError::CpiGuardTransferBlocked.into());
+        }
+    }
+
     if token_account.base.is_frozen() {
         return Err(TokenError::AccountFrozen.into());
     }