mev-beta/tests/contracts/FlashLoanReceiverSecure.t.sol

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

import "forge-std/Test.sol";
import "../../contracts/balancer/FlashLoanReceiverSecure.sol";
import "@openzeppelin/contracts/token/ERC20/ERC20.sol";

/// @title Mock ERC20 Token for Testing
contract MockERC20 is ERC20 {
    constructor() ERC20("Mock Token", "MOCK") {
        _mint(msg.sender, 1000000 * 10**18);
    }

    function mint(address to, uint256 amount) external {
        _mint(to, amount);
    }
}

/// @title Mock Balancer Vault for Testing
contract MockBalancerVault {
    function flashLoan(
        address recipient,
        address[] memory tokens,
        uint256[] memory amounts,
        bytes memory userData
    ) external {
        // Transfer tokens to recipient
        for (uint i = 0; i < tokens.length; i++) {
            IERC20(tokens[i]).transfer(recipient, amounts[i]);
        }

        // Call receiveFlashLoan
        uint256[] memory feeAmounts = new uint256[](tokens.length);
        for (uint i = 0; i < tokens.length; i++) {
            feeAmounts[i] = 0; // Balancer has 0 fees
        }

        IFlashLoanRecipient(recipient).receiveFlashLoan(
            _convertToIERC20Array(tokens),
            amounts,
            feeAmounts,
            userData
        );

        // Verify repayment
        for (uint i = 0; i < tokens.length; i++) {
            require(
                IERC20(tokens[i]).balanceOf(address(this)) >= amounts[i],
                "Flash loan not repaid"
            );
        }
    }

    function _convertToIERC20Array(address[] memory tokens)
        private
        pure
        returns (IERC20[] memory)
    {
        IERC20[] memory ierc20Array = new IERC20[](tokens.length);
        for (uint i = 0; i < tokens.length; i++) {
            ierc20Array[i] = IERC20(tokens[i]);
        }
        return ierc20Array;
    }
}

interface IFlashLoanRecipient {
    function receiveFlashLoan(
        IERC20[] memory tokens,
        uint256[] memory amounts,
        uint256[] memory feeAmounts,
        bytes memory userData
    ) external;
}

/// @title Comprehensive Test Suite for FlashLoanReceiverSecure
contract FlashLoanReceiverSecureTest is Test {
    FlashLoanReceiverSecure public flashLoan;
    MockBalancerVault public vault;
    MockERC20 public token1;
    MockERC20 public token2;

    address public owner;
    address public user1;
    address public attacker;

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

    // Allow test contract to receive ETH
    receive() external payable {}

    function setUp() public {
        owner = address(this);
        user1 = address(0x1);
        attacker = address(0x2);

        // Deploy mock contracts
        vault = new MockBalancerVault();
        token1 = new MockERC20();
        token2 = new MockERC20();

        // Deploy FlashLoanReceiverSecure
        flashLoan = new FlashLoanReceiverSecure(address(vault));

        // Fund vault with tokens for flash loans
        token1.transfer(address(vault), 100000 * 10**18);
        token2.transfer(address(vault), 100000 * 10**18);

        // Fund flash loan contract with some tokens/ETH for testing withdrawals
        token1.transfer(address(flashLoan), 1000 * 10**18);
        token2.transfer(address(flashLoan), 500 * 10**18);
        vm.deal(address(flashLoan), 10 ether);
    }

    /// ============================================
    /// WITHDRAWAL TESTS
    /// ============================================

    function testWithdrawProfit_Success() public {
        uint256 withdrawAmount = 100 * 10**18;
        uint256 initialBalance = token1.balanceOf(owner);

        flashLoan.withdrawProfit(address(token1), withdrawAmount);

        assertEq(token1.balanceOf(owner), initialBalance + withdrawAmount);
        assertEq(token1.balanceOf(address(flashLoan)), 900 * 10**18);
    }

    function testWithdrawProfit_MultipleWithdrawals() public {
        uint256 firstWithdraw = 200 * 10**18;
        uint256 secondWithdraw = 300 * 10**18;

        flashLoan.withdrawProfit(address(token1), firstWithdraw);
        flashLoan.withdrawProfit(address(token1), secondWithdraw);

        assertEq(token1.balanceOf(address(flashLoan)), 500 * 10**18);
    }

    function testWithdrawProfit_RevertInvalidToken() public {
        vm.expectRevert("Invalid token address");
        flashLoan.withdrawProfit(address(0), 100);
    }

    function testWithdrawProfit_RevertZeroAmount() public {
        vm.expectRevert("Amount must be positive");
        flashLoan.withdrawProfit(address(token1), 0);
    }

    function testWithdrawProfit_RevertInsufficientBalance() public {
        vm.expectRevert("Insufficient balance");
        flashLoan.withdrawProfit(address(token1), 10000 * 10**18); // More than balance
    }

    function testWithdrawProfit_RevertNotOwner() public {
        vm.prank(attacker);
        vm.expectRevert("Not owner");
        flashLoan.withdrawProfit(address(token1), 100 * 10**18);
    }

    /// ============================================
    /// EMERGENCY WITHDRAW - TOKEN TESTS
    /// ============================================

    function testEmergencyWithdraw_Token_Success() public {
        uint256 initialBalance = token1.balanceOf(owner);
        uint256 contractBalance = token1.balanceOf(address(flashLoan));

        flashLoan.emergencyWithdraw(address(token1));

        assertEq(token1.balanceOf(owner), initialBalance + contractBalance);
        assertEq(token1.balanceOf(address(flashLoan)), 0);
    }

    function testEmergencyWithdraw_MultipleTokens() public {
        uint256 token1Balance = token1.balanceOf(address(flashLoan));
        uint256 token2Balance = token2.balanceOf(address(flashLoan));

        flashLoan.emergencyWithdraw(address(token1));
        flashLoan.emergencyWithdraw(address(token2));

        assertEq(token1.balanceOf(address(flashLoan)), 0);
        assertEq(token2.balanceOf(address(flashLoan)), 0);
        assertGt(token1.balanceOf(owner), token1Balance);
        assertGt(token2.balanceOf(owner), token2Balance);
    }

    function testEmergencyWithdraw_Token_RevertNoBalance() public {
        // First withdraw all tokens
        flashLoan.emergencyWithdraw(address(token1));

        // Try to withdraw again
        vm.expectRevert("No tokens to withdraw");
        flashLoan.emergencyWithdraw(address(token1));
    }

    function testEmergencyWithdraw_Token_RevertNotOwner() public {
        vm.prank(attacker);
        vm.expectRevert("Not owner");
        flashLoan.emergencyWithdraw(address(token1));
    }

    /// ============================================
    /// EMERGENCY WITHDRAW - ETH TESTS
    /// ============================================

    function testEmergencyWithdraw_ETH_Success() public {
        uint256 initialBalance = owner.balance;
        uint256 contractBalance = address(flashLoan).balance;

        flashLoan.emergencyWithdraw(address(0));

        assertEq(owner.balance, initialBalance + contractBalance);
        assertEq(address(flashLoan).balance, 0);
    }

    function testEmergencyWithdraw_ETH_RevertNoBalance() public {
        // First withdraw all ETH
        flashLoan.emergencyWithdraw(address(0));

        // Try to withdraw again
        vm.expectRevert("No ETH to withdraw");
        flashLoan.emergencyWithdraw(address(0));
    }

    function testEmergencyWithdraw_ETH_RevertNotOwner() public {
        vm.prank(attacker);
        vm.expectRevert("Not owner");
        flashLoan.emergencyWithdraw(address(0));
    }

    /// ============================================
    /// MIXED WITHDRAWAL TESTS
    /// ============================================

    function testWithdrawProfit_ThenEmergencyWithdraw() public {
        // First do controlled withdrawal
        flashLoan.withdrawProfit(address(token1), 500 * 10**18);
        assertEq(token1.balanceOf(address(flashLoan)), 500 * 10**18);

        // Then emergency withdraw remaining
        flashLoan.emergencyWithdraw(address(token1));
        assertEq(token1.balanceOf(address(flashLoan)), 0);
    }

    function testEmergencyWithdraw_AllAssets() public {
        uint256 initialETH = owner.balance;
        uint256 initialToken1 = token1.balanceOf(owner);
        uint256 initialToken2 = token2.balanceOf(owner);

        flashLoan.emergencyWithdraw(address(0)); // ETH
        flashLoan.emergencyWithdraw(address(token1));
        flashLoan.emergencyWithdraw(address(token2));

        assertEq(address(flashLoan).balance, 0);
        assertEq(token1.balanceOf(address(flashLoan)), 0);
        assertEq(token2.balanceOf(address(flashLoan)), 0);

        assertGt(owner.balance, initialETH);
        assertGt(token1.balanceOf(owner), initialToken1);
        assertGt(token2.balanceOf(owner), initialToken2);
    }

    /// ============================================
    /// ACCESS CONTROL TESTS
    /// ============================================

    function testTransferOwnership_Success() public {
        address newOwner = address(0x999);

        flashLoan.transferOwnership(newOwner);
        assertEq(flashLoan.owner(), newOwner);

        // Old owner can't withdraw anymore
        vm.expectRevert("Not owner");
        flashLoan.withdrawProfit(address(token1), 100);

        // New owner can withdraw
        vm.prank(newOwner);
        flashLoan.withdrawProfit(address(token1), 100 * 10**18);
    }

    function testTransferOwnership_RevertInvalidAddress() public {
        vm.expectRevert("Invalid new owner");
        flashLoan.transferOwnership(address(0));
    }

    function testTransferOwnership_RevertSameOwner() public {
        vm.expectRevert("Already owner");
        flashLoan.transferOwnership(owner);
    }

    function testTransferOwnership_RevertNotOwner() public {
        vm.prank(attacker);
        vm.expectRevert("Not owner");
        flashLoan.transferOwnership(attacker);
    }

    /// ============================================
    /// REENTRANCY PROTECTION TESTS
    /// ============================================

    function testWithdrawProfit_ReentrancyProtection() public {
        // Note: Full reentrancy testing requires malicious contract
        // This is a basic check that nonReentrant modifier is present

        // Attempting to call withdrawProfit during its execution should fail
        // This would require a custom malicious ERC20 token that calls back
        // For now, we verify the modifier is working by checking sequential calls work

        flashLoan.withdrawProfit(address(token1), 100 * 10**18);
        flashLoan.withdrawProfit(address(token1), 100 * 10**18);
        // If reentrancy guard broken, second call would fail or behave incorrectly
        assertEq(token1.balanceOf(address(flashLoan)), 800 * 10**18);
    }

    /// ============================================
    /// VIEW FUNCTION TESTS
    /// ============================================

    function testGetBalance_Token() public {
        uint256 balance = flashLoan.getBalance(address(token1));
        assertEq(balance, 1000 * 10**18);
    }

    function testGetBalance_ETH() public {
        // For ETH, getBalance uses address(this).balance internally
        assertEq(address(flashLoan).balance, 10 ether);
    }

    function testOwner() public {
        assertEq(flashLoan.owner(), owner);
    }

    function testVault() public {
        assertEq(address(flashLoan.vault()), address(vault));
    }

    function testConstants() public {
        assertEq(flashLoan.MAX_SLIPPAGE_BPS(), 50);
        assertEq(flashLoan.MAX_PATH_LENGTH(), 5);
        assertEq(flashLoan.BASIS_POINTS(), 10000);
    }

    /// ============================================
    /// EDGE CASES
    /// ============================================

    function testWithdrawProfit_ExactBalance() public {
        uint256 exactBalance = token1.balanceOf(address(flashLoan));

        flashLoan.withdrawProfit(address(token1), exactBalance);
        assertEq(token1.balanceOf(address(flashLoan)), 0);
    }

    function testWithdrawProfit_MultipleTokenTypes() public {
        flashLoan.withdrawProfit(address(token1), 100 * 10**18);
        flashLoan.withdrawProfit(address(token2), 50 * 10**18);

        assertEq(token1.balanceOf(address(flashLoan)), 900 * 10**18);
        assertEq(token2.balanceOf(address(flashLoan)), 450 * 10**18);
    }

    function testEmergencyWithdraw_AfterReceivingETH() public {
        // Send more ETH to contract
        vm.deal(address(flashLoan), 20 ether);

        uint256 initialBalance = owner.balance;
        flashLoan.emergencyWithdraw(address(0));

        assertEq(owner.balance, initialBalance + 20 ether);
    }

    /// ============================================
    /// RECEIVE ETH TESTS
    /// ============================================

    function testReceiveETH() public {
        uint256 initialBalance = address(flashLoan).balance;

        (bool success,) = address(flashLoan).call{value: 1 ether}("");
        assertTrue(success);

        assertEq(address(flashLoan).balance, initialBalance + 1 ether);
    }

    function testReceiveETH_CanWithdrawAfter() public {
        // Send ETH
        payable(address(flashLoan)).transfer(5 ether);

        // Withdraw
        uint256 initialBalance = owner.balance;
        flashLoan.emergencyWithdraw(address(0));

        assertEq(owner.balance, initialBalance + 15 ether); // 10 initial + 5 sent
    }

    /// ============================================
    /// FUZZ TESTS
    /// ============================================

    function testFuzz_WithdrawProfit(uint256 amount) public {
        uint256 contractBalance = token1.balanceOf(address(flashLoan));
        vm.assume(amount > 0 && amount <= contractBalance);

        flashLoan.withdrawProfit(address(token1), amount);
        assertEq(token1.balanceOf(address(flashLoan)), contractBalance - amount);
    }

    function testFuzz_EmergencyWithdrawETH(uint96 ethAmount) public {
        vm.assume(ethAmount > 0);
        vm.deal(address(flashLoan), ethAmount);

        uint256 initialBalance = owner.balance;
        flashLoan.emergencyWithdraw(address(0));

        assertEq(owner.balance, initialBalance + ethAmount);
        assertEq(address(flashLoan).balance, 0);
    }

    /// ============================================
    /// GAS OPTIMIZATION TESTS
    /// ============================================

    function testGas_WithdrawProfit() public {
        uint256 gasBefore = gasleft();
        flashLoan.withdrawProfit(address(token1), 100 * 10**18);
        uint256 gasUsed = gasBefore - gasleft();

        // Should use less than 100k gas
        assertLt(gasUsed, 100000);
    }

    function testGas_EmergencyWithdrawToken() public {
        uint256 gasBefore = gasleft();
        flashLoan.emergencyWithdraw(address(token1));
        uint256 gasUsed = gasBefore - gasleft();

        // Should use less than 100k gas
        assertLt(gasUsed, 100000);
    }

    function testGas_EmergencyWithdrawETH() public {
        uint256 gasBefore = gasleft();
        flashLoan.emergencyWithdraw(address(0));
        uint256 gasUsed = gasBefore - gasleft();

        // Should use less than 50k gas
        assertLt(gasUsed, 50000);
    }
}