token-2022: Sync native balance during realloc (#4582)

* token-2022: Sync native balance during realloc

* Fix non-empty case and add tests, address other feedback

* Don't sync native, add tests for empty extensions

* Simplify code, add more test cases

* Rename, don't always unpack, use saturating_add

* Rename to new_rent_exempt_reserve

* Rename variables

* Go back to checked math, reorder

Author: joncinque

Cargo.lock

@@ -29,3 +29,4 @@ spl-instruction-padding = { version = "0.1.0", path="../../instruction-padding/p
 spl-token-client = { version = "0.5", path = "../client" }
 spl-transfer-hook-example = { version = "0.1", path="../transfer-hook-example", features = ["no-entrypoint"] }
 spl-transfer-hook-interface = { version = "0.1", path="../transfer-hook-interface" }
+test-case = "3.1"

token/program-2022-test/Cargo.toml

@@ -5,12 +5,19 @@ use {
     program_test::{TestContext, TokenContext},
     solana_program_test::tokio,
     solana_sdk::{
-        instruction::InstructionError, program_option::COption, pubkey::Pubkey, signature::Signer,
-        signer::keypair::Keypair, transaction::TransactionError, transport::TransportError,
+        instruction::InstructionError,
+        program_option::COption,
+        pubkey::Pubkey,
+        signature::Signer,
+        signer::keypair::Keypair,
+        system_instruction,
+        transaction::{Transaction, TransactionError},
+        transport::TransportError,
     },
     spl_token_2022::{error::TokenError, extension::ExtensionType, state::Account},
     spl_token_client::token::{ExtensionInitializationParams, TokenError as TokenClientError},
     std::convert::TryInto,
+    test_case::test_case,
 };
 
 #[tokio::test]
@@ -188,3 +195,90 @@ async fn reallocate_without_current_extension_knowledge() {
         ])
     );
 }
+
+#[test_case(&[ExtensionType::CpiGuard], 1_000_000_000, true ; "transfer more than new rent and sync")]
+#[test_case(&[ExtensionType::CpiGuard], 1_000_000_000, false ; "transfer more than new rent")]
+#[test_case(&[ExtensionType::CpiGuard], 1, true ; "transfer less than new rent and sync")]
+#[test_case(&[ExtensionType::CpiGuard], 1, false ; "transfer less than new rent")]
+#[test_case(&[ExtensionType::CpiGuard], 0, false ; "no transfer with extension")]
+#[test_case(&[], 1_000_000_000, true ; "transfer lamports and sync without extension")]
+#[test_case(&[], 1_000_000_000, false ; "transfer lamports without extension")]
+#[test_case(&[], 0, false ; "no transfer without extension")]
+#[tokio::test]
+async fn reallocate_updates_native_rent_exemption(
+    extensions: &[ExtensionType],
+    transfer_lamports: u64,
+    sync_native: bool,
+) {
+    let mut context = TestContext::new().await;
+    context.init_token_with_native_mint().await.unwrap();
+    let TokenContext { token, alice, .. } = context.token_context.unwrap();
+    let context = context.context.clone();
+
+    let alice_account = Keypair::new();
+    token
+        .create_auxiliary_token_account(&alice_account, &alice.pubkey())
+        .await
+        .unwrap();
+    let alice_account = alice_account.pubkey();
+
+    // transfer more lamports
+    if transfer_lamports > 0 {
+        let mut context = context.lock().await;
+        let instructions = vec![system_instruction::transfer(
+            &context.payer.pubkey(),
+            &alice_account,
+            transfer_lamports,
+        )];
+        let tx = Transaction::new_signed_with_payer(
+            &instructions,
+            Some(&context.payer.pubkey()),
+            &[&context.payer],
+            context.last_blockhash,
+        );
+        context.banks_client.process_transaction(tx).await.unwrap();
+    }
+
+    // amount in the account should be 0 no matter what
+    let account_info = token.get_account_info(&alice_account).await.unwrap();
+    assert_eq!(account_info.base.amount, 0);
+
+    if sync_native {
+        token.sync_native(&alice_account).await.unwrap();
+        let account_info = token.get_account_info(&alice_account).await.unwrap();
+        assert_eq!(account_info.base.amount, transfer_lamports);
+    }
+
+    let token_account = token.get_account_info(&alice_account).await.unwrap();
+    let pre_amount = token_account.base.amount;
+    let pre_rent_exempt_reserve = token_account.base.is_native.unwrap();
+
+    // reallocate resizes account to accommodate new extension
+    token
+        .reallocate(&alice_account, &alice.pubkey(), extensions, &[&alice])
+        .await
+        .unwrap();
+
+    let account = token.get_account(alice_account).await.unwrap();
+    assert_eq!(
+        account.data.len(),
+        ExtensionType::get_account_len::<Account>(extensions)
+    );
+    let expected_rent_exempt_reserve = {
+        let mut context = context.lock().await;
+        let rent = context.banks_client.get_rent().await.unwrap();
+        rent.minimum_balance(account.data.len())
+    };
+    let token_account = token.get_account_info(&alice_account).await.unwrap();
+    let post_amount = token_account.base.amount;
+    let post_rent_exempt_reserve = token_account.base.is_native.unwrap();
+    // amount of lamports should be totally unchanged
+    assert_eq!(pre_amount, post_amount);
+    // but rent exempt reserve should change
+    assert_eq!(post_rent_exempt_reserve, expected_rent_exempt_reserve);
+    if extensions.is_empty() {
+        assert_eq!(pre_rent_exempt_reserve, post_rent_exempt_reserve);
+    } else {
+        assert!(pre_rent_exempt_reserve < post_rent_exempt_reserve);
+    }
+}

token/program-2022-test/tests/reallocate.rs

@@ -3,7 +3,7 @@ use {
         error::TokenError,
         extension::{
             set_account_type, AccountType, BaseStateWithExtensions, ExtensionType,
-            StateWithExtensions,
+            StateWithExtensions, StateWithExtensionsMut,
         },
         processor::Processor,
         state::Account,
@@ -13,6 +13,7 @@ use {
         entrypoint::ProgramResult,
         msg,
         program::invoke,
+        program_option::COption,
         pubkey::Pubkey,
         system_instruction,
         sysvar::{rent::Rent, Sysvar},
@@ -33,7 +34,7 @@ pub fn process_reallocate(
     let authority_info_data_len = authority_info.data_len();
 
     // check that account is the right type and validate owner
-    let mut current_extension_types = {
+    let (mut current_extension_types, native_token_amount) = {
         let token_account = token_account_info.data.borrow();
         let account = StateWithExtensions::<Account>::unpack(&token_account)?;
         Processor::validate_owner(
@@ -43,7 +44,8 @@ pub fn process_reallocate(
             authority_info_data_len,
             account_info_iter.as_slice(),
         )?;
-        account.get_extension_types()?
+        let native_token_amount = account.base.is_native().then(|| account.base.amount);
+        (account.get_extension_types()?, native_token_amount)
     };
 
     // check that all desired extensions are for the right account type
@@ -71,20 +73,42 @@ pub fn process_reallocate(
 
     // if additional lamports needed to remain rent-exempt, transfer them
     let rent = Rent::get()?;
-    let new_minimum_balance = rent.minimum_balance(needed_account_len);
-    let lamports_diff = new_minimum_balance.saturating_sub(token_account_info.lamports());
-    invoke(
-        &system_instruction::transfer(payer_info.key, token_account_info.key, lamports_diff),
-        &[
-            payer_info.clone(),
-            token_account_info.clone(),
-            system_program_info.clone(),
-        ],
-    )?;
+    let new_rent_exempt_reserve = rent.minimum_balance(needed_account_len);
 
-    // unpack to set account_type, if needed
-    let mut token_account = token_account_info.data.borrow_mut();
-    set_account_type::<Account>(&mut token_account)?;
+    let current_lamport_reserve = token_account_info
+        .lamports()
+        .checked_sub(native_token_amount.unwrap_or(0))
+        .ok_or(TokenError::Overflow)?;
+    let lamports_diff = new_rent_exempt_reserve.saturating_sub(current_lamport_reserve);
+    if lamports_diff > 0 {
+        invoke(
+            &system_instruction::transfer(payer_info.key, token_account_info.key, lamports_diff),
+            &[
+                payer_info.clone(),
+                token_account_info.clone(),
+                system_program_info.clone(),
+            ],
+        )?;
+    }
+
+    // set account_type, if needed
+    let mut token_account_data = token_account_info.data.borrow_mut();
+    set_account_type::<Account>(&mut token_account_data)?;
+
+    // sync the rent exempt reserve for native accounts
+    if let Some(native_token_amount) = native_token_amount {
+        let mut token_account = StateWithExtensionsMut::<Account>::unpack(&mut token_account_data)?;
+        // sanity check that there are enough lamports to cover the token amount
+        // and the rent exempt reserve
+        let minimum_lamports = new_rent_exempt_reserve
+            .checked_add(native_token_amount)
+            .ok_or(TokenError::Overflow)?;
+        if token_account_info.lamports() < minimum_lamports {
+            return Err(TokenError::InvalidState.into());
+        }
+        token_account.base.is_native = COption::Some(new_rent_exempt_reserve);
+        token_account.pack_base();
+    }
 
     Ok(())
 }