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

joncinque · web-flow · commit 129b35609adf · 2022-02-04T23:07:48.000+01:00
diff --git a/token/program-2022/src/processor.rs b/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())?;
-            if expected_decimals != mint.decimals {
+            let mint_data = mint_info.data.borrow();
+            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.delegated_amount = amount;
-
-        Account::pack(source_account, &mut source_account_info.data.borrow_mut())?;
+        source_account.base.delegate = COption::Some(*delegate_info.key);
+        source_account.base.delegated_amount = amount;
+        source_account.pack_base();
 
         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())?;
-        if source_account.is_frozen() {
+        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.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.delegated_amount = 0;
-
-        Account::pack(source_account, &mut source_account_info.data.borrow_mut())?;
+        source_account.base.delegate = COption::None;
+        source_account.base.delegated_amount = 0;
+        source_account.pack_base();
 
         Ok(())
     }
diff --git a/token/program-2022/tests/delegate.rs b/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;
+}
diff --git a/token/rust/src/token.rs b/token/rust/src/token.rs
