FlashLoan2: API with Begin and End instructions

For FlashLoan users had to pass the target cpi programs, accounts and
data to the FlashLoan instruction itself.

The new API allows existing instructions to be used unchanged, they
just need to be bracketed by FlashLoan2Begin and FlashLoan2End.

programs/mango-v4/src/instructions/flash_loan2.rs

@@ -0,0 +1,288 @@
+use crate::accounts_zerocopy::*;
+use crate::error::MangoError;
+use crate::group_seeds;
+use crate::state::{compute_health_from_fixed_accounts, Bank, Group, HealthType, MangoAccount};
+use crate::util::checked_math as cm;
+use anchor_lang::prelude::*;
+use anchor_lang::solana_program::sysvar::instructions as tx_instructions;
+use anchor_spl::token::{self, Token, TokenAccount};
+use fixed::types::I80F48;
+
+/// Sets up mango vaults for flash loan
+///
+/// In addition to these accounts, there must be a sequence of remaining_accounts:
+/// 1. N banks
+/// 2. N vaults, matching the banks
+#[derive(Accounts)]
+pub struct FlashLoan2Begin<'info> {
+    pub group: AccountLoader<'info, Group>,
+    pub temporary_vault_authority: Signer<'info>,
+    pub token_program: Program<'info, Token>,
+
+    #[account(address = tx_instructions::ID)]
+    pub instructions: UncheckedAccount<'info>,
+}
+
+/// Finalizes a flash loan
+///
+/// In addition to these accounts, there must be a sequence of remaining_accounts:
+/// 1. health accounts
+/// 2. N vaults, matching what was in FlashLoan2Begin
+#[derive(Accounts)]
+pub struct FlashLoan2End<'info> {
+    pub group: AccountLoader<'info, Group>,
+    #[account(
+        mut,
+        has_one = owner,
+        has_one = group,
+    )]
+    pub account: AccountLoader<'info, MangoAccount>,
+    pub owner: Signer<'info>,
+    pub token_program: Program<'info, Token>,
+}
+
+pub fn flash_loan2_begin<'key, 'accounts, 'remaining, 'info>(
+    ctx: Context<'key, 'accounts, 'remaining, 'info, FlashLoan2Begin<'info>>,
+    loan_amounts: Vec<u64>,
+) -> Result<()> {
+    let num_loans = loan_amounts.len();
+    require_eq!(
+        ctx.remaining_accounts.len(),
+        2 * num_loans,
+        MangoError::SomeError
+    );
+    let banks = &ctx.remaining_accounts[..num_loans];
+    let vaults = &ctx.remaining_accounts[num_loans..];
+
+    let group = ctx.accounts.group.load()?;
+    let group_seeds = group_seeds!(group);
+    let seeds = [&group_seeds[..]];
+
+    // Check that the banks and vaults correspond
+    for ((bank_ai, vault_ai), amount) in banks.iter().zip(vaults.iter()).zip(loan_amounts.iter()) {
+        let mut bank = bank_ai.load_mut::<Bank>()?;
+        require_keys_eq!(bank.group, ctx.accounts.group.key());
+        require_keys_eq!(bank.vault, *vault_ai.key);
+
+        let token_account = Account::<TokenAccount>::try_from(vault_ai)?;
+
+        bank.flash_loan_approved_amount = *amount;
+        bank.flash_loan_vault_initial = token_account.amount;
+
+        // Approve the withdraw
+        if *amount > 0 {
+            let approve_ctx = CpiContext::new(
+                ctx.accounts.token_program.to_account_info(),
+                token::Approve {
+                    to: vault_ai.clone(),
+                    delegate: ctx.accounts.temporary_vault_authority.to_account_info(),
+                    authority: ctx.accounts.group.to_account_info(),
+                },
+            )
+            .with_signer(&seeds);
+            token::approve(approve_ctx, *amount)?;
+        }
+    }
+
+    // Check if the other instructions in the transactions are compatible
+    {
+        let ixs = ctx.accounts.instructions.as_ref();
+        let current_index = tx_instructions::load_current_index_checked(ixs)? as usize;
+
+        // Forbid FlashLoan2Begin to be called from CPI (it does not have to be the first instruction)
+        let current_ix = tx_instructions::load_instruction_at_checked(current_index, ixs)?;
+        require_keys_eq!(
+            current_ix.program_id,
+            *ctx.program_id,
+            MangoError::SomeError
+        );
+
+        // The only other mango instruction that must appear before the end of the tx is
+        // the FlashLoan2End instruction. No other mango instructions are allowed.
+        let mut index = current_index + 1;
+        let mut found_end = false;
+        loop {
+            let ix = match tx_instructions::load_instruction_at_checked(index, ixs) {
+                Ok(ix) => ix,
+                Err(ProgramError::InvalidArgument) => break, // past the last instruction
+                Err(e) => Err(e)?,
+            };
+
+            // Check that the mango program key is not used
+            if ix.program_id == crate::id() {
+                // must be the last mango ix -- this could possibly be relaxed, but right now
+                // we need to guard against multiple FlashLoanEnds
+                require!(!found_end, MangoError::SomeError);
+                found_end = true;
+
+                // must be the FlashLoan2End instruction
+                require!(
+                    &ix.data[0..8] == &[187, 107, 239, 212, 18, 21, 145, 171],
+                    MangoError::SomeError
+                );
+
+                // check that the same vaults are passed
+                let end_vaults = &ix.accounts[ix.accounts.len() - num_loans..];
+                for (start_vault, end_vault) in vaults.iter().zip(end_vaults.iter()) {
+                    require_keys_eq!(*start_vault.key, end_vault.pubkey);
+                }
+            } else {
+                // ensure no one can cpi into mango either
+                for meta in ix.accounts.iter() {
+                    require_keys_neq!(meta.pubkey, crate::id());
+                }
+            }
+
+            index += 1;
+        }
+        require!(found_end, MangoError::SomeError);
+    }
+
+    Ok(())
+}
+
+struct TokenVaultChange {
+    bank_index: usize,
+    raw_token_index: usize,
+    amount: I80F48,
+}
+
+// Remaining accounts:
+// 1. health
+// 2. vaults (must be same as FlashLoanStart)
+pub fn flash_loan2_end<'key, 'accounts, 'remaining, 'info>(
+    ctx: Context<'key, 'accounts, 'remaining, 'info, FlashLoan2End<'info>>,
+) -> Result<()> {
+    let group = ctx.accounts.group.load()?;
+    let group_seeds = group_seeds!(group);
+
+    let mut account = ctx.accounts.account.load_mut()?;
+    require!(account.is_bankrupt == 0, MangoError::IsBankrupt);
+
+    // Find index at which vaults start
+    let vaults_index = ctx
+        .remaining_accounts
+        .iter()
+        .position(|ai| {
+            let maybe_token_account = Account::<TokenAccount>::try_from(ai);
+            if maybe_token_account.is_err() {
+                return false;
+            }
+
+            maybe_token_account.unwrap().owner == account.group
+        })
+        .ok_or_else(|| error!(MangoError::SomeError))?;
+
+    // First initialize to the remaining delegated amount
+    let health_ais = &ctx.remaining_accounts[..vaults_index];
+    let vaults = &ctx.remaining_accounts[vaults_index..];
+    let mut vaults_with_banks = vec![false; vaults.len()];
+
+    // Loop over the banks, finding matching vaults
+    // TODO: must be moved into health.rs, because it assumes something about the health accounts structure
+    let mut changes = vec![];
+    for (i, bank_ai) in health_ais.iter().enumerate() {
+        // iterate until the first non-bank
+        let bank = match bank_ai.load::<Bank>() {
+            Ok(b) => b,
+            Err(_) => break,
+        };
+
+        // find a vault -- if there's none, skip
+        let (vault_index, vault_ai) = match vaults
+            .iter()
+            .enumerate()
+            .find(|(_, vault_ai)| vault_ai.key == &bank.vault)
+        {
+            Some(v) => v,
+            None => continue,
+        };
+
+        let vault = Account::<TokenAccount>::try_from(vault_ai)?;
+        vaults_with_banks[vault_index] = true;
+
+        // Ensure this bank/vault combination was mentioned in the Begin instruction:
+        // The Begin instruction only checks that End ends with the same vault accounts -
+        // but there could be an extra vault account in End, or a different bank could be
+        // used for the same vault.
+        require_neq!(bank.flash_loan_vault_initial, u64::MAX);
+
+        // Create the token position now, so we can compute the pre-health with fixed order health accounts
+        let (_, raw_token_index) = account.tokens.get_mut_or_create(bank.token_index)?;
+
+        // Revoke delegation
+        let ix = token::spl_token::instruction::revoke(
+            &token::spl_token::ID,
+            vault_ai.key,
+            &ctx.accounts.group.key(),
+            &[],
+        )?;
+        solana_program::program::invoke_signed(
+            &ix,
+            &[vault_ai.clone(), ctx.accounts.group.to_account_info()],
+            &[group_seeds],
+        )?;
+
+        // Track vault difference
+        let new_amount = I80F48::from(vault.amount);
+        let old_amount = I80F48::from(bank.flash_loan_vault_initial);
+        let change = cm!(new_amount - old_amount);
+
+        changes.push(TokenVaultChange {
+            bank_index: i,
+            raw_token_index,
+            amount: change,
+        });
+    }
+
+    // all vaults must have had matching banks
+    require!(vaults_with_banks.iter().all(|&b| b), MangoError::SomeError);
+
+    // Check pre-cpi health
+    // NOTE: This health check isn't strictly necessary. It will be, later, when
+    // we want to have reduce_only or be able to move an account out of bankruptcy.
+    let pre_cpi_health =
+        compute_health_from_fixed_accounts(&account, HealthType::Init, health_ais)?;
+    require!(pre_cpi_health >= 0, MangoError::HealthMustBePositive);
+    msg!("pre_cpi_health {:?}", pre_cpi_health);
+
+    // Apply the vault diffs to the bank positions
+    let mut deactivated_token_positions = vec![];
+    for change in changes {
+        let mut bank = health_ais[change.bank_index].load_mut::<Bank>()?;
+        let position = account.tokens.get_mut_raw(change.raw_token_index);
+        let native = position.native(&bank);
+        let approved_amount = I80F48::from(bank.flash_loan_approved_amount);
+
+        let loan = if native.is_positive() {
+            cm!(approved_amount - native).max(I80F48::ZERO)
+        } else {
+            approved_amount
+        };
+
+        let loan_origination_fee = cm!(loan * bank.loan_origination_fee_rate);
+        bank.collected_fees_native = cm!(bank.collected_fees_native + loan_origination_fee);
+
+        let is_active =
+            bank.change_without_fee(position, cm!(change.amount - loan_origination_fee))?;
+        if !is_active {
+            deactivated_token_positions.push(change.raw_token_index);
+        }
+
+        bank.flash_loan_approved_amount = 0;
+        bank.flash_loan_vault_initial = u64::MAX;
+    }
+
+    // Check post-cpi health
+    let post_cpi_health =
+        compute_health_from_fixed_accounts(&account, HealthType::Init, health_ais)?;
+    require!(post_cpi_health >= 0, MangoError::HealthMustBePositive);
+    msg!("post_cpi_health {:?}", post_cpi_health);
+
+    // Deactivate inactive token accounts after health check
+    for raw_token_index in deactivated_token_positions {
+        account.tokens.deactivate(raw_token_index);
+    }
+
+    Ok(())
+}

programs/mango-v4/src/instructions/mod.rs

@@ -7,6 +7,7 @@ pub use create_account::*;
 pub use create_group::*;
 pub use create_stub_oracle::*;
 pub use flash_loan::*;
+pub use flash_loan2::*;
 pub use liq_token_with_token::*;
 pub use perp_cancel_all_orders::*;
 pub use perp_cancel_all_orders_by_side::*;
@@ -43,6 +44,7 @@ mod create_account;
 mod create_group;
 mod create_stub_oracle;
 mod flash_loan;
+mod flash_loan2;
 mod liq_token_with_token;
 mod perp_cancel_all_orders;
 mod perp_cancel_all_orders_by_side;

programs/mango-v4/src/instructions/token_register.rs

@@ -136,6 +136,8 @@ pub fn token_register(
         init_liab_weight: I80F48::from_num(init_liab_weight),
         liquidation_fee: I80F48::from_num(liquidation_fee),
         dust: I80F48::ZERO,
+        flash_loan_vault_initial: u64::MAX,
+        flash_loan_approved_amount: 0,
         token_index,
         bump: *ctx.bumps.get("bank").ok_or(MangoError::SomeError)?,
         mint_decimals: ctx.accounts.mint.decimals,

programs/mango-v4/src/lib.rs

@@ -142,6 +142,19 @@ pub mod mango_v4 {
         instructions::flash_loan(ctx, withdraws, cpi_datas)
     }
 
+    pub fn flash_loan2_begin<'key, 'accounts, 'remaining, 'info>(
+        ctx: Context<'key, 'accounts, 'remaining, 'info, FlashLoan2Begin<'info>>,
+        loan_amounts: Vec<u64>,
+    ) -> Result<()> {
+        instructions::flash_loan2_begin(ctx, loan_amounts)
+    }
+
+    pub fn flash_loan2_end<'key, 'accounts, 'remaining, 'info>(
+        ctx: Context<'key, 'accounts, 'remaining, 'info, FlashLoan2End<'info>>,
+    ) -> Result<()> {
+        instructions::flash_loan2_end(ctx)
+    }
+
     ///
     /// Serum
     ///

programs/mango-v4/src/state/bank.rs

@@ -63,6 +63,9 @@ pub struct Bank {
     // Collection of all fractions-of-native-tokens that got rounded away
     pub dust: I80F48,
 
+    pub flash_loan_vault_initial: u64,
+    pub flash_loan_approved_amount: u64,
+
     // Index into TokenInfo on the group
     pub token_index: TokenIndex,
 
@@ -78,7 +81,7 @@ pub struct Bank {
 }
 const_assert_eq!(
     size_of::<Bank>(),
-    16 + 32 * 4 + 8 + 16 * 21 + 2 + 1 + 1 + 4 + 8
+    16 + 32 * 4 + 8 + 16 * 21 + 2 * 8 + 2 + 1 + 1 + 4 + 8
 );
 const_assert_eq!(size_of::<Bank>() % 8, 0);
 
@@ -113,6 +116,11 @@ impl std::fmt::Debug for Bank {
             .field("liquidation_fee", &self.liquidation_fee)
             .field("dust", &self.dust)
             .field("token_index", &self.token_index)
+            .field(
+                "flash_loan_approved_amount",
+                &self.flash_loan_approved_amount,
+            )
+            .field("flash_loan_vault_initial", &self.flash_loan_vault_initial)
             .field("reserved", &self.reserved)
             .finish()
     }
@@ -148,6 +156,8 @@ impl Bank {
             init_liab_weight: existing_bank.init_liab_weight,
             liquidation_fee: existing_bank.liquidation_fee,
             dust: I80F48::ZERO,
+            flash_loan_approved_amount: 0,
+            flash_loan_vault_initial: u64::MAX,
             token_index: existing_bank.token_index,
             bump: existing_bank.bump,
             mint_decimals: existing_bank.mint_decimals,

programs/mango-v4/src/state/health.rs

@@ -50,6 +50,14 @@ pub struct FixedOrderAccountRetriever<T: KeyedAccountReader> {
     pub begin_serum3: usize,
 }
 
+impl<T: KeyedAccountReader> FixedOrderAccountRetriever<T> {
+    fn bank(&self, group: &Pubkey, account_index: usize) -> Result<&Bank> {
+        let bank = self.ais[account_index].load::<Bank>()?;
+        require!(&bank.group == group, MangoError::SomeError);
+        Ok(bank)
+    }
+}
+
 impl<T: KeyedAccountReader> AccountRetriever for FixedOrderAccountRetriever<T> {
     fn bank_and_oracle(
         &self,
@@ -57,8 +65,7 @@ impl<T: KeyedAccountReader> AccountRetriever for FixedOrderAccountRetriever<T> {
         account_index: usize,
         token_index: TokenIndex,
     ) -> Result<(&Bank, I80F48)> {
-        let bank = self.ais[account_index].load::<Bank>()?;
-        require!(&bank.group == group, MangoError::SomeError);
+        let bank = self.bank(group, account_index)?;
         require!(bank.token_index == token_index, MangoError::SomeError);
         let oracle = &self.ais[cm!(self.n_banks + account_index)];
         require!(&bank.oracle == oracle.key(), MangoError::SomeError);

programs/mango-v4/tests/program_test/mango_client.rs

@@ -15,8 +15,10 @@ use solana_sdk::signature::{Keypair, Signer};
 use solana_sdk::transport::TransportError;
 
 use std::str::FromStr;
+use std::sync::Arc;
 
 use super::solana::SolanaCookie;
+use super::utils::clone_keypair;
 use mango_v4::state::*;
 
 #[async_trait::async_trait(?Send)]
@@ -49,6 +51,47 @@ pub async fn send_tx<CI: ClientInstruction>(
     Ok(accounts)
 }
 
+/// Build a transaction from multiple instructions
+pub struct ClientTransaction {
+    solana: Arc<SolanaCookie>,
+    instructions: Vec<instruction::Instruction>,
+    signers: Vec<Keypair>,
+}
+
+impl<'a> ClientTransaction {
+    pub fn new(solana: &Arc<SolanaCookie>) -> Self {
+        Self {
+            solana: solana.clone(),
+            instructions: vec![],
+            signers: vec![],
+        }
+    }
+
+    pub async fn add_instruction<CI: ClientInstruction>(&mut self, ix: CI) -> CI::Accounts {
+        let solana: &SolanaCookie = &self.solana;
+        let (accounts, instruction) = ix.to_instruction(solana).await;
+        self.instructions.push(instruction);
+        self.signers
+            .extend(ix.signers().iter().map(|k| clone_keypair(k)));
+        accounts
+    }
+
+    pub fn add_instruction_direct(&mut self, ix: instruction::Instruction) {
+        self.instructions.push(ix);
+    }
+
+    pub fn add_signer(&mut self, keypair: &Keypair) {
+        self.signers.push(clone_keypair(keypair));
+    }
+
+    pub async fn send(&self) -> std::result::Result<(), TransportError> {
+        let signer_refs = self.signers.iter().map(|k| k).collect::<Vec<&Keypair>>();
+        self.solana
+            .process_transaction(&self.instructions, Some(&signer_refs[..]))
+            .await
+    }
+}
+
 #[async_trait::async_trait(?Send)]
 pub trait ClientInstruction {
     type Accounts: anchor_lang::ToAccountMetas;
@@ -371,6 +414,105 @@ impl<'keypair> ClientInstruction for FlashLoanInstruction<'keypair> {
     }
 }
 
+pub struct FlashLoan2BeginInstruction<'keypair> {
+    pub group: Pubkey,
+    pub mango_token_bank: Pubkey,
+    pub mango_token_vault: Pubkey,
+    pub withdraw_amount: u64,
+    pub temporary_vault_authority: &'keypair Keypair,
+}
+#[async_trait::async_trait(?Send)]
+impl<'keypair> ClientInstruction for FlashLoan2BeginInstruction<'keypair> {
+    type Accounts = mango_v4::accounts::FlashLoan2Begin;
+    type Instruction = mango_v4::instruction::FlashLoan2Begin;
+    async fn to_instruction(
+        &self,
+        _account_loader: impl ClientAccountLoader + 'async_trait,
+    ) -> (Self::Accounts, instruction::Instruction) {
+        let program_id = mango_v4::id();
+
+        let accounts = Self::Accounts {
+            group: self.group,
+            temporary_vault_authority: self.temporary_vault_authority.pubkey(),
+            token_program: Token::id(),
+            instructions: solana_program::sysvar::instructions::id(),
+        };
+
+        let instruction = Self::Instruction {
+            loan_amounts: vec![self.withdraw_amount],
+        };
+
+        let mut instruction = make_instruction(program_id, &accounts, instruction);
+        instruction.accounts.push(AccountMeta {
+            pubkey: self.mango_token_bank,
+            is_writable: true,
+            is_signer: false,
+        });
+        instruction.accounts.push(AccountMeta {
+            pubkey: self.mango_token_vault,
+            is_writable: true,
+            is_signer: false,
+        });
+
+        (accounts, instruction)
+    }
+
+    fn signers(&self) -> Vec<&Keypair> {
+        vec![self.temporary_vault_authority]
+    }
+}
+
+pub struct FlashLoan2EndInstruction<'keypair> {
+    pub account: Pubkey,
+    pub owner: &'keypair Keypair,
+    pub mango_token_bank: Pubkey,
+    pub mango_token_vault: Pubkey,
+}
+#[async_trait::async_trait(?Send)]
+impl<'keypair> ClientInstruction for FlashLoan2EndInstruction<'keypair> {
+    type Accounts = mango_v4::accounts::FlashLoan2End;
+    type Instruction = mango_v4::instruction::FlashLoan2End;
+    async fn to_instruction(
+        &self,
+        account_loader: impl ClientAccountLoader + 'async_trait,
+    ) -> (Self::Accounts, instruction::Instruction) {
+        let program_id = mango_v4::id();
+        let instruction = Self::Instruction {};
+
+        let account: MangoAccount = account_loader.load(&self.account).await.unwrap();
+
+        let health_check_metas = derive_health_check_remaining_account_metas(
+            &account_loader,
+            &account,
+            Some(self.mango_token_bank),
+            true,
+            None,
+        )
+        .await;
+
+        let accounts = Self::Accounts {
+            group: account.group,
+            account: self.account,
+            owner: self.owner.pubkey(),
+            token_program: Token::id(),
+        };
+
+        let mut instruction = make_instruction(program_id, &accounts, instruction);
+        instruction.accounts.extend(health_check_metas.into_iter());
+        instruction.accounts.push(AccountMeta {
+            pubkey: self.mango_token_vault,
+            is_writable: true,
+            is_signer: false,
+        });
+
+        (accounts, instruction)
+    }
+
+    fn signers(&self) -> Vec<&Keypair> {
+        vec![self.owner]
+    }
+}
+
 pub struct TokenWithdrawInstruction<'keypair> {
     pub amount: u64,
     pub allow_borrow: bool,

programs/mango-v4/tests/test_margin_trade.rs

@@ -13,7 +13,7 @@ mod program_test;
 // This is an unspecific happy-case test that just runs a few instructions to check
 // that they work in principle. It should be split up / renamed.
 #[tokio::test]
-async fn test_margin_trade() -> Result<(), BanksClientError> {
+async fn test_margin_trade1() -> Result<(), BanksClientError> {
     let mut builder = TestContextBuilder::new();
     let margin_trade = builder.add_margin_trade_program();
     let context = builder.start_default().await;
@@ -323,3 +323,261 @@ async fn test_margin_trade() -> Result<(), BanksClientError> {
 
     Ok(())
 }
+
+// This is an unspecific happy-case test that just runs a few instructions to check
+// that they work in principle. It should be split up / renamed.
+#[tokio::test]
+async fn test_margin_trade2() -> Result<(), BanksClientError> {
+    let builder = TestContextBuilder::new();
+    let context = builder.start_default().await;
+    let solana = &context.solana.clone();
+
+    let admin = &Keypair::new();
+    let owner = &context.users[0].key;
+    let payer = &context.users[1].key;
+    let mints = &context.mints[0..2];
+    let payer_mint0_account = context.users[1].token_accounts[0];
+    let payer_mint1_account = context.users[1].token_accounts[1];
+    let loan_origination_fee = 0.0005;
+
+    // higher resolution that the loan_origination_fee for one token
+    let balance_f64eq = |a: f64, b: f64| (a - b).abs() < 0.0001;
+
+    //
+    // SETUP: Create a group, account, register a token (mint0)
+    //
+
+    let mango_setup::GroupWithTokens { group, tokens } = mango_setup::GroupWithTokensConfig {
+        admin,
+        payer,
+        mints,
+    }
+    .create(solana)
+    .await;
+    let bank = tokens[0].bank;
+    let vault = tokens[0].vault;
+
+    //
+    // provide some funds for tokens, so the test user can borrow
+    //
+    let provided_amount = 1000;
+
+    let provider_account = send_tx(
+        solana,
+        CreateAccountInstruction {
+            account_num: 1,
+            group,
+            owner,
+            payer,
+        },
+    )
+    .await
+    .unwrap()
+    .account;
+
+    send_tx(
+        solana,
+        TokenDepositInstruction {
+            amount: provided_amount,
+            account: provider_account,
+            token_account: payer_mint0_account,
+            token_authority: payer,
+        },
+    )
+    .await
+    .unwrap();
+    send_tx(
+        solana,
+        TokenDepositInstruction {
+            amount: provided_amount,
+            account: provider_account,
+            token_account: payer_mint1_account,
+            token_authority: payer,
+        },
+    )
+    .await
+    .unwrap();
+
+    //
+    // create thes test user account
+    //
+
+    let account = send_tx(
+        solana,
+        CreateAccountInstruction {
+            account_num: 0,
+            group,
+            owner,
+            payer,
+        },
+    )
+    .await
+    .unwrap()
+    .account;
+
+    //
+    // TEST: Deposit funds
+    //
+    let deposit_amount_initial = 100;
+    {
+        let start_balance = solana.token_account_balance(payer_mint0_account).await;
+
+        send_tx(
+            solana,
+            TokenDepositInstruction {
+                amount: deposit_amount_initial,
+                account,
+                token_account: payer_mint0_account,
+                token_authority: payer,
+            },
+        )
+        .await
+        .unwrap();
+
+        assert_eq!(
+            solana.token_account_balance(vault).await,
+            provided_amount + deposit_amount_initial
+        );
+        assert_eq!(
+            solana.token_account_balance(payer_mint0_account).await,
+            start_balance - deposit_amount_initial
+        );
+        assert_eq!(
+            account_position(solana, account, bank).await,
+            deposit_amount_initial as i64,
+        );
+    }
+
+    //
+    // TEST: Margin trade
+    //
+    let margin_account = payer_mint0_account;
+    let margin_account_initial = solana.token_account_balance(margin_account).await;
+    let withdraw_amount = 2;
+    let deposit_amount = 1;
+    let send_flash_loan_tx = |solana, withdraw_amount, deposit_amount| async move {
+        let temporary_vault_authority = &Keypair::new();
+
+        let mut tx = ClientTransaction::new(solana);
+        tx.add_instruction(FlashLoan2BeginInstruction {
+            group,
+            temporary_vault_authority,
+            mango_token_bank: bank,
+            mango_token_vault: vault,
+            withdraw_amount,
+        })
+        .await;
+        if withdraw_amount > 0 {
+            tx.add_instruction_direct(
+                spl_token::instruction::transfer(
+                    &spl_token::ID,
+                    &vault,
+                    &margin_account,
+                    &temporary_vault_authority.pubkey(),
+                    &[&temporary_vault_authority.pubkey()],
+                    withdraw_amount,
+                )
+                .unwrap(),
+            );
+        }
+        if deposit_amount > 0 {
+            tx.add_instruction_direct(
+                spl_token::instruction::transfer(
+                    &spl_token::ID,
+                    &margin_account,
+                    &vault,
+                    &payer.pubkey(),
+                    &[&payer.pubkey()],
+                    deposit_amount,
+                )
+                .unwrap(),
+            );
+            tx.add_signer(&payer);
+        }
+        tx.add_instruction(FlashLoan2EndInstruction {
+            account,
+            owner,
+            mango_token_bank: bank,
+            mango_token_vault: vault,
+        })
+        .await;
+        tx.send().await.unwrap();
+    };
+    send_flash_loan_tx(solana, withdraw_amount, deposit_amount).await;
+
+    assert_eq!(
+        solana.token_account_balance(vault).await,
+        provided_amount + deposit_amount_initial - withdraw_amount + deposit_amount
+    );
+    assert_eq!(
+        solana.token_account_balance(margin_account).await,
+        margin_account_initial + withdraw_amount - deposit_amount
+    );
+    // no fee because user had positive balance
+    assert!(balance_f64eq(
+        account_position_f64(solana, account, bank).await,
+        (deposit_amount_initial - withdraw_amount + deposit_amount) as f64
+    ));
+
+    //
+    // TEST: Bringing the balance to 0 deactivates the token
+    //
+    let deposit_amount_initial = account_position(solana, account, bank).await;
+    let margin_account_initial = solana.token_account_balance(margin_account).await;
+    let withdraw_amount = deposit_amount_initial as u64;
+    let deposit_amount = 0;
+    send_flash_loan_tx(solana, withdraw_amount, deposit_amount).await;
+    assert_eq!(solana.token_account_balance(vault).await, provided_amount);
+    assert_eq!(
+        solana.token_account_balance(margin_account).await,
+        margin_account_initial + withdraw_amount
+    );
+    // Check that position is fully deactivated
+    let account_data: MangoAccount = solana.get_account(account).await;
+    assert_eq!(account_data.tokens.iter_active().count(), 0);
+
+    //
+    // TEST: Activating a token via margin trade
+    //
+    let margin_account_initial = solana.token_account_balance(margin_account).await;
+    let withdraw_amount = 0;
+    let deposit_amount = 100;
+    send_flash_loan_tx(solana, withdraw_amount, deposit_amount).await;
+    assert_eq!(
+        solana.token_account_balance(vault).await,
+        provided_amount + deposit_amount
+    );
+    assert_eq!(
+        solana.token_account_balance(margin_account).await,
+        margin_account_initial - deposit_amount
+    );
+    assert!(balance_f64eq(
+        account_position_f64(solana, account, bank).await,
+        deposit_amount as f64
+    ));
+
+    //
+    // TEST: Try loan fees by withdrawing more than the user balance
+    //
+    let margin_account_initial = solana.token_account_balance(margin_account).await;
+    let deposit_amount_initial = account_position(solana, account, bank).await as u64;
+    let withdraw_amount = 500;
+    let deposit_amount = 450;
+    println!("{}", deposit_amount_initial);
+    send_flash_loan_tx(solana, withdraw_amount, deposit_amount).await;
+    assert_eq!(
+        solana.token_account_balance(vault).await,
+        provided_amount + deposit_amount_initial + deposit_amount - withdraw_amount
+    );
+    assert_eq!(
+        solana.token_account_balance(margin_account).await,
+        margin_account_initial + withdraw_amount - deposit_amount
+    );
+    assert!(balance_f64eq(
+        account_position_f64(solana, account, bank).await,
+        (deposit_amount_initial + deposit_amount - withdraw_amount) as f64
+            - (withdraw_amount - deposit_amount_initial) as f64 * loan_origination_fee
+    ));
+
+    Ok(())
+}