token-2022: impl all remaining cpi guard checks

also place all checks after validate_owner, for consistency
solana-labs · Oct 18, 2022 · 2c415e6 · 2c415e6
1 parent 5730242
commit 2c415e6
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 15 deletions.
diff --git a/token/program-2022/src/error.rs b/token/program-2022/src/error.rs
@@ -179,8 +179,8 @@ pub enum TokenError {
     /// CPI Guard is enabled, and a program attempted to approve a delegate
     #[error("CPI Guard is enabled, and a program attempted to approve a delegate")]
     CpiGuardApproveBlocked,
-    /// CPI Guard is enabled, and a program attempted to add or change an authority
-    #[error("CPI Guard is enabled, and a program attempted to add or change an authority")]
+    /// CPI Guard is enabled, and a program attempted to add or replace an authority
+    #[error("CPI Guard is enabled, and a program attempted to add or replace an authority")]
     CpiGuardSetAuthorityBlocked,
     /// Account ownership cannot be changed while CPI Guard is enabled
     #[error("Account ownership cannot be changed while CPI Guard is enabled")]

diff --git a/token/program-2022/src/processor.rs b/token/program-2022/src/processor.rs
@@ -350,19 +350,19 @@ impl Processor {
                 }
             }
             _ => {
-                if let Ok(cpi_guard) = source_account.get_extension::<CpiGuard>() {
-                    if cpi_guard.lock_cpi.into() && in_cpi() {
-                        return Err(TokenError::CpiGuardTransferBlocked.into());
-                    }
-                }
-
                 Self::validate_owner(
                     program_id,
                     &source_account.base.owner,
                     authority_info,
                     authority_info_data_len,
                     account_info_iter.as_slice(),
                 )?;
+
+                if let Ok(cpi_guard) = source_account.get_extension::<CpiGuard>() {
+                    if cpi_guard.lock_cpi.into() && in_cpi() {
+                        return Err(TokenError::CpiGuardTransferBlocked.into());
+                    }
+                }
             }
         };
 
@@ -485,6 +485,12 @@ impl Processor {
             account_info_iter.as_slice(),
         )?;
 
+        if let Ok(cpi_guard) = source_account.get_extension::<CpiGuard>() {
+            if cpi_guard.lock_cpi.into() && in_cpi() {
+                return Err(TokenError::CpiGuardApproveBlocked.into());
+            }
+        }
+
         source_account.base.delegate = COption::Some(*delegate_info.key);
         source_account.base.delegated_amount = amount;
         source_account.pack_base();
@@ -558,6 +564,14 @@ impl Processor {
                         return Err(TokenError::ImmutableOwner.into());
                     }
 
+                    if let Ok(cpi_guard) = account.get_extension::<CpiGuard>() {
+                        if cpi_guard.lock_cpi.into() && in_cpi() {
+                            return Err(TokenError::CpiGuardSetAuthorityBlocked.into());
+                        } else if cpi_guard.lock_cpi.into() {
+                            return Err(TokenError::CpiGuardOwnerChangeBlocked.into());
+                        }
+                    }
+
                     if let COption::Some(authority) = new_authority {
                         account.base.owner = authority;
                     } else {
@@ -580,6 +594,13 @@ impl Processor {
                         authority_info_data_len,
                         account_info_iter.as_slice(),
                     )?;
+
+                    if let Ok(cpi_guard) = account.get_extension::<CpiGuard>() {
+                        if cpi_guard.lock_cpi.into() && in_cpi() && new_authority != COption::None {
+                            return Err(TokenError::CpiGuardSetAuthorityBlocked.into());
+                        }
+                    }
+
                     account.base.close_authority = new_authority;
                 }
                 _ => {
@@ -838,13 +859,21 @@ impl Processor {
                         source_account.base.delegate = COption::None;
                     }
                 }
-                _ => Self::validate_owner(
-                    program_id,
-                    &source_account.base.owner,
-                    authority_info,
-                    authority_info_data_len,
-                    account_info_iter.as_slice(),
-                )?,
+                _ => {
+                    if let Ok(cpi_guard) = source_account.get_extension::<CpiGuard>() {
+                        if cpi_guard.lock_cpi.into() && in_cpi() {
+                            return Err(TokenError::CpiGuardTransferBlocked.into());
+                        }
+                    }
+
+                    Self::validate_owner(
+                        program_id,
+                        &source_account.base.owner,
+                        authority_info,
+                        authority_info_data_len,
+                        account_info_iter.as_slice(),
+                    )?;
+                }
             }
         }
 
@@ -898,6 +927,15 @@ impl Processor {
                 .base
                 .is_owned_by_system_program_or_incinerator()
             {
+                if let Ok(cpi_guard) = source_account.get_extension::<CpiGuard>() {
+                    if cpi_guard.lock_cpi.into()
+                        && in_cpi()
+                        && !cmp_pubkeys(destination_account_info.key, &source_account.base.owner)
+                    {
+                        return Err(TokenError::CpiGuardCloseAccountBlocked.into());
+                    }
+                }
+
                 Self::validate_owner(
                     program_id,
                     &authority,