refactor: defer deposit/borrow limit checks to the end_flashloan ix

programs/marginfi/src/instructions/marginfi_account/borrow.rs

@@ -4,7 +4,9 @@ use crate::{
     events::{AccountEventHeader, LendingAccountBorrowEvent},
     prelude::{MarginfiError, MarginfiGroup, MarginfiResult},
     state::{
-        marginfi_account::{BankAccountWrapper, MarginfiAccount, RiskEngine, DISABLED_FLAG},
+        marginfi_account::{
+            BankAccountWrapper, MarginfiAccount, RiskEngine, DISABLED_FLAG, IN_FLASHLOAN_FLAG,
+        },
         marginfi_group::{Bank, BankVaultType},
     },
     utils,
@@ -57,6 +59,8 @@ pub fn lending_account_borrow<'info>(
     {
         let mut bank = bank_loader.load_mut()?;
 
+        let is_flashloan = marginfi_account.get_flag(IN_FLASHLOAN_FLAG);
+
         let liquidity_vault_authority_bump = bank.liquidity_vault_authority_bump;
 
         let mut bank_account = BankAccountWrapper::find_or_create(
@@ -78,7 +82,12 @@ pub fn lending_account_borrow<'info>(
             .transpose()?
             .unwrap_or(amount);
 
-        bank_account.borrow(I80F48::from_num(amount_pre_fee))?;
+        if is_flashloan {
+            bank_account.borrow_skip_limit_checks(I80F48::from_num(amount_pre_fee))?;
+        } else {
+            bank_account.borrow(I80F48::from_num(amount_pre_fee))?;
+        }
+
         bank_account.withdraw_spl_transfer(
             amount_pre_fee,
             bank_liquidity_vault.to_account_info(),

programs/marginfi/src/instructions/marginfi_account/flashloan.rs

@@ -133,6 +133,8 @@ pub fn lending_account_end_flashloan<'info>(
 
     marginfi_account.unset_flag(IN_FLASHLOAN_FLAG);
 
+    // check bank deposit/borrow limits
+    RiskEngine::check_bank_deposit_borrow_limit(ctx.remaining_accounts)?;
     RiskEngine::check_account_init_health(&marginfi_account, ctx.remaining_accounts)?;
 
     Ok(())

programs/marginfi/src/state/marginfi_account.rs

@@ -128,6 +128,7 @@ pub enum BalanceDecreaseType {
     WithdrawOnly,
     BorrowOnly,
     BypassBorrowLimit,
+    BypassBorrowAndDepositLimit,
 }
 
 #[derive(Copy, Clone)]
@@ -174,7 +175,7 @@ impl<'info> BankAccountWithPriceFeed<'_, 'info> {
 
         debug!("Expecting {} remaining accounts", active_balances.len() * 2);
         debug!("Got {} remaining accounts", remaining_ais.len());
-
+        
         check!(
             active_balances.len() * 2 <= remaining_ais.len(),
             MarginfiError::MissingPythOrBankAccount
@@ -473,6 +474,28 @@ impl<'info> RiskEngine<'_, 'info> {
         Ok(())
     }
 
+    pub fn check_bank_deposit_borrow_limit<'a>(
+        remaining_ais: &'info [AccountInfo<'info>],
+    ) -> MarginfiResult<()> {
+        let bank_len = remaining_ais.len();
+        if bank_len == 0 {
+            return Ok(());
+        }
+
+        for index in 0..bank_len {
+            let bank_index = index * 2;
+            if bank_index >= remaining_ais.len() {
+                break;
+            }
+            let bank_ai = remaining_ais.get(bank_index).unwrap();
+            let bank_al = AccountLoader::<Bank>::try_from(bank_ai)?;
+            let bank = bank_al.load()?;
+            bank.check_deposit_borrow_limits()?;
+        }
+
+        Ok(())
+    }
+
     /// Returns the total assets and liabilities of the account in the form of (assets, liabilities)
     pub fn get_account_health_components(
         &self,
@@ -920,6 +943,11 @@ impl<'a> BankAccountWrapper<'a> {
         self.decrease_balance_internal(amount, BalanceDecreaseType::Any)
     }
 
+    /// Like borrow fn, but defer deposit/borrow limit checks to the end_flashloan ix.
+    pub fn borrow_skip_limit_checks(&mut self, amount: I80F48) -> MarginfiResult {
+        self.decrease_balance_internal(amount, BalanceDecreaseType::BypassBorrowAndDepositLimit)
+    }
+
     // ------------ Hybrid operations for seamless repay + deposit / withdraw + borrow
 
     /// Repay liability and deposit/increase asset depending on
@@ -1188,13 +1216,24 @@ impl<'a> BankAccountWrapper<'a> {
 
         let asset_shares_decrease = bank.get_asset_shares(asset_amount_decrease)?;
         balance.change_asset_shares(-asset_shares_decrease)?;
-        bank.change_asset_shares(-asset_shares_decrease, false)?;
+
+        bank.change_asset_shares(
+            -asset_shares_decrease,
+            matches!(
+                operation_type,
+                BalanceDecreaseType::BypassBorrowAndDepositLimit
+            ),
+        )?;
 
         let liability_shares_increase = bank.get_liability_shares(liability_amount_increase)?;
         balance.change_liability_shares(liability_shares_increase)?;
         bank.change_liability_shares(
             liability_shares_increase,
-            matches!(operation_type, BalanceDecreaseType::BypassBorrowLimit),
+            matches!(operation_type, BalanceDecreaseType::BypassBorrowLimit)
+                || matches!(
+                    operation_type,
+                    BalanceDecreaseType::BypassBorrowAndDepositLimit
+                ),
         )?;
 
         bank.check_utilization_ratio()?;

programs/marginfi/src/state/marginfi_group.rs

@@ -441,6 +441,32 @@ impl Bank {
         Ok(())
     }
 
+    pub fn check_deposit_borrow_limits(&self) -> MarginfiResult {
+        debug!("=== check_deposit_borrow_limits ===");
+        if self.config.is_deposit_limit_active() && self.config.is_borrow_limit_active() {
+            let total_deposits_amount = self.get_asset_amount(self.total_asset_shares.into())?;
+            let deposit_limit = I80F48::from_num(self.config.deposit_limit);
+            let total_liability_amount =
+                self.get_liability_amount(self.total_liability_shares.into())?;
+            let borrow_limit = I80F48::from_num(self.config.borrow_limit);
+
+            debug!("total_deposits_amount: {}", total_deposits_amount);
+            debug!("deposit_limit: {}", deposit_limit);
+            debug!("total_liability_amount: {}", total_liability_amount);
+            debug!("borrow_limit: {}", borrow_limit);
+
+            check!(
+                total_deposits_amount < deposit_limit,
+                crate::prelude::MarginfiError::BankAssetCapacityExceeded
+            );
+            check!(
+                total_liability_amount < borrow_limit,
+                crate::prelude::MarginfiError::BankLiabilityCapacityExceeded
+            )
+        }
+        Ok(())
+    }
+
     pub fn maybe_get_asset_weight_init_discount(
         &self,
         price: I80F48,

programs/marginfi/tests/user_actions/flash_loan.rs

@@ -18,6 +18,9 @@ use solana_sdk::{
 // 7. Flashloan fails because of invalid `end_flashloan` ix order
 // 8. Flashloan fails because `end_flashloan` ix is for another account
 // 9. Flashloan fails because account is already in a flashloan
+// 10. Flashloan success (1 action) defer deposit/borrow limit checks to the end_flashloan ix
+
+
 
 #[tokio::test]
 async fn flashloan_success_1op() -> anyhow::Result<()> {
@@ -535,3 +538,52 @@ async fn flashloan_fail_already_in_flashloan() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn flashloan_success_defer_deposit_borrow_limit_check() -> anyhow::Result<()> {
+    // Setup test executor with non-admin payer
+    let test_f = TestFixture::new(Some(TestSettings::all_banks_payer_not_admin())).await;
+
+    let sol_bank = test_f.get_bank(&BankMint::Sol);
+
+    // Fund SOL lender
+    let lender_mfi_account_f = test_f.create_marginfi_account().await;
+    let lender_token_account_f_sol = test_f
+        .sol_mint
+        .create_token_account_and_mint_to(1_000)
+        .await;
+    lender_mfi_account_f
+        .try_bank_deposit(lender_token_account_f_sol.key, sol_bank, 1_000)
+        .await?;
+
+    // Fund SOL borrower
+    let borrower_mfi_account_f = test_f.create_marginfi_account().await;
+
+    borrower_mfi_account_f
+        .try_set_flag(FLASHLOAN_ENABLED_FLAG)
+        .await?;
+
+    let borrower_token_account_f_sol = test_f.sol_mint.create_empty_token_account().await;
+
+    // Borrow SOL
+    let borrow_ix = borrower_mfi_account_f
+        .make_bank_borrow_ix(borrower_token_account_f_sol.key, sol_bank, 1_000)
+        .await;
+
+    let repay_ix = borrower_mfi_account_f
+        .make_bank_repay_ix(
+            borrower_token_account_f_sol.key,
+            sol_bank,
+            1_000,
+            Some(true),
+        )
+        .await;
+
+    let flash_loan_result = borrower_mfi_account_f
+        .try_flashloan(vec![borrow_ix, repay_ix], vec![], vec![sol_bank.key])
+        .await;
+
+    assert!(flash_loan_result.is_ok());
+
+    Ok(())
+}