token-2022: Prohibit self-delegated transfers with CPI Guard (#6724)

solana-labs · May 14, 2024 · dc85c98 · dc85c98
1 parent d1c03cc
commit dc85c98
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 0 deletions.
diff --git a/token/program-2022-test/tests/cpi_guard.rs b/token/program-2022-test/tests/cpi_guard.rs
@@ -280,6 +280,27 @@ async fn test_cpi_guard_transfer() {
         let alice_state = token_obj.get_account_info(&alice.pubkey()).await.unwrap();
         assert_eq!(alice_state.base.amount, amount);
 
+        // delegate-auth cpi transfer to self does not work
+        token_obj
+            .approve(
+                &alice.pubkey(),
+                &alice.pubkey(),
+                &alice.pubkey(),
+                1,
+                &[&alice],
+            )
+            .await
+            .unwrap();
+
+        let error = token_obj
+            .process_ixs(&[mk_transfer(alice.pubkey(), do_checked)], &[&alice])
+            .await
+            .unwrap_err();
+        assert_eq!(error, client_error(TokenError::CpiGuardTransferBlocked));
+
+        let alice_state = token_obj.get_account_info(&alice.pubkey()).await.unwrap();
+        assert_eq!(alice_state.base.amount, amount);
+
         // delegate-auth cpi transfer with cpi guard works
         token_obj
             .approve(

diff --git a/token/program-2022/src/processor.rs b/token/program-2022/src/processor.rs
@@ -402,6 +402,15 @@ impl Processor {
                     authority_info_data_len,
                     account_info_iter.as_slice(),
                 )?;
+                if let Ok(cpi_guard) = source_account.get_extension::<CpiGuard>() {
+                    // If delegated to self, don't allow a transfer with CPI Guard
+                    if delegate == source_account.base.owner
+                        && cpi_guard.lock_cpi.into()
+                        && in_cpi()
+                    {
+                        return Err(TokenError::CpiGuardTransferBlocked.into());
+                    }
+                }
                 let delegated_amount = u64::from(source_account.base.delegated_amount);
                 if delegated_amount < amount {
                     return Err(TokenError::InsufficientFunds.into());