token-2022: Add freeze / thaw through StateWithExtensions (#2848)

Author: joncinque

token/program-2022/src/processor.rs

@@ -882,19 +882,23 @@ 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 freeze && source_account.is_frozen() || !freeze && !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 freeze && source_account.base.is_frozen() || !freeze && !source_account.base.is_frozen()
+        {
             return Err(TokenError::InvalidState.into());
         }
-        if source_account.is_native() {
+        if source_account.base.is_native() {
             return Err(TokenError::NativeNotSupported.into());
         }
-        if !cmp_pubkeys(mint_info.key, &source_account.mint) {
+        if !cmp_pubkeys(mint_info.key, &source_account.base.mint) {
             return Err(TokenError::MintMismatch.into());
         }
 
-        let mint = Mint::unpack(&mint_info.data.borrow_mut())?;
-        match mint.freeze_authority {
+        let mint_data = mint_info.data.borrow();
+        let mint = StateWithExtensions::<Mint>::unpack(&mint_data)?;
+        match mint.base.freeze_authority {
             COption::Some(authority) => Self::validate_owner(
                 program_id,
                 &authority,
@@ -905,13 +909,13 @@ impl Processor {
             COption::None => Err(TokenError::MintCannotFreeze.into()),
         }?;
 
-        source_account.state = if freeze {
+        source_account.base.state = if freeze {
             AccountState::Frozen
         } else {
             AccountState::Initialized
         };
 
-        Account::pack(source_account, &mut source_account_info.data.borrow_mut())?;
+        source_account.pack_base();
 
         Ok(())
     }

token/program-2022/tests/freeze.rs

@@ -0,0 +1,44 @@
+#![cfg(feature = "test-bpf")]
+
+mod program_test;
+use {
+    program_test::{TestContext, TokenContext},
+    solana_program_test::tokio,
+    solana_sdk::{signature::Signer, signer::keypair::Keypair},
+    spl_token_2022::state::AccountState,
+};
+
+#[tokio::test]
+async fn basic() {
+    let mut context = TestContext::new().await;
+    context.init_token_with_freezing_mint(vec![]).await.unwrap();
+    let TokenContext {
+        freeze_authority,
+        token,
+        alice,
+        ..
+    } = context.token_context.unwrap();
+    let freeze_authority = freeze_authority.unwrap();
+
+    let account = Keypair::new();
+    let account = token
+        .create_auxiliary_token_account(&account, &alice.pubkey())
+        .await
+        .unwrap();
+    let state = token.get_account_info(&account).await.unwrap();
+    assert_eq!(state.base.state, AccountState::Initialized);
+
+    token
+        .freeze_account(&account, &freeze_authority)
+        .await
+        .unwrap();
+    let state = token.get_account_info(&account).await.unwrap();
+    assert_eq!(state.base.state, AccountState::Frozen);
+
+    token
+        .thaw_account(&account, &freeze_authority)
+        .await
+        .unwrap();
+    let state = token.get_account_info(&account).await.unwrap();
+    assert_eq!(state.base.state, AccountState::Initialized);
+}

token/rust/src/token.rs

@@ -481,6 +481,44 @@ where
         .await
     }
 
+    /// Freeze a token account
+    pub async fn freeze_account<S2: Signer>(
+        &self,
+        account: &Pubkey,
+        authority: &S2,
+    ) -> TokenResult<T::Output> {
+        self.process_ixs(
+            &[instruction::freeze_account(
+                &self.program_id,
+                account,
+                &self.pubkey,
+                &authority.pubkey(),
+                &[],
+            )?],
+            &[authority],
+        )
+        .await
+    }
+
+    /// Thaw / unfreeze a token account
+    pub async fn thaw_account<S2: Signer>(
+        &self,
+        account: &Pubkey,
+        authority: &S2,
+    ) -> TokenResult<T::Output> {
+        self.process_ixs(
+            &[instruction::thaw_account(
+                &self.program_id,
+                account,
+                &self.pubkey,
+                &authority.pubkey(),
+                &[],
+            )?],
+            &[authority],
+        )
+        .await
+    }
+
     /// Set transfer fee
     pub async fn set_transfer_fee<S2: Signer>(
         &self,
@@ -538,4 +576,25 @@ where
         )
         .await
     }
+
+    /// Withdraw withheld tokens from mint
+    pub async fn withdraw_withheld_tokens_from_mint<S2: Signer>(
+        &self,
+        destination: &Pubkey,
+        authority: &S2,
+    ) -> TokenResult<T::Output> {
+        self.process_ixs(
+            &[
+                transfer_fee::instruction::withdraw_withheld_tokens_from_mint(
+                    &self.program_id,
+                    &self.pubkey,
+                    destination,
+                    &authority.pubkey(),
+                    &[],
+                )?,
+            ],
+            &[authority],
+        )
+        .await
+    }
 }