Crispheaney/withdraw in rm liquidity

programs/drift/src/instructions/keeper.rs

@@ -3334,7 +3334,12 @@ pub fn handle_update_amm_cache<'c: 'info, 'info>(
         )?;
 
         cached_info.update_perp_market_fields(&perp_market)?;
-        cached_info.update_oracle_info(slot, &mm_oracle_price_data, &perp_market, &state.oracle_guard_rails)?;
+        cached_info.update_oracle_info(
+            slot,
+            &mm_oracle_price_data,
+            &perp_market,
+            &state.oracle_guard_rails,
+        )?;
 
         if perp_market.lp_status != 0 {
             amm_cache.update_amount_owed_from_lp_pool(&perp_market, &quote_market)?;

programs/drift/src/instructions/lp_pool.rs

@@ -44,6 +44,8 @@ use crate::{
     validate,
 };
 use std::convert::TryFrom;
+use std::iter::Peekable;
+use std::slice::Iter;
 
 use solana_program::sysvar::clock::Clock;
 
@@ -1055,17 +1057,39 @@ pub fn handle_lp_pool_remove_liquidity<'c: 'info, 'info>(
     } else {
         out_amount.safe_add(out_fee_amount.unsigned_abs())?
     };
-    let out_amount_net_fees =
-        out_amount_net_fees.min(ctx.accounts.constituent_out_token_account.amount as u128);
 
     validate!(
         out_amount_net_fees >= min_amount_out,
         ErrorCode::SlippageOutsideLimit,
-        format!(
-            "Slippage outside limit: lp_mint_amount_net_fees({}) < min_mint_amount({})",
-            out_amount_net_fees, min_amount_out
-        )
-        .as_str()
+        "Slippage outside limit: out_amount_net_fees({}) < min_amount_out({})",
+        out_amount_net_fees,
+        min_amount_out
+    )?;
+
+    if out_amount_net_fees > out_constituent.vault_token_balance.cast()? {
+        let transfer_amount = out_amount_net_fees.cast::<u64>()?.safe_sub(out_constituent.vault_token_balance)?;
+        msg!("transfering from program vault to constituent vault: {}", transfer_amount);
+        transfer_from_program_vault(
+            transfer_amount,
+            &mut out_spot_market,
+            &mut out_constituent,
+            out_oracle.price,
+            &ctx.accounts.state,
+            &mut ctx.accounts.spot_market_token_account,
+            &mut ctx.accounts.constituent_out_token_account,
+            &ctx.accounts.token_program,
+            &ctx.accounts.drift_signer,
+            &None,
+            Some(remaining_accounts),
+        )?;
+    }
+
+    validate!(
+        out_amount_net_fees <= out_constituent.vault_token_balance.cast()?,
+        ErrorCode::InsufficientConstituentTokenBalance,
+        "Insufficient out constituent balance: out_amount_net_fees({}) > out_constituent.token_balance({})",
+        out_amount_net_fees,
+        out_constituent.vault_token_balance
     )?;
 
     out_constituent.record_swap_fees(out_fee_amount)?;
@@ -1404,7 +1428,6 @@ pub fn handle_withdraw_from_program_vault<'c: 'info, 'info>(
 
     let mut spot_market = ctx.accounts.spot_market.load_mut()?;
 
-    let spot_market_vault = &ctx.accounts.spot_market_vault;
     let oracle_id = spot_market.oracle_id();
     let mut oracle_map = OracleMap::load_one(
         &ctx.accounts.oracle,
@@ -1430,61 +1453,85 @@ pub fn handle_withdraw_from_program_vault<'c: 'info, 'info>(
         constituent.last_oracle_price = oracle_data.price;
         constituent.last_oracle_slot = oracle_data_slot;
     }
-    constituent.sync_token_balance(ctx.accounts.constituent_token_account.amount);
+
+    let mint = &Some(*ctx.accounts.mint.clone());
+    transfer_from_program_vault(
+        amount,
+        &mut spot_market,
+        &mut constituent,
+        oracle_data.price,
+        &state,
+        &mut ctx.accounts.spot_market_vault,
+        &mut ctx.accounts.constituent_token_account,
+        &ctx.accounts.token_program,
+        &ctx.accounts.drift_signer,
+        mint,
+        Some(remaining_accounts),
+    )?;
+
+    Ok(())
+}
+
+fn transfer_from_program_vault<'info>(
+    amount: u64,
+    spot_market: &mut SpotMarket,
+    constituent: &mut Constituent,
+    oracle_price: i64,
+    state: &State,
+    spot_market_vault: &mut InterfaceAccount<'info, TokenAccount>,
+    constituent_token_account: &mut InterfaceAccount<'info, TokenAccount>,
+    token_program: &Interface<'info, TokenInterface>,
+    drift_signer: &AccountInfo<'info>,
+    mint: &Option<InterfaceAccount<'info, Mint>>,
+    remaining_accounts: Option<&mut Peekable<Iter<'info, AccountInfo<'info>>>>,
+) -> Result<()> {
+    constituent.sync_token_balance(constituent_token_account.amount);
+
     let balance_before = constituent.get_full_token_amount(&spot_market)?;
 
-    // Can only borrow up to the max
-    let bl_token_balance = constituent.spot_balance.get_token_amount(&spot_market)?;
-    let amount_to_transfer = if constituent.spot_balance.balance_type == SpotBalanceType::Borrow {
-        amount.min(
-            constituent
-                .max_borrow_token_amount
-                .saturating_sub(bl_token_balance as u64),
-        )
-    } else {
-        amount.min(
-            constituent
-                .max_borrow_token_amount
-                .saturating_add(bl_token_balance as u64),
-        )
-    };
+    let max_transfer = constituent.get_max_transfer(&spot_market)?;
+
+    validate!(
+        max_transfer >= amount,
+        ErrorCode::LpInvariantFailed,
+        "Max transfer ({} is less than amount ({})",
+        max_transfer,
+        amount
+    )?;
 
     // Execute transfer and sync new balance in the constituent account
     controller::token::send_from_program_vault(
-        &ctx.accounts.token_program,
+        &token_program,
         &spot_market_vault,
-        &ctx.accounts.constituent_token_account,
-        &ctx.accounts.drift_signer,
+        &constituent_token_account,
+        &drift_signer,
         state.signer_nonce,
-        amount_to_transfer,
-        &Some(*ctx.accounts.mint.clone()),
-        Some(remaining_accounts),
+        amount,
+        mint,
+        remaining_accounts,
     )?;
-    ctx.accounts.constituent_token_account.reload()?;
-    constituent.sync_token_balance(ctx.accounts.constituent_token_account.amount);
+    constituent_token_account.reload()?;
+    constituent.sync_token_balance(constituent_token_account.amount);
 
     // Adjust BLPosition for the new deposits
     let spot_position = &mut constituent.spot_balance;
     update_spot_balances(
-        amount_to_transfer as u128,
+        amount as u128,
         &SpotBalanceType::Borrow,
-        &mut spot_market,
+        spot_market,
         spot_position,
         true,
     )?;
 
     safe_decrement!(
         spot_position.cumulative_deposits,
-        amount_to_transfer.cast()?
+        amount.cast()?
     );
 
     // Re-check spot market invariants
-    ctx.accounts.spot_market_vault.reload()?;
+    spot_market_vault.reload()?;
     spot_market.validate_max_token_deposits_and_borrows(true)?;
-    math::spot_withdraw::validate_spot_market_vault_amount(
-        &spot_market,
-        ctx.accounts.spot_market_vault.amount,
-    )?;
+    math::spot_withdraw::validate_spot_market_vault_amount(&spot_market, spot_market_vault.amount)?;
 
     // Verify withdraw fully accounted for in BLPosition
     let balance_after = constituent.get_full_token_amount(&spot_market)?;
@@ -1493,20 +1540,18 @@ pub fn handle_withdraw_from_program_vault<'c: 'info, 'info>(
         balance_after
             .abs_diff(balance_before)
             .cast::<i64>()?
-            .safe_mul(oracle_data.price)?
+            .safe_mul(oracle_price)?
             .safe_div(PRICE_PRECISION_I64)?
             .safe_div(10_i64.pow(spot_market.decimals - 6))?
     } else {
         balance_after
             .abs_diff(balance_before)
             .cast::<i64>()?
             .safe_mul(10_i64.pow(6 - spot_market.decimals))?
-            .safe_mul(oracle_data.price)?
+            .safe_mul(oracle_price)?
             .safe_div(PRICE_PRECISION_I64)?
     };
 
-    msg!("Balance difference (notional): {}", balance_diff_notional);
-
     validate!(
         balance_diff_notional <= PRICE_PRECISION_I64 / 100,
         ErrorCode::LpInvariantFailed,
@@ -1836,17 +1881,22 @@ pub struct ViewLPPoolAddLiquidityFees<'info> {
 
 #[derive(Accounts)]
 #[instruction(
-    in_market_index: u16,
+    out_market_index: u16,
 )]
 pub struct LPPoolRemoveLiquidity<'info> {
     pub state: Box<Account<'info, State>>,
+    #[account(
+        constraint = drift_signer.key() == state.signer
+    )]
+    /// CHECK: drift_signer
+    pub drift_signer: AccountInfo<'info>,
     #[account(mut)]
     pub lp_pool: AccountLoader<'info, LPPool>,
     pub authority: Signer<'info>,
     pub out_market_mint: Box<InterfaceAccount<'info, Mint>>,
     #[account(
         mut,
-        seeds = [CONSTITUENT_PDA_SEED.as_ref(), lp_pool.key().as_ref(), in_market_index.to_le_bytes().as_ref()],
+        seeds = [CONSTITUENT_PDA_SEED.as_ref(), lp_pool.key().as_ref(), out_market_index.to_le_bytes().as_ref()],
         bump,
         constraint =
             out_constituent.load()?.mint.eq(&constituent_out_token_account.mint)
@@ -1860,7 +1910,7 @@ pub struct LPPoolRemoveLiquidity<'info> {
     pub user_out_token_account: Box<InterfaceAccount<'info, TokenAccount>>,
     #[account(
         mut,
-        seeds = ["CONSTITUENT_VAULT".as_ref(), lp_pool.key().as_ref(), in_market_index.to_le_bytes().as_ref()],
+        seeds = ["CONSTITUENT_VAULT".as_ref(), lp_pool.key().as_ref(), out_market_index.to_le_bytes().as_ref()],
         bump,
     )]
     pub constituent_out_token_account: Box<InterfaceAccount<'info, TokenAccount>>,
@@ -1869,6 +1919,12 @@ pub struct LPPoolRemoveLiquidity<'info> {
         constraint = user_lp_token_account.mint.eq(&lp_mint.key())
     )]
     pub user_lp_token_account: Box<InterfaceAccount<'info, TokenAccount>>,
+    #[account(
+        mut,
+        seeds = [b"spot_market_vault".as_ref(), out_market_index.to_le_bytes().as_ref()],
+        bump,
+    )]
+    pub spot_market_token_account: Box<InterfaceAccount<'info, TokenAccount>>,
 
     #[account(
         mut,

programs/drift/src/state/lp_pool.rs

@@ -946,6 +946,18 @@ impl Constituent {
             bytemuck::bytes_of(bump),
         ]
     }
+
+    pub fn get_max_transfer(&self, spot_market: &SpotMarket) -> DriftResult<u64> {
+        let token_amount = self.get_full_token_amount(spot_market)?;
+
+        let max_transfer = if self.spot_balance.balance_type == SpotBalanceType::Borrow {
+            self.max_borrow_token_amount.saturating_sub(token_amount as u64)
+        } else {
+            self.max_borrow_token_amount.saturating_add(token_amount as u64)
+        };
+
+        Ok(max_transfer)
+    }
 }
 
 #[zero_copy]

sdk/src/driftClient.ts

@@ -10891,6 +10891,7 @@ export class DriftClient {
 						userOutTokenAccount,
 						constituentOutTokenAccount,
 						userLpTokenAccount,
+						spotMarketTokenAccount: spotMarket.vault, 
 						lpMint,
 						lpPoolTokenVault: getLpPoolTokenVaultPublicKey(
 							this.program.programId,

sdk/src/idl/drift.json

@@ -8563,6 +8563,11 @@
           "isMut": false,
           "isSigner": false
         },
+        {
+          "name": "driftSigner",
+          "isMut": false,
+          "isSigner": false
+        },
         {
           "name": "lpPool",
           "isMut": true,
@@ -8598,6 +8603,11 @@
           "isMut": true,
           "isSigner": false
         },
+        {
+          "name": "spotMarketTokenAccount",
+          "isMut": true,
+          "isSigner": false
+        },
         {
           "name": "lpMint",
           "isMut": true,

tests/lpPool.ts

@@ -1494,26 +1494,16 @@ describe('LP Pool', () => {
 		const balanceBefore = constituent.vaultTokenBalance;
 		const spotBalanceBefore = constituent.spotBalance;
 
+		try {
 		await adminClient.withdrawFromProgramVault(
 			encodeName(lpPoolName),
-			0,
-			new BN(100).mul(QUOTE_PRECISION)
-		);
-
-		constituent = (await adminClient.program.account.constituent.fetch(
-			getConstituentPublicKey(program.programId, lpPoolKey, 0)
-		)) as ConstituentAccount;
-
-		assert(
-			constituent.vaultTokenBalance
-				.sub(balanceBefore)
-				.eq(new BN(10).mul(QUOTE_PRECISION))
-		);
-		expect(
-			constituent.spotBalance.scaledBalance
-				.sub(spotBalanceBefore.scaledBalance)
-				.toNumber()
-		).to.be.approximately(10 * 10 ** 9, 1);
+				0,
+				new BN(100).mul(QUOTE_PRECISION)
+			);
+		} catch (e) {
+			console.log(e);
+			assert(e.toString().includes('0x18b9')); // invariant failed
+		}
 	});
 
 	it('cant disable lp pool settling', async () => {