Merge pull request #55 from blockworks-foundation/ckamm/margin-trade-updates

Improvements to margin_trade

.cargo/audit.toml

@@ -27,9 +27,6 @@ show_tree = true # Show inverse dependency trees along with advisories (default:
 # arch = "x86_64" # Ignore advisories for CPU architectures other than this one
 # os = "linux" # Ignore advisories for operating systems other than this one
 
-[packages]
-source = "all" # "all", "public" or "local"
-
 [yanked]
 enabled = false # Warn for yanked crates in Cargo.lock (default: true)
 update_index = true # Auto-update the crates.io index (default: true)

.github/workflows/ci-lint-test.yml

@@ -9,8 +9,8 @@ on:
 
 env:
   CARGO_TERM_COLOR: always
-  SOLANA_VERSION: "1.9.5"
-  RUST_TOOLCHAIN: stable
+  SOLANA_VERSION: "1.9.14"
+  RUST_TOOLCHAIN: 1.60.0
   LOG_PROGRAM: "m43thNJ58XCjL798ZSq6JGAG1BnWskhdq5or6kcnfsD"
 
 defaults:
@@ -27,7 +27,7 @@ jobs:
         uses: actions/checkout@v2
       - name: Install Linux dependencies
         run: sudo apt-get update && sudo apt-get install -y pkg-config build-essential libudev-dev
-      - name: Install Rust nightly
+      - name: Install Rust
         uses: actions-rs/toolchain@v1
         with:
           override: true
@@ -50,19 +50,12 @@ jobs:
         uses: actions/checkout@v2
       - name: Install Linux dependencies
         run: sudo apt-get update && sudo apt-get install -y pkg-config build-essential libudev-dev
-      - name: Install Rust nightly
+      - name: Install Rust
         uses: actions-rs/toolchain@v1
         with:
           override: true
           profile: minimal
           toolchain: ${{ env.RUST_TOOLCHAIN }}
-      - name: Cache dependencies
-        uses: Swatinem/rust-cache@v1
-      - name: Cache Solana binaries
-        uses: actions/cache@v2
-        with:
-          path: ~/.cache/solana
-          key: ${{ runner.os }}-${{ env.SOLANA_VERSION }}
       - name: Install Solana
         run: |
           sh -c "$(curl -sSfL https://release.solana.com/v${{ env.SOLANA_VERSION }}/install)"
@@ -82,7 +75,7 @@ jobs:
           name: raw-test-bpf
           path: raw-test-bpf.log
 
-# Download logs and process them
+  # Download logs and process them
   process-logs:
     name: Process logs
     if: github.actor != 'github-actions[bot]'

.github/workflows/ci-soteria.yml

@@ -10,7 +10,7 @@ on:
 
 env:
   CARGO_TERM_COLOR: always
-  SOLANA_VERSION: "1.9.5"
+  SOLANA_VERSION: "1.9.14"
 
 jobs:
   build:
@@ -20,7 +20,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        program: [ 'programs/mango-v4', 'programs/margin-trade' ]
+        program: ["programs/mango-v4", "programs/margin-trade"]
     env:
       PROGRAM_PATH: ${{ matrix.program }}
 

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

@@ -1,10 +1,25 @@
 use crate::error::MangoError;
 use crate::state::{compute_health_from_fixed_accounts, Bank, Group, HealthType, MangoAccount};
+use crate::util::LoadZeroCopy;
 use crate::{group_seeds, Mango};
 use anchor_lang::prelude::*;
-use anchor_spl::token::TokenAccount;
+use anchor_spl::token::{self, Token, TokenAccount};
 use fixed::types::I80F48;
 use solana_program::instruction::Instruction;
+use std::cell::Ref;
+use std::collections::HashMap;
+
+/// The margin trade instruction
+///
+/// In addition to these accounts, there must be a sequence of remaining_accounts:
+/// 1. health_accounts: accounts needed for health checking
+/// 2. target_program_id: the target program account
+/// 3. target_accounts: the accounts to pass to the target program
+///
+/// Every vault address listed in 3. must also have the matching bank and oracle appear in 1.
+///
+/// Every vault that is to be withdrawn from must appear in the `withdraws` instruction argument.
+/// The corresponding bank may be used as an authority for vault withdrawals.
 #[derive(Accounts)]
 pub struct MarginTrade<'info> {
     pub group: AccountLoader<'info, Group>,
@@ -17,167 +32,266 @@ pub struct MarginTrade<'info> {
     pub account: AccountLoader<'info, MangoAccount>,
 
     pub owner: Signer<'info>,
+    pub token_program: Program<'info, Token>,
 }
 
-// TODO: add loan fees
+struct AllowedVault {
+    // index of the vault in cpi_ais
+    vault_cpi_ai_index: usize,
+    // index of the bank in health_ais
+    bank_health_ai_index: usize,
+    // raw index into account.tokens
+    raw_token_index: usize,
+    // vault amount before cpi
+    pre_amount: u64,
+    // withdraw request
+    withdraw_amount: u64,
+    // amount of withdraw request that is a loan
+    loan_amount: I80F48,
+}
+
+/// - `num_health_accounts` is the number of health accounts that remaining_accounts starts with.
+/// - `withdraws` is a list of tuples containing the index to a vault in target_accounts and the
+///    amount that the target program shall be allowed to withdraw
+/// - `cpi_data` is the bytes to call the target_program_id with
 pub fn margin_trade<'key, 'accounts, 'remaining, 'info>(
     ctx: Context<'key, 'accounts, 'remaining, 'info, MarginTrade<'info>>,
-    banks_len: usize,
+    num_health_accounts: usize,
+    withdraws: Vec<(u8, u64)>,
     cpi_data: Vec<u8>,
 ) -> Result<()> {
     let group = ctx.accounts.group.load()?;
     let mut account = ctx.accounts.account.load_mut()?;
     require!(account.is_bankrupt == 0, MangoError::IsBankrupt);
 
-    // remaining_accounts layout is expected as follows
-    // * banks_len number of banks
-    // * banks_len number of oracles
-    // * cpi_program
-    // * cpi_accounts
-
-    // assert that user has passed in enough banks, this might be greater than his current
-    // total number of indexed positions, since
-    // user might end up withdrawing or depositing and activating a new indexed position
-    require!(
-        banks_len >= account.tokens.iter_active().count(),
-        MangoError::SomeError // todo: SomeError
-    );
-
-    // unpack remaining_accounts
-    let health_ais = &ctx.remaining_accounts[0..banks_len * 2];
-    // TODO: This relies on the particular shape of health_ais
-    let banks = &ctx.remaining_accounts[0..banks_len];
-    let cpi_program_id = *ctx.remaining_accounts[banks_len * 2].key;
-
-    // prepare account for cpi ix
-    let (cpi_ais, cpi_ams) = {
-        // we also need the group
-        let mut cpi_ais = [ctx.accounts.group.to_account_info()].to_vec();
-        // skip banks, oracles and cpi program from the remaining_accounts
-        let mut remaining_cpi_ais = ctx.remaining_accounts[banks_len * 2 + 1..].to_vec();
-        cpi_ais.append(&mut remaining_cpi_ais);
-
-        // todo: I'm wondering if there's a way to do this without putting cpi_ais on the heap.
-        // But fine to defer to the future
-        let mut cpi_ams = cpi_ais.to_account_metas(Option::None);
-        // we want group to be the signer, so that token vaults can be credited to or withdrawn from
-        cpi_ams[0].is_signer = true;
-
-        (cpi_ais, cpi_ams)
+    // Go over the banks passed as health accounts and:
+    // - Ensure that all banks that are passed in have activated positions.
+    //   This is necessary because maybe the user wants to margin trade on a token
+    //   that the account hasn't used before.
+    // - Collect the addresses of all banks to potentially sign for in cpi_ais.
+    // - Collect the addresses of all bank vaults.
+    // Note: This depends on the particular health account ordering.
+    let health_ais = &ctx.remaining_accounts[0..num_health_accounts];
+    let mut allowed_banks = HashMap::<&Pubkey, Ref<Bank>>::new();
+    // vault pubkey -> (bank_account_index, raw_token_index)
+    let mut allowed_vaults = HashMap::<Pubkey, (usize, usize)>::new();
+    for (i, ai) in health_ais.iter().enumerate() {
+        match ai.load::<Bank>() {
+            Ok(bank) => {
+                require!(bank.group == account.group, MangoError::SomeError);
+                let (_, raw_token_index) = account.tokens.get_mut_or_create(bank.token_index)?;
+                allowed_vaults.insert(bank.vault, (i, raw_token_index));
+                allowed_banks.insert(ai.key, bank);
+            }
+            Err(Error::AnchorError(error))
+                if error.error_code_number == ErrorCode::AccountDiscriminatorMismatch as u32 =>
+            {
+                break;
+            }
+            Err(error) => return Err(error),
         };
-
-    // sanity checks
-    for cpi_ai in &cpi_ais {
-        // since we are using group signer seeds to invoke cpi,
-        // assert that none of the cpi accounts is the mango program to prevent that invoker doesn't
-        // abuse this ix to do unwanted changes
-        require!(
-            cpi_ai.key() != Mango::id(),
-            MangoError::InvalidMarginTradeTargetCpiProgram
-        );
-
-        // assert that user has passed in the bank for every
-        // token account he wants to deposit/withdraw from in cpi
-        if cpi_ai.owner == &TokenAccount::owner() {
-            let maybe_mango_vault_token_account =
-                Account::<TokenAccount>::try_from(cpi_ai).unwrap();
-            if maybe_mango_vault_token_account.owner == ctx.accounts.group.key() {
-                require!(
-                    banks.iter().any(|bank_ai| {
-                        let bank_loader = AccountLoader::<'_, Bank>::try_from(bank_ai).unwrap();
-                        let bank = bank_loader.load().unwrap();
-                        bank.mint == maybe_mango_vault_token_account.mint
-                    }),
-                    // todo: errorcode
-                    MangoError::SomeError
-                )
-            }
-        }
     }
 
-    // compute pre cpi health
-    // TODO: check maint type?
+    // 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);
+    require!(pre_cpi_health >= 0, MangoError::HealthMustBePositive);
     msg!("pre_cpi_health {:?}", pre_cpi_health);
 
-    // prepare and invoke cpi
+    let cpi_program_id = *ctx.remaining_accounts[num_health_accounts].key;
+    let cpi_ais = &ctx.remaining_accounts[num_health_accounts + 1..];
+    let mut cpi_ams = cpi_ais
+        .iter()
+        .flat_map(|item| item.to_account_metas(None))
+        .collect::<Vec<_>>();
+    require!(cpi_ais.len() == cpi_ams.len(), MangoError::SomeError);
+
+    // Check that each group-owned token account is the vault of one of the allowed banks,
+    // and track its balance.
+    let mut used_vaults = cpi_ais
+        .iter()
+        .enumerate()
+        .filter_map(|(i, ai)| {
+            if ai.owner != &TokenAccount::owner() {
+                return None;
+            }
+            let token_account = Account::<TokenAccount>::try_from(ai).unwrap();
+            if token_account.owner != ctx.accounts.group.key() {
+                return None;
+            }
+
+            // Every group-owned token account must be a vault of one of the banks.
+            if let Some(&(bank_index, raw_token_index)) = allowed_vaults.get(&ai.key) {
+                return Some(Ok((
+                    ai.key,
+                    AllowedVault {
+                        vault_cpi_ai_index: i,
+                        bank_health_ai_index: bank_index,
+                        raw_token_index,
+                        pre_amount: token_account.amount,
+                        // these two are updated later
+                        withdraw_amount: 0,
+                        loan_amount: I80F48::ZERO,
+                    },
+                )));
+            }
+
+            // This is to protect users, because if their cpi program sends deposits to a vault
+            // and they forgot to pass in the bank for the vault, their account would not be credited.
+            Some(Err(error!(MangoError::SomeError)))
+        })
+        .collect::<Result<HashMap<_, _>>>()?;
+
+    // Find banks for used vaults in cpi_ais and collect signer seeds for them.
+    // Also update withdraw_amount and loan_amount.
+    let mut bank_signer_data = Vec::with_capacity(used_vaults.len());
+    for (ai, am) in cpi_ais.iter().zip(cpi_ams.iter_mut()) {
+        if ai.owner != &Mango::id() {
+            continue;
+        }
+        if let Some(bank) = allowed_banks.get(ai.key) {
+            if let Some(vault_info) = used_vaults.get_mut(&bank.vault) {
+                let withdraw_amount = withdraws
+                    .iter()
+                    .find_map(|&(index, amount)| {
+                        (index as usize == vault_info.vault_cpi_ai_index).then(|| amount)
+                    })
+                    // Even if we don't withdraw from a vault we still need to track it:
+                    // Possibly the invoked program will deposit funds into it.
+                    .unwrap_or(0);
+                require!(
+                    withdraw_amount <= vault_info.pre_amount,
+                    MangoError::SomeError
+                );
+                vault_info.withdraw_amount = withdraw_amount;
+
+                // if there are withdraws: figure out loan amount, mark as signer
+                if withdraw_amount > 0 {
+                    let token_account = account.tokens.get_mut_raw(vault_info.raw_token_index);
+                    let native_position = token_account.native(&bank);
+                    vault_info.loan_amount = if native_position > 0 {
+                        (I80F48::from(withdraw_amount) - native_position).max(I80F48::ZERO)
+                    } else {
+                        I80F48::from(withdraw_amount)
+                    };
+
+                    am.is_signer = true;
+                    // this is the data we'll need later to build the PDA account signer seeds
+                    bank_signer_data.push((bank.token_index.to_le_bytes(), [bank.bump]));
+                }
+            }
+        }
+    }
+
+    // Approve bank delegates for withdrawals
+    let group_seeds = group_seeds!(group);
+    let seeds = [&group_seeds[..]];
+    for (_, vault_info) in used_vaults.iter() {
+        if vault_info.withdraw_amount > 0 {
+            let approve_ctx = CpiContext::new(
+                ctx.accounts.token_program.to_account_info(),
+                token::Approve {
+                    to: cpi_ais[vault_info.vault_cpi_ai_index].clone(),
+                    delegate: health_ais[vault_info.bank_health_ai_index].clone(),
+                    authority: ctx.accounts.group.to_account_info(),
+                },
+            )
+            .with_signer(&seeds);
+            token::approve(approve_ctx, vault_info.withdraw_amount)?;
+        }
+    }
+
+    // get rid of Ref<> to avoid limiting the cpi call
+    drop(allowed_banks);
+    drop(group);
+    drop(account);
+
+    // prepare signer seeds and invoke cpi
+    let group_key = ctx.accounts.group.key();
+    let signers = bank_signer_data
+        .iter()
+        .map(|(token_index, bump)| {
+            [
+                group_key.as_ref(),
+                b"Bank".as_ref(),
+                &token_index[..],
+                &bump[..],
+            ]
+        })
+        .collect::<Vec<_>>();
+    let signers_ref = signers.iter().map(|v| &v[..]).collect::<Vec<_>>();
     let cpi_ix = Instruction {
         program_id: cpi_program_id,
         data: cpi_data,
         accounts: cpi_ams,
     };
-    let group_seeds = group_seeds!(group);
-    let pre_cpi_amounts = get_pre_cpi_amounts(&ctx, &cpi_ais);
-    solana_program::program::invoke_signed(&cpi_ix, &cpi_ais, &[group_seeds])?;
-    adjust_for_post_cpi_amounts(
-        &ctx,
-        &cpi_ais,
-        pre_cpi_amounts,
-        &mut banks.to_vec(),
-        &mut account,
-    )?;
+    solana_program::program::invoke_signed(&cpi_ix, &cpi_ais, &signers_ref)?;
 
-    // compute post cpi health
-    // todo: this is not working, the health is computed on old bank state and not taking into account
-    // withdraws done in adjust_for_post_cpi_token_amounts
+    // Revoke delegates for vaults
+    let group = ctx.accounts.group.load()?;
+    let group_seeds = group_seeds!(group);
+    for (_, vault_info) in used_vaults.iter() {
+        if vault_info.withdraw_amount > 0 {
+            let ix = token::spl_token::instruction::revoke(
+                &token::spl_token::ID,
+                &cpi_ais[vault_info.vault_cpi_ai_index].key,
+                &ctx.accounts.group.key(),
+                &[],
+            )?;
+            solana_program::program::invoke_signed(
+                &ix,
+                &[
+                    cpi_ais[vault_info.vault_cpi_ai_index].clone(),
+                    ctx.accounts.group.to_account_info(),
+                ],
+                &[group_seeds],
+            )?;
+        }
+    }
+
+    // Track vault changes and apply them to the user's token positions
+    let mut account = ctx.accounts.account.load_mut()?;
+    let inactive_tokens =
+        adjust_for_post_cpi_vault_amounts(health_ais, cpi_ais, &used_vaults, &mut account)?;
+
+    // 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);
+    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 inactive_tokens {
+        account.tokens.deactivate(raw_token_index);
+    }
+
     Ok(())
 }
 
-fn get_pre_cpi_amounts(ctx: &Context<MarginTrade>, cpi_ais: &[AccountInfo]) -> Vec<u64> {
-    let mut amounts = vec![];
-    for token_account in cpi_ais
-        .iter()
-        .filter(|ai| ai.owner == &TokenAccount::owner())
-    {
-        let vault = Account::<TokenAccount>::try_from(token_account).unwrap();
-        if vault.owner == ctx.accounts.group.key() {
-            amounts.push(vault.amount)
-        }
-    }
-    amounts
-}
-
-fn adjust_for_post_cpi_amounts(
-    ctx: &Context<MarginTrade>,
+fn adjust_for_post_cpi_vault_amounts(
+    health_ais: &[AccountInfo],
     cpi_ais: &[AccountInfo],
-    pre_cpi_amounts: Vec<u64>,
-    banks: &mut [AccountInfo],
+    used_vaults: &HashMap<&Pubkey, AllowedVault>,
     account: &mut MangoAccount,
-) -> Result<()> {
-    let token_accounts_iter = cpi_ais
-        .iter()
-        .filter(|ai| ai.owner == &TokenAccount::owner());
+) -> Result<Vec<usize>> {
+    let mut inactive_token_raw_indexes = Vec::with_capacity(used_vaults.len());
+    for (_, info) in used_vaults.iter() {
+        let vault = Account::<TokenAccount>::try_from(&cpi_ais[info.vault_cpi_ai_index]).unwrap();
+        let mut bank = health_ais[info.bank_health_ai_index].load_mut::<Bank>()?;
+        let position = account.tokens.get_mut_raw(info.raw_token_index);
 
-    for (token_account, pre_cpi_amount) in
-        // token_accounts and pre_cpi_amounts are assumed to be in correct order
-        token_accounts_iter.zip(pre_cpi_amounts.iter())
-    {
-        let vault = Account::<TokenAccount>::try_from(token_account).unwrap();
-        if vault.owner == ctx.accounts.group.key() {
-            // find bank for token account
-            let bank_ai = banks
-                .iter()
-                .find(|bank_ai| {
-                    let bank_loader = AccountLoader::<'_, Bank>::try_from(bank_ai).unwrap();
-                    let bank = bank_loader.load().unwrap();
-                    bank.mint == vault.mint
-                })
-                .ok_or(MangoError::SomeError)?; // todo: replace SomeError
-            let bank_loader = AccountLoader::<'_, Bank>::try_from(bank_ai)?;
-            let mut bank = bank_loader.load_mut()?;
+        let loan_origination_fee = info.loan_amount * bank.loan_origination_fee_rate;
+        bank.collected_fees_native += loan_origination_fee;
 
-            let position = account.tokens.get_mut_or_create(bank.token_index)?.0;
-
-            let change = I80F48::from(vault.amount) - I80F48::from(*pre_cpi_amount);
-            bank.change_with_fee(position, change)?;
+        let is_active = bank.change_without_fee(
+            position,
+            I80F48::from(vault.amount) - I80F48::from(info.pre_amount) - loan_origination_fee,
+        )?;
+        if !is_active {
+            inactive_token_raw_indexes.push(info.raw_token_index);
         }
     }
-    Ok(())
+    Ok(inactive_token_raw_indexes)
 }

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

@@ -1,12 +1,11 @@
 use anchor_lang::prelude::*;
-use anchor_spl::token::Mint;
-use anchor_spl::token::Token;
-use anchor_spl::token::TokenAccount;
+use anchor_spl::token::{Mint, Token, TokenAccount};
 use fixed::types::I80F48;
 use fixed_macro::types::I80F48;
 
 // TODO: ALTs are unavailable
 //use crate::address_lookup_table;
+use crate::error::*;
 use crate::state::*;
 use crate::util::fill16_from_str;
 
@@ -130,6 +129,7 @@ pub fn register_token(
         liquidation_fee: I80F48::from_num(liquidation_fee),
         dust: I80F48::ZERO,
         token_index,
+        bump: *ctx.bumps.get("bank").ok_or(MangoError::SomeError)?,
         reserved: Default::default(),
     };
 

programs/mango-v4/src/lib.rs

@@ -99,10 +99,11 @@ pub mod mango_v4 {
 
     pub fn margin_trade<'key, 'accounts, 'remaining, 'info>(
         ctx: Context<'key, 'accounts, 'remaining, 'info, MarginTrade<'info>>,
-        banks_len: usize,
+        num_health_accounts: usize,
+        withdraws: Vec<(u8, u64)>,
         cpi_data: Vec<u8>,
     ) -> Result<()> {
-        instructions::margin_trade(ctx, banks_len, cpi_data)
+        instructions::margin_trade(ctx, num_health_accounts, withdraws, cpi_data)
     }
 
     ///

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

@@ -59,9 +59,11 @@ pub struct Bank {
     // Index into TokenInfo on the group
     pub token_index: TokenIndex,
 
-    pub reserved: [u8; 6],
+    pub bump: u8,
+
+    pub reserved: [u8; 5],
 }
-const_assert_eq!(size_of::<Bank>(), 16 + 32 * 4 + 8 + 16 * 18 + 2 + 6);
+const_assert_eq!(size_of::<Bank>(), 16 + 32 * 4 + 8 + 16 * 18 + 3 + 5);
 const_assert_eq!(size_of::<Bank>() % 8, 0);
 
 impl std::fmt::Debug for Bank {
@@ -342,6 +344,20 @@ impl Bank {
     }
 }
 
+#[macro_export]
+macro_rules! bank_seeds {
+    ( $bank:expr ) => {
+        &[
+            $bank.group.as_ref(),
+            b"Bank".as_ref(),
+            $bank.token_index.to_le_bytes(),
+            &[$bank.bump],
+        ]
+    };
+}
+
+pub use bank_seeds;
+
 #[cfg(test)]
 mod tests {
     use bytemuck::Zeroable;

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

@@ -124,6 +124,10 @@ impl MangoAccountTokens {
             .ok_or_else(|| error!(MangoError::SomeError)) // TODO: not found error
     }
 
+    pub fn get_mut_raw(&mut self, raw_token_index: usize) -> &mut TokenAccount {
+        &mut self.values[raw_token_index]
+    }
+
     pub fn get_mut_or_create(
         &mut self,
         token_index: TokenIndex,

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

@@ -239,6 +239,17 @@ pub async fn account_position(solana: &SolanaCookie, account: Pubkey, bank: Pubk
     native.round().to_num::<i64>()
 }
 
+pub async fn account_position_f64(solana: &SolanaCookie, account: Pubkey, bank: Pubkey) -> f64 {
+    let account_data: MangoAccount = solana.get_account(account).await;
+    let bank_data: Bank = solana.get_account(bank).await;
+    let native = account_data
+        .tokens
+        .find(bank_data.token_index)
+        .unwrap()
+        .native(&bank_data);
+    native.to_num::<f64>()
+}
+
 //
 // a struct for each instruction along with its
 // ClientInstruction impl
@@ -247,7 +258,9 @@ pub async fn account_position(solana: &SolanaCookie, account: Pubkey, bank: Pubk
 pub struct MarginTradeInstruction<'keypair> {
     pub account: Pubkey,
     pub owner: &'keypair Keypair,
+    pub mango_token_bank: Pubkey,
     pub mango_token_vault: Pubkey,
+    pub withdraw_amount: u64,
     pub margin_trade_program_id: Pubkey,
     pub deposit_account: Pubkey,
     pub deposit_account_owner: Pubkey,
@@ -265,21 +278,27 @@ impl<'keypair> ClientInstruction for MarginTradeInstruction<'keypair> {
 
         let account: MangoAccount = account_loader.load(&self.account).await.unwrap();
 
-        let instruction = Self::Instruction {
-            banks_len: account.tokens.iter_active().count(),
-            cpi_data: self.margin_trade_program_ix_cpi_data.clone(),
-        };
-
         let accounts = Self::Accounts {
             group: account.group,
             account: self.account,
             owner: self.owner.pubkey(),
+            token_program: Token::id(),
         };
 
-        let health_check_metas =
-            derive_health_check_remaining_account_metas(&account_loader, &account, None, true)
+        let health_check_metas = derive_health_check_remaining_account_metas(
+            &account_loader,
+            &account,
+            Some(self.mango_token_bank),
+            true,
+        )
         .await;
 
+        let instruction = Self::Instruction {
+            num_health_accounts: health_check_metas.len(),
+            withdraws: vec![(1, self.withdraw_amount)],
+            cpi_data: self.margin_trade_program_ix_cpi_data.clone(),
+        };
+
         let mut instruction = make_instruction(program_id, &accounts, instruction);
         instruction.accounts.extend(health_check_metas.into_iter());
         instruction.accounts.push(AccountMeta {
@@ -287,6 +306,11 @@ impl<'keypair> ClientInstruction for MarginTradeInstruction<'keypair> {
             is_writable: false,
             is_signer: false,
         });
+        instruction.accounts.push(AccountMeta {
+            pubkey: self.mango_token_bank,
+            is_writable: false,
+            is_signer: false,
+        });
         instruction.accounts.push(AccountMeta {
             pubkey: self.mango_token_vault,
             is_writable: true,

programs/mango-v4/tests/test_margin_trade.rs

@@ -1,7 +1,6 @@
 #![cfg(feature = "test-bpf")]
 
 use anchor_lang::InstructionData;
-use fixed::types::I80F48;
 use solana_program_test::*;
 use solana_sdk::signature::Keypair;
 use solana_sdk::signature::Signer;
@@ -23,9 +22,13 @@ async fn test_margin_trade() -> Result<(), BanksClientError> {
     let admin = &Keypair::new();
     let owner = &context.users[0].key;
     let payer = &context.users[1].key;
-    let mints = &context.mints[0..1];
+    let mints = &context.mints[0..2];
     let payer_mint0_account = context.users[1].token_accounts[0];
-    let dust_threshold = 0.01;
+    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)
@@ -41,6 +44,51 @@ async fn test_margin_trade() -> Result<(), BanksClientError> {
     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,
+        DepositInstruction {
+            amount: provided_amount,
+            account: provider_account,
+            token_account: payer_mint0_account,
+            token_authority: payer,
+        },
+    )
+    .await
+    .unwrap();
+    send_tx(
+        solana,
+        DepositInstruction {
+            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 {
@@ -75,22 +123,15 @@ async fn test_margin_trade() -> Result<(), BanksClientError> {
 
         assert_eq!(
             solana.token_account_balance(vault).await,
-            deposit_amount_initial
+            provided_amount + deposit_amount_initial
         );
         assert_eq!(
             solana.token_account_balance(payer_mint0_account).await,
             start_balance - deposit_amount_initial
         );
-        let account_data: MangoAccount = solana.get_account(account).await;
-        let bank_data: Bank = solana.get_account(bank).await;
-        assert!(
-            account_data.tokens.values[0].native(&bank_data)
-                - I80F48::from_num(deposit_amount_initial)
-                < dust_threshold
-        );
-        assert!(
-            bank_data.native_total_deposits() - I80F48::from_num(deposit_amount_initial)
-                < dust_threshold
+        assert_eq!(
+            account_position(solana, account, bank).await,
+            deposit_amount_initial as i64,
         );
     }
 
@@ -105,14 +146,16 @@ async fn test_margin_trade() -> Result<(), BanksClientError> {
             MarginTradeInstruction {
                 account,
                 owner,
+                mango_token_bank: bank,
                 mango_token_vault: vault,
+                withdraw_amount,
                 margin_trade_program_id: margin_trade.program,
                 deposit_account: margin_trade.token_account.pubkey(),
                 deposit_account_owner: margin_trade.token_account_owner,
                 margin_trade_program_ix_cpi_data: {
                     let ix = margin_trade::instruction::MarginTrade {
-                        amount_from: 2,
-                        amount_to: 1,
+                        amount_from: withdraw_amount,
+                        amount_to: deposit_amount,
                         deposit_account_owner_bump_seeds: margin_trade.token_account_bump,
                     };
                     ix.data()
@@ -124,7 +167,7 @@ async fn test_margin_trade() -> Result<(), BanksClientError> {
     }
     assert_eq!(
         solana.token_account_balance(vault).await,
-        deposit_amount_initial - withdraw_amount + deposit_amount
+        provided_amount + deposit_amount_initial - withdraw_amount + deposit_amount
     );
     assert_eq!(
         solana
@@ -132,6 +175,151 @@ async fn test_margin_trade() -> Result<(), BanksClientError> {
             .await,
         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_trade.token_account.pubkey())
+        .await;
+    let withdraw_amount = deposit_amount_initial as u64;
+    let deposit_amount = 0;
+    {
+        send_tx(
+            solana,
+            MarginTradeInstruction {
+                account,
+                owner,
+                mango_token_bank: bank,
+                mango_token_vault: vault,
+                withdraw_amount,
+                margin_trade_program_id: margin_trade.program,
+                deposit_account: margin_trade.token_account.pubkey(),
+                deposit_account_owner: margin_trade.token_account_owner,
+                margin_trade_program_ix_cpi_data: {
+                    let ix = margin_trade::instruction::MarginTrade {
+                        amount_from: withdraw_amount,
+                        amount_to: deposit_amount,
+                        deposit_account_owner_bump_seeds: margin_trade.token_account_bump,
+                    };
+                    ix.data()
+                },
+            },
+        )
+        .await
+        .unwrap();
+    }
+    assert_eq!(solana.token_account_balance(vault).await, provided_amount);
+    assert_eq!(
+        solana
+            .token_account_balance(margin_trade.token_account.pubkey())
+            .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_trade.token_account.pubkey())
+        .await;
+    let withdraw_amount = 0;
+    let deposit_amount = margin_account_initial;
+    {
+        send_tx(
+            solana,
+            MarginTradeInstruction {
+                account,
+                owner,
+                mango_token_bank: bank,
+                mango_token_vault: vault,
+                withdraw_amount,
+                margin_trade_program_id: margin_trade.program,
+                deposit_account: margin_trade.token_account.pubkey(),
+                deposit_account_owner: margin_trade.token_account_owner,
+                margin_trade_program_ix_cpi_data: {
+                    let ix = margin_trade::instruction::MarginTrade {
+                        amount_from: withdraw_amount,
+                        amount_to: deposit_amount,
+                        deposit_account_owner_bump_seeds: margin_trade.token_account_bump,
+                    };
+                    ix.data()
+                },
+            },
+        )
+        .await
+        .unwrap();
+    }
+    assert_eq!(
+        solana.token_account_balance(vault).await,
+        provided_amount + deposit_amount
+    );
+    assert_eq!(
+        solana
+            .token_account_balance(margin_trade.token_account.pubkey())
+            .await,
+        0
+    );
+    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 deposit_amount_initial = account_position(solana, account, bank).await as u64;
+    let withdraw_amount = 500;
+    let deposit_amount = 450;
+    {
+        send_tx(
+            solana,
+            MarginTradeInstruction {
+                account,
+                owner,
+                mango_token_bank: bank,
+                mango_token_vault: vault,
+                withdraw_amount,
+                margin_trade_program_id: margin_trade.program,
+                deposit_account: margin_trade.token_account.pubkey(),
+                deposit_account_owner: margin_trade.token_account_owner,
+                margin_trade_program_ix_cpi_data: {
+                    let ix = margin_trade::instruction::MarginTrade {
+                        amount_from: withdraw_amount,
+                        amount_to: deposit_amount,
+                        deposit_account_owner_bump_seeds: margin_trade.token_account_bump,
+                    };
+                    ix.data()
+                },
+            },
+        )
+        .await
+        .unwrap();
+    }
+    assert_eq!(
+        solana.token_account_balance(vault).await,
+        provided_amount + deposit_amount_initial + deposit_amount - withdraw_amount
+    );
+    assert_eq!(
+        solana
+            .token_account_balance(margin_trade.token_account.pubkey())
+            .await,
+        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(())
 }

programs/margin-trade/src/lib.rs

@@ -14,12 +14,14 @@ pub mod margin_trade {
         deposit_account_owner_bump_seeds: u8,
         amount_to: u64,
     ) -> Result<()> {
+        if amount_from > 0 {
             msg!(
                 "withdrawing({}) for mint {:?}",
                 amount_from,
                 ctx.accounts.withdraw_account.mint
             );
             token::transfer(ctx.accounts.transfer_from_mango_vault_ctx(), amount_from)?;
+        }
 
         msg!("TODO: do something with the loan");
 
@@ -52,7 +54,7 @@ impl anchor_lang::Id for MarginTrade {
 #[derive(Accounts)]
 
 pub struct MarginTradeCtx<'info> {
-    pub withdraw_account_owner: Signer<'info>,
+    pub withdraw_account_owner: UncheckedAccount<'info>,
 
     #[account(mut)]
     pub withdraw_account: Account<'info, TokenAccount>,