mev-beta/contracts/balancer/FlashLoanReceiverSecure.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import "@openzeppelin/contracts/utils/ReentrancyGuard.sol";

interface IBalancerVault {
    function flashLoan(
        address recipient,
        IERC20[] memory tokens,
        uint256[] memory amounts,
        bytes memory userData
    ) external;
}

interface IUniswapV2Router {
    function swapExactTokensForTokens(
        uint256 amountIn,
        uint256 amountOutMin,
        address[] calldata path,
        address to,
        uint256 deadline
    ) external returns (uint256[] memory amounts);
}

interface IUniswapV3Router {
    struct ExactInputSingleParams {
        address tokenIn;
        address tokenOut;
        uint24 fee;
        address recipient;
        uint256 deadline;
        uint256 amountIn;
        uint256 amountOutMinimum;
        uint160 sqrtPriceLimitX96;
    }

    function exactInputSingle(ExactInputSingleParams calldata params)
        external
        payable
        returns (uint256 amountOut);
}

/// @title Balancer Flash Loan Receiver for Arbitrage Execution (SECURE VERSION)
/// @notice Receives flash loans from Balancer and executes arbitrage paths with comprehensive security
/// @dev FIXED: All critical security vulnerabilities from audit
contract FlashLoanReceiverSecure is ReentrancyGuard {
    using SafeERC20 for IERC20;

    address public owner;
    IBalancerVault public immutable vault;

    // SECURITY FIX #4: Flash loan initiation flag
    bool private _flashLoanActive;

    // SECURITY FIX #5: Maximum path length to prevent gas limit DoS
    uint256 public constant MAX_PATH_LENGTH = 5;

    // SECURITY FIX #1: Maximum slippage in basis points (0.5% = 50 bps)
    uint256 public constant MAX_SLIPPAGE_BPS = 50;
    uint256 public constant BASIS_POINTS = 10000;

    struct ArbitragePath {
        address[] tokens;        // Token path
        address[] exchanges;     // DEX addresses
        uint24[] fees;          // Uniswap V3 fees (0 for V2)
        bool[] isV3;            // true if Uniswap V3, false if V2
        uint256 minProfit;      // Minimum profit required
        uint256 slippageBps;    // Slippage tolerance in basis points
    }

    event ArbitrageExecuted(
        address indexed initiator,
        uint256 profit,
        uint8 pathLength
    );

    event FlashLoanInitiated(
        address indexed token,
        uint256 amount
    );

    event SlippageProtectionTriggered(
        uint256 expectedMin,
        uint256 actualReceived
    );

    modifier onlyOwner() {
        require(msg.sender == owner, "Not owner");
        _;
    }

    constructor(address _vault) {
        require(_vault != address(0), "Invalid vault address");
        owner = msg.sender;
        vault = IBalancerVault(_vault);
    }

    /// @notice Execute arbitrage using Balancer flash loan
    /// @param tokens Token addresses to borrow
    /// @param amounts Amounts to borrow
    /// @param path Encoded arbitrage path
    /// @dev SECURITY FIX #2: Added nonReentrant modifier
    /// @dev SECURITY FIX #4: Sets flash loan initiation flag
    function executeArbitrage(
        IERC20[] memory tokens,
        uint256[] memory amounts,
        bytes memory path
    ) external onlyOwner nonReentrant {
        require(tokens.length > 0, "No tokens specified");
        require(tokens.length == amounts.length, "Array length mismatch");
        require(!_flashLoanActive, "Flash loan already active");

        // SECURITY FIX #4: Set flash loan active flag
        _flashLoanActive = true;

        emit FlashLoanInitiated(address(tokens[0]), amounts[0]);

        // Request flash loan from Balancer Vault
        vault.flashLoan(address(this), tokens, amounts, path);

        // SECURITY FIX #4: Clear flash loan active flag
        _flashLoanActive = false;
    }

    /// @notice Callback from Balancer Vault after flash loan
    /// @param tokens Tokens received
    /// @param amounts Amounts received
    /// @param feeAmounts Fee amounts (always 0 for Balancer)
    /// @param userData Encoded arbitrage path
    /// @dev SECURITY FIX #2: Added nonReentrant modifier
    /// @dev SECURITY FIX #4: Validates flash loan was initiated by this contract
    function receiveFlashLoan(
        IERC20[] memory tokens,
        uint256[] memory amounts,
        uint256[] memory feeAmounts,
        bytes memory userData
    ) external nonReentrant {
        // Validate callback is from vault
        require(msg.sender == address(vault), "Only vault can call");

        // SECURITY FIX #4: Validate flash loan was initiated by this contract
        require(_flashLoanActive, "Flash loan not initiated by contract");

        // Decode arbitrage path
        ArbitragePath memory path = abi.decode(userData, (ArbitragePath));

        // SECURITY FIX #5: Validate path length
        require(path.tokens.length >= 2, "Path too short");
        require(path.tokens.length <= MAX_PATH_LENGTH, "Path exceeds maximum length");
        require(path.tokens.length == path.exchanges.length + 1, "Invalid path structure");

        // SECURITY FIX #1: Validate slippage tolerance
        require(path.slippageBps <= MAX_SLIPPAGE_BPS, "Slippage tolerance too high");

        // Execute arbitrage swaps
        uint256 currentAmount = amounts[0];
        address currentToken = address(tokens[0]);

        for (uint256 i = 0; i < path.tokens.length - 1; i++) {
            address tokenIn = path.tokens[i];
            address tokenOut = path.tokens[i + 1];
            address exchange = path.exchanges[i];

            require(exchange != address(0), "Invalid exchange address");

            // SECURITY FIX #3: Use SafeERC20 for approvals (forceApprove in OZ v5)
            IERC20(tokenIn).forceApprove(exchange, currentAmount);

            if (path.isV3[i]) {
                // Uniswap V3 swap
                currentAmount = _executeV3Swap(
                    tokenIn,
                    tokenOut,
                    exchange,
                    currentAmount,
                    path.fees[i],
                    path.slippageBps
                );
            } else {
                // Uniswap V2 swap
                currentAmount = _executeV2Swap(
                    tokenIn,
                    tokenOut,
                    exchange,
                    currentAmount,
                    path.slippageBps
                );
            }

            currentToken = tokenOut;
        }

        // Calculate profit
        uint256 loanAmount = amounts[0];
        uint256 totalRepayment = loanAmount + feeAmounts[0]; // feeAmounts is 0 for Balancer
        require(currentAmount >= totalRepayment, "Insufficient funds for repayment");

        uint256 profit = currentAmount - totalRepayment;
        require(profit >= path.minProfit, "Profit below minimum threshold");

        // SECURITY FIX #3: Use SafeERC20 for repayment
        for (uint256 i = 0; i < tokens.length; i++) {
            tokens[i].safeTransfer(address(vault), amounts[i] + feeAmounts[i]);
        }

        // Emit event
        emit ArbitrageExecuted(owner, profit, uint8(path.tokens.length));

        // Profit remains in contract for withdrawal
    }

    /// @notice Execute Uniswap V3 swap with slippage protection
    /// @dev SECURITY FIX #1: Implements proper slippage protection
    function _executeV3Swap(
        address tokenIn,
        address tokenOut,
        address exchange,
        uint256 amountIn,
        uint24 fee,
        uint256 slippageBps
    ) private returns (uint256 amountOut) {
        // SECURITY FIX #1: Calculate minimum acceptable output
        uint256 minAmountOut = _calculateMinAmountOut(amountIn, slippageBps);

        IUniswapV3Router.ExactInputSingleParams memory params = IUniswapV3Router
            .ExactInputSingleParams({
                tokenIn: tokenIn,
                tokenOut: tokenOut,
                fee: fee,
                recipient: address(this),
                deadline: block.timestamp + 300, // 5 minute deadline
                amountIn: amountIn,
                amountOutMinimum: minAmountOut, // SECURITY FIX #1: NOT 0!
                sqrtPriceLimitX96: 0
            });

        amountOut = IUniswapV3Router(exchange).exactInputSingle(params);

        // Validate output meets minimum
        require(amountOut >= minAmountOut, "Slippage tolerance exceeded");

        emit SlippageProtectionTriggered(minAmountOut, amountOut);
    }

    /// @notice Execute Uniswap V2 swap with slippage protection
    /// @dev SECURITY FIX #1: Implements proper slippage protection
    function _executeV2Swap(
        address tokenIn,
        address tokenOut,
        address exchange,
        uint256 amountIn,
        uint256 slippageBps
    ) private returns (uint256 amountOut) {
        // SECURITY FIX #1: Calculate minimum acceptable output
        uint256 minAmountOut = _calculateMinAmountOut(amountIn, slippageBps);

        address[] memory swapPath = new address[](2);
        swapPath[0] = tokenIn;
        swapPath[1] = tokenOut;

        uint256[] memory swapAmounts = IUniswapV2Router(exchange)
            .swapExactTokensForTokens(
                amountIn,
                minAmountOut, // SECURITY FIX #1: NOT 0!
                swapPath,
                address(this),
                block.timestamp + 300 // 5 minute deadline
            );

        amountOut = swapAmounts[swapAmounts.length - 1];

        // Validate output meets minimum
        require(amountOut >= minAmountOut, "Slippage tolerance exceeded");

        emit SlippageProtectionTriggered(minAmountOut, amountOut);
    }

    /// @notice Calculate minimum acceptable output amount based on slippage tolerance
    /// @dev SECURITY FIX #1: Helper function for slippage calculations
    /// @param amountIn Input amount
    /// @param slippageBps Slippage tolerance in basis points
    /// @return minAmount Minimum acceptable output amount
    function _calculateMinAmountOut(
        uint256 amountIn,
        uint256 slippageBps
    ) private pure returns (uint256 minAmount) {
        require(slippageBps <= MAX_SLIPPAGE_BPS, "Slippage too high");
        // Calculate: amountIn * (10000 - slippageBps) / 10000
        minAmount = (amountIn * (BASIS_POINTS - slippageBps)) / BASIS_POINTS;
    }

    /// @notice Withdraw profits
    /// @param token Token to withdraw
    /// @param amount Amount to withdraw
    /// @dev SECURITY FIX #3: Use SafeERC20 for transfers
    function withdrawProfit(address token, uint256 amount) external onlyOwner nonReentrant {
        require(token != address(0), "Invalid token address");
        require(amount > 0, "Amount must be positive");

        uint256 balance = IERC20(token).balanceOf(address(this));
        require(balance >= amount, "Insufficient balance");

        // SECURITY FIX #3: Use SafeERC20
        IERC20(token).safeTransfer(owner, amount);
    }

    /// @notice Emergency withdraw
    /// @param token Token address (or 0x0 for ETH)
    /// @dev SECURITY FIX #3: Use SafeERC20 for transfers
    function emergencyWithdraw(address token) external onlyOwner nonReentrant {
        if (token == address(0)) {
            // Withdraw ETH
            uint256 balance = address(this).balance;
            require(balance > 0, "No ETH to withdraw");
            payable(owner).transfer(balance);
        } else {
            // Withdraw ERC20
            uint256 balance = IERC20(token).balanceOf(address(this));
            require(balance > 0, "No tokens to withdraw");

            // SECURITY FIX #3: Use SafeERC20
            IERC20(token).safeTransfer(owner, balance);
        }
    }

    /// @notice Transfer ownership
    /// @param newOwner New owner address
    function transferOwnership(address newOwner) external onlyOwner {
        require(newOwner != address(0), "Invalid new owner");
        require(newOwner != owner, "Already owner");
        owner = newOwner;
    }

    /// @notice Get contract balance for a token
    /// @param token Token address
    /// @return balance Token balance of this contract
    function getBalance(address token) external view returns (uint256 balance) {
        balance = IERC20(token).balanceOf(address(this));
    }

    receive() external payable {}
}