token-2022: Support approve / revoke with extensions (#2870)

token/program-2022/src/processor.rs

@@ -434,35 +434,37 @@ impl Processor {
         let owner_info = next_account_info(account_info_iter)?;
         let owner_info_data_len = owner_info.data_len();
 
-        let mut source_account = Account::unpack(&source_account_info.data.borrow())?;
+        let mut source_account_data = source_account_info.data.borrow_mut();
+        let mut source_account =
+            StateWithExtensionsMut::<Account>::unpack(&mut source_account_data)?;
 
-        if source_account.is_frozen() {
+        if source_account.base.is_frozen() {
             return Err(TokenError::AccountFrozen.into());
         }
 
         if let Some((mint_info, expected_decimals)) = expected_mint_info {
-            if !cmp_pubkeys(&source_account.mint, mint_info.key) {
+            if !cmp_pubkeys(&source_account.base.mint, mint_info.key) {
                 return Err(TokenError::MintMismatch.into());
             }
 
-            let mint = Mint::unpack(&mint_info.data.borrow_mut())?;
+            let mint_data = mint_info.data.borrow();
-            if expected_decimals != mint.decimals {
+            let mint = StateWithExtensions::<Mint>::unpack(&mint_data)?;
+            if expected_decimals != mint.base.decimals {
                 return Err(TokenError::MintDecimalsMismatch.into());
             }
         }
 
         Self::validate_owner(
             program_id,
-            &source_account.owner,
+            &source_account.base.owner,
             owner_info,
             owner_info_data_len,
             account_info_iter.as_slice(),
         )?;
 
-        source_account.delegate = COption::Some(*delegate_info.key);
+        source_account.base.delegate = COption::Some(*delegate_info.key);
-        source_account.delegated_amount = amount;
+        source_account.base.delegated_amount = amount;
-
+        source_account.pack_base();
-        Account::pack(source_account, &mut source_account_info.data.borrow_mut())?;
 
         Ok(())
     }
@@ -474,28 +476,29 @@ impl Processor {
         let authority_info = next_account_info(account_info_iter)?;
         let authority_info_data_len = authority_info.data_len();
 
-        let mut source_account = Account::unpack(&source_account_info.data.borrow())?;
+        let mut source_account_data = source_account_info.data.borrow_mut();
-        if source_account.is_frozen() {
+        let mut source_account =
+            StateWithExtensionsMut::<Account>::unpack(&mut source_account_data)?;
+        if source_account.base.is_frozen() {
             return Err(TokenError::AccountFrozen.into());
         }
 
         Self::validate_owner(
             program_id,
-            match source_account.delegate {
+            match source_account.base.delegate {
                 COption::Some(ref delegate) if cmp_pubkeys(authority_info.key, delegate) => {
                     delegate
                 }
-                _ => &source_account.owner,
+                _ => &source_account.base.owner,
             },
             authority_info,
             authority_info_data_len,
             account_info_iter.as_slice(),
         )?;
 
-        source_account.delegate = COption::None;
+        source_account.base.delegate = COption::None;
-        source_account.delegated_amount = 0;
+        source_account.base.delegated_amount = 0;
-
+        source_account.pack_base();
-        Account::pack(source_account, &mut source_account_info.data.borrow_mut())?;
 
         Ok(())
     }

token/program-2022/tests/delegate.rs

@@ -0,0 +1,268 @@
+#![cfg(feature = "test-bpf")]
+
+mod program_test;
+use {
+    program_test::{TestContext, TokenContext},
+    solana_program_test::tokio,
+    solana_sdk::{
+        instruction::InstructionError, pubkey::Pubkey, signature::Signer, signer::keypair::Keypair,
+        transaction::TransactionError, transport::TransportError,
+    },
+    spl_token_2022::error::TokenError,
+    spl_token_client::token::{ExtensionInitializationParams, TokenError as TokenClientError},
+};
+
+#[derive(PartialEq)]
+enum TransferMode {
+    All,
+    CheckedOnly,
+}
+
+#[derive(PartialEq)]
+enum ApproveMode {
+    Unchecked,
+    Checked,
+}
+
+#[derive(PartialEq)]
+enum OwnerMode {
+    SelfOwned,
+    External,
+}
+
+async fn run_basic(
+    context: TestContext,
+    owner_mode: OwnerMode,
+    transfer_mode: TransferMode,
+    approve_mode: ApproveMode,
+) {
+    let TokenContext {
+        decimals,
+        mint_authority,
+        token,
+        alice,
+        bob,
+        ..
+    } = context.token_context.unwrap();
+
+    let alice_account = match owner_mode {
+        OwnerMode::SelfOwned => token
+            .create_auxiliary_token_account(&alice, &alice.pubkey())
+            .await
+            .unwrap(),
+        OwnerMode::External => {
+            let alice_account = Keypair::new();
+            token
+                .create_auxiliary_token_account(&alice_account, &alice.pubkey())
+                .await
+                .unwrap()
+        }
+    };
+    let bob_account = Keypair::new();
+    let bob_account = token
+        .create_auxiliary_token_account(&bob_account, &bob.pubkey())
+        .await
+        .unwrap();
+
+    // mint tokens
+    let amount = 100;
+    token
+        .mint_to(&alice_account, &mint_authority, amount)
+        .await
+        .unwrap();
+
+    // delegate to bob
+    let delegated_amount = 10;
+    match approve_mode {
+        ApproveMode::Unchecked => token
+            .approve(&alice_account, &bob.pubkey(), &alice, delegated_amount)
+            .await
+            .unwrap(),
+        ApproveMode::Checked => token
+            .approve_checked(
+                &alice_account,
+                &bob.pubkey(),
+                &alice,
+                delegated_amount,
+                decimals,
+            )
+            .await
+            .unwrap(),
+    }
+
+    // transfer too much is not ok
+    let error = token
+        .transfer_checked(
+            &alice_account,
+            &bob_account,
+            &bob,
+            delegated_amount + 1,
+            decimals,
+        )
+        .await
+        .unwrap_err();
+    assert_eq!(
+        error,
+        TokenClientError::Client(Box::new(TransportError::TransactionError(
+            TransactionError::InstructionError(
+                0,
+                InstructionError::Custom(TokenError::InsufficientFunds as u32)
+            )
+        )))
+    );
+
+    // transfer is ok
+    if transfer_mode == TransferMode::All {
+        token
+            .transfer_unchecked(&alice_account, &bob_account, &bob, 1)
+            .await
+            .unwrap();
+    }
+
+    token
+        .transfer_checked(&alice_account, &bob_account, &bob, 1, decimals)
+        .await
+        .unwrap();
+
+    // burn is ok
+    token.burn(&alice_account, &bob, 1).await.unwrap();
+    token
+        .burn_checked(&alice_account, &bob, 1, decimals)
+        .await
+        .unwrap();
+
+    // wrong signer
+    let error = token
+        .transfer_checked(&alice_account, &bob_account, &Keypair::new(), 1, decimals)
+        .await
+        .unwrap_err();
+    assert_eq!(
+        error,
+        TokenClientError::Client(Box::new(TransportError::TransactionError(
+            TransactionError::InstructionError(
+                0,
+                InstructionError::Custom(TokenError::OwnerMismatch as u32)
+            )
+        )))
+    );
+
+    // revoke
+    token.revoke(&alice_account, &alice).await.unwrap();
+
+    // now fails
+    let error = token
+        .transfer_checked(&alice_account, &bob_account, &bob, 1, decimals)
+        .await
+        .unwrap_err();
+    assert_eq!(
+        error,
+        TokenClientError::Client(Box::new(TransportError::TransactionError(
+            TransactionError::InstructionError(
+                0,
+                InstructionError::Custom(TokenError::OwnerMismatch as u32)
+            )
+        )))
+    );
+}
+
+#[tokio::test]
+async fn basic() {
+    let mut context = TestContext::new().await;
+    context.init_token_with_mint(vec![]).await.unwrap();
+    run_basic(
+        context,
+        OwnerMode::External,
+        TransferMode::All,
+        ApproveMode::Unchecked,
+    )
+    .await;
+}
+
+#[tokio::test]
+async fn basic_checked() {
+    let mut context = TestContext::new().await;
+    context.init_token_with_mint(vec![]).await.unwrap();
+    run_basic(
+        context,
+        OwnerMode::External,
+        TransferMode::All,
+        ApproveMode::Checked,
+    )
+    .await;
+}
+
+#[tokio::test]
+async fn basic_self_owned() {
+    let mut context = TestContext::new().await;
+    context.init_token_with_mint(vec![]).await.unwrap();
+    run_basic(
+        context,
+        OwnerMode::SelfOwned,
+        TransferMode::All,
+        ApproveMode::Checked,
+    )
+    .await;
+}
+
+#[tokio::test]
+async fn basic_with_extension() {
+    let mut context = TestContext::new().await;
+    context
+        .init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig {
+            transfer_fee_config_authority: Some(Pubkey::new_unique()),
+            withdraw_withheld_authority: Some(Pubkey::new_unique()),
+            transfer_fee_basis_points: 100u16,
+            maximum_fee: 1_000u64,
+        }])
+        .await
+        .unwrap();
+    run_basic(
+        context,
+        OwnerMode::External,
+        TransferMode::CheckedOnly,
+        ApproveMode::Unchecked,
+    )
+    .await;
+}
+
+#[tokio::test]
+async fn basic_with_extension_checked() {
+    let mut context = TestContext::new().await;
+    context
+        .init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig {
+            transfer_fee_config_authority: Some(Pubkey::new_unique()),
+            withdraw_withheld_authority: Some(Pubkey::new_unique()),
+            transfer_fee_basis_points: 100u16,
+            maximum_fee: 1_000u64,
+        }])
+        .await
+        .unwrap();
+    run_basic(
+        context,
+        OwnerMode::External,
+        TransferMode::CheckedOnly,
+        ApproveMode::Checked,
+    )
+    .await;
+}
+
+#[tokio::test]
+async fn basic_self_owned_with_extension() {
+    let mut context = TestContext::new().await;
+    context
+        .init_token_with_mint(vec![ExtensionInitializationParams::TransferFeeConfig {
+            transfer_fee_config_authority: Some(Pubkey::new_unique()),
+            withdraw_withheld_authority: Some(Pubkey::new_unique()),
+            transfer_fee_basis_points: 100u16,
+            maximum_fee: 1_000u64,
+        }])
+        .await
+        .unwrap();
+    run_basic(
+        context,
+        OwnerMode::SelfOwned,
+        TransferMode::CheckedOnly,
+        ApproveMode::Checked,
+    )
+    .await;
+}

token/rust/src/token.rs

@@ -505,6 +505,71 @@ where
         .await
     }
 
+    /// Approve a delegate to spend tokens
+    pub async fn approve<S2: Signer>(
+        &self,
+        source: &Pubkey,
+        delegate: &Pubkey,
+        authority: &S2,
+        amount: u64,
+    ) -> TokenResult<T::Output> {
+        self.process_ixs(
+            &[instruction::approve(
+                &self.program_id,
+                source,
+                delegate,
+                &authority.pubkey(),
+                &[],
+                amount,
+            )?],
+            &[authority],
+        )
+        .await
+    }
+
+    /// Approve a delegate to spend tokens, with decimal check
+    pub async fn approve_checked<S2: Signer>(
+        &self,
+        source: &Pubkey,
+        delegate: &Pubkey,
+        authority: &S2,
+        amount: u64,
+        decimals: u8,
+    ) -> TokenResult<T::Output> {
+        self.process_ixs(
+            &[instruction::approve_checked(
+                &self.program_id,
+                source,
+                &self.pubkey,
+                delegate,
+                &authority.pubkey(),
+                &[],
+                amount,
+                decimals,
+            )?],
+            &[authority],
+        )
+        .await
+    }
+
+    /// Revoke a delegate
+    pub async fn revoke<S2: Signer>(
+        &self,
+        source: &Pubkey,
+        authority: &S2,
+    ) -> TokenResult<T::Output> {
+        self.process_ixs(
+            &[instruction::revoke(
+                &self.program_id,
+                source,
+                &authority.pubkey(),
+                &[],
+            )?],
+            &[authority],
+        )
+        .await
+    }
+
     /// Close account into another
     pub async fn close_account<S2: Signer>(
         &self,