Flash Loan Validation Remediation

Overview

Related Detector: Flash Loan Validation

Flash loan validation vulnerabilities occur when a program lends tokens and then fails to verify that the borrower repaid the loan (plus fees) within the same transaction. The fix is to record the pool balance before lending, then verify it is at least pre_balance + fee after the borrower’s callback completes.

Recommended Fix

Before (Vulnerable)

pub fn flash_borrow(accounts: &[AccountInfo], amount: u64) -> ProgramResult {
    let pool = &accounts[0];
    let borrower = &accounts[1];

    // Lend tokens
    **pool.try_borrow_mut_lamports()? -= amount;
    **borrower.try_borrow_mut_lamports()? += amount;

    // Borrower callback
    invoke(&callback_ix, accounts)?;

    // MISSING: no repayment check
    Ok(())
}

After (Fixed)

pub fn flash_borrow(accounts: &[AccountInfo], amount: u64) -> ProgramResult {
    let pool = &accounts[0];
    let borrower = &accounts[1];

    // FIXED: record pre-loan balance
    let pre_balance = **pool.try_borrow_lamports()?;

    // FIXED: calculate minimum repayment
    let fee = amount
        .checked_mul(FLASH_FEE_BPS)
        .ok_or(ProgramError::ArithmeticOverflow)?
        .checked_div(10_000)
        .ok_or(ProgramError::ArithmeticOverflow)?
        .max(1);

    // Lend tokens
    **pool.try_borrow_mut_lamports()? -= amount;
    **borrower.try_borrow_mut_lamports()? += amount;

    // Borrower callback
    invoke(&callback_ix, accounts)?;

    // FIXED: verify repayment
    let post_balance = **pool.try_borrow_lamports()?;
    let min_balance = pre_balance
        .checked_add(fee)
        .ok_or(ProgramError::ArithmeticOverflow)?;

    if post_balance < min_balance {
        return Err(ProgramError::InsufficientFunds);
    }
    Ok(())
}

Alternative Mitigations

1. Two-instruction flash loan pattern

Split the flash loan into borrow and repay instructions, with the repay instruction enforcing the invariant:

pub fn flash_borrow(ctx: Context<FlashBorrow>, amount: u64) -> Result<()> {
    let pool = &mut ctx.accounts.pool;

    // Record the loan in pool state
    pool.outstanding_flash_loan = amount;
    pool.flash_loan_fee = calculate_fee(amount);
    pool.flash_loan_borrower = *ctx.accounts.borrower.key;

    // Transfer to borrower
    // ...
    Ok(())
}

pub fn flash_repay(ctx: Context<FlashRepay>) -> Result<()> {
    let pool = &mut ctx.accounts.pool;

    // Verify outstanding loan exists
    require!(pool.outstanding_flash_loan > 0, ErrorCode::NoOutstandingLoan);

    // Verify repayment was made (pool balance check)
    let required = pool.outstanding_flash_loan + pool.flash_loan_fee;
    // ... verify balance

    // Clear loan state
    pool.outstanding_flash_loan = 0;
    pool.flash_loan_fee = 0;
    Ok(())
}

2. CPI callback with program-controlled repayment

Instead of relying on the borrower to repay, use a CPI callback pattern where the program transfers funds back automatically:

pub fn flash_loan(accounts: &[AccountInfo], amount: u64) -> ProgramResult {
    let pool = &accounts[0];
    let borrower = &accounts[1];

    let pre_balance = **pool.try_borrow_lamports()?;

    // Lend
    **pool.try_borrow_mut_lamports()? -= amount;
    **borrower.try_borrow_mut_lamports()? += amount;

    // Borrower callback
    invoke(&callback_ix, accounts)?;

    // Program-controlled repayment: pull funds back
    let repayment = amount + calculate_fee(amount);
    **borrower.try_borrow_mut_lamports()? -= repayment;
    **pool.try_borrow_mut_lamports()? += repayment;

    Ok(())
}

3. Invariant check at instruction end

Use a post-instruction invariant check that verifies pool solvency regardless of what operations occurred:

pub fn verify_pool_invariant(pool: &AccountInfo, expected_min: u64) -> ProgramResult {
    let balance = **pool.try_borrow_lamports()?;
    if balance < expected_min {
        msg!("Pool invariant violated: {} < {}", balance, expected_min);
        return Err(ProgramError::InvalidAccountData);
    }
    Ok(())
}

Common Mistakes

Mistake 1: Checking balance before the callback instead of after

// WRONG: balance checked before callback, not after
let balance = **pool.try_borrow_lamports()?;
assert!(balance >= min_required);  // This checks the pre-lend balance
invoke(&callback_ix, accounts)?;   // Borrower defaults after this

The balance check must occur after all external calls have completed.

Mistake 2: Using the wrong account for balance verification

invoke(&callback_ix, accounts)?;
// WRONG: checking borrower's balance instead of pool's
let borrower_balance = **borrower.try_borrow_lamports()?;

Verify the lending pool’s balance, not the borrower’s.

Mistake 3: Fee calculation that rounds to zero

// WRONG: for small amounts, fee rounds to zero
let fee = amount * FEE_BPS / 10_000;  // fee = 0 when amount < 10000/FEE_BPS

Enforce a minimum fee of at least 1 lamport: .max(1).

Mistake 4: No cap on flash loan amount

// WRONG: no maximum borrow amount
**pool.try_borrow_mut_lamports()? -= amount;  // amount could be pool's entire balance

Cap the flash loan to a percentage of the pool balance to limit exposure.