smom-dbis-138/test/flash/AaveQuotePushCollateralHooks.t.sol

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

import {Test} from "forge-std/Test.sol";
import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import {
    AaveQuotePushFlashReceiver,
    IAaveDODOQuotePushSwapExactIn,
    IAaveExternalUnwinder,
    IAavePoolLike
} from "../../contracts/flash/AaveQuotePushFlashReceiver.sol";

contract MockQuoteToken is ERC20 {
    constructor() ERC20("Mock Quote", "MQ") {}

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

contract MockBaseToken is ERC20 {
    constructor() ERC20("Mock Base", "MB") {}

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

contract MockPmmIntegration is IAaveDODOQuotePushSwapExactIn {
    MockQuoteToken internal immutable quote;
    MockBaseToken internal immutable base;

    constructor(MockQuoteToken quote_, MockBaseToken base_) {
        quote = quote_;
        base = base_;
    }

    function swapExactIn(address, address tokenIn, uint256 amountIn, uint256 minAmountOut)
        external
        override
        returns (uint256 amountOut)
    {
        require(tokenIn == address(quote), "quote only");
        quote.transferFrom(msg.sender, address(this), amountIn);
        amountOut = amountIn;
        require(amountOut >= minAmountOut, "minOut");
        base.mint(msg.sender, amountOut);
    }
}

contract MockUnwinder is IAaveExternalUnwinder {
    MockQuoteToken internal immutable quote;
    MockBaseToken internal immutable base;

    constructor(MockQuoteToken quote_, MockBaseToken base_) {
        quote = quote_;
        base = base_;
    }

    function unwind(address tokenIn, address tokenOut, uint256 amountIn, uint256 minAmountOut, bytes calldata)
        external
        override
        returns (uint256 amountOut)
    {
        require(tokenIn == address(base), "base only");
        require(tokenOut == address(quote), "quote only");
        base.transferFrom(msg.sender, address(this), amountIn);
        amountOut = amountIn;
        require(amountOut >= minAmountOut, "min unwind");
        quote.mint(msg.sender, amountOut);
    }
}

contract MockAavePool is IAavePoolLike {
    struct ToggleCall {
        address asset;
        bool useAsCollateral;
    }

    uint256 public premiumBps = 5;
    uint256 public toggleCallCount;
    ToggleCall[] internal _toggleCalls;
    uint256 public supplyCount;

    function flashLoan(
        address receiverAddress,
        address[] calldata assets,
        uint256[] calldata amounts,
        uint256[] calldata,
        address,
        bytes calldata params,
        uint16
    ) external override {
        uint256 premium = amounts[0] * premiumBps / 10_000;
        uint256[] memory premiums = new uint256[](1);
        premiums[0] = premium;
        IERC20(assets[0]).transfer(receiverAddress, amounts[0]);
        bool ok = IAaveFlashLoanReceiver(receiverAddress).executeOperation(
            assets, amounts, premiums, receiverAddress, params
        );
        require(ok, "callback failed");
        IERC20(assets[0]).transferFrom(receiverAddress, address(this), amounts[0] + premium);
    }

    function flashLoanSimple(address, address, uint256, bytes calldata, uint16) external pure override {
        revert("use flashLoan");
    }

    function supply(address asset, uint256 amount, address onBehalfOf, uint16) external override {
        IERC20(asset).transferFrom(msg.sender, address(this), amount);
        supplyCount++;
        onBehalfOf;
    }

    function setUserUseReserveAsCollateral(address asset, bool useAsCollateral) external override {
        _toggleCalls.push(ToggleCall({asset: asset, useAsCollateral: useAsCollateral}));
        toggleCallCount++;
    }

    function toggleCallAt(uint256 index) external view returns (address asset, bool useAsCollateral) {
        ToggleCall memory t = _toggleCalls[index];
        return (t.asset, t.useAsCollateral);
    }
}

interface IAaveFlashLoanReceiver {
    function executeOperation(
        address[] calldata assets,
        uint256[] calldata amounts,
        uint256[] calldata premiums,
        address initiator,
        bytes calldata params
    ) external returns (bool);
}

contract AaveQuotePushCollateralHooksTest is Test {
    MockQuoteToken internal quote;
    MockBaseToken internal base;
    MockPmmIntegration internal integration;
    MockUnwinder internal unwinder;
    MockAavePool internal pool;
    AaveQuotePushFlashReceiver internal receiver;

    address internal constant WETH = address(0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2);

    function setUp() public {
        quote = new MockQuoteToken();
        base = new MockBaseToken();
        integration = new MockPmmIntegration(quote, base);
        unwinder = new MockUnwinder(quote, base);
        pool = new MockAavePool();
        receiver = new AaveQuotePushFlashReceiver(address(pool), address(this));
        quote.mint(address(pool), 10_000_000);
        quote.mint(address(receiver), 1_000_000);
    }

    function testCollateralToggleRunsInFlashCallbackBeforeAndAfterSwap() public {
        AaveQuotePushFlashReceiver.CollateralToggleStep[] memory beforeSwap =
            new AaveQuotePushFlashReceiver.CollateralToggleStep[](1);
        beforeSwap[0] = AaveQuotePushFlashReceiver.CollateralToggleStep({asset: WETH, useAsCollateral: false});

        AaveQuotePushFlashReceiver.CollateralToggleStep[] memory afterSwap =
            new AaveQuotePushFlashReceiver.CollateralToggleStep[](1);
        afterSwap[0] = AaveQuotePushFlashReceiver.CollateralToggleStep({asset: WETH, useAsCollateral: true});

        AaveQuotePushFlashReceiver.QuotePushParams memory p = AaveQuotePushFlashReceiver.QuotePushParams({
            integration: address(integration),
            pmmPool: address(0xBEEF),
            baseToken: address(base),
            externalUnwinder: address(unwinder),
            minOutPmm: 1,
            minOutUnwind: 1,
            unwindData: "",
            atomicBridge: AaveQuotePushFlashReceiver.AtomicBridgeParams({
                coordinator: address(0),
                sourceChain: 0,
                destinationChain: 0,
                destinationAsset: address(0),
                bridgeAmount: 0,
                minDestinationAmount: 0,
                destinationRecipient: address(0),
                destinationDeadline: 0,
                routeId: bytes32(0),
                settlementMode: bytes32(0),
                submitCommitment: false
            }),
            collateral: AaveQuotePushFlashReceiver.CollateralParams({
                supplyBeforeSwap: new AaveQuotePushFlashReceiver.CollateralSupplyStep[](0),
                toggleBeforeSwap: beforeSwap,
                toggleAfterSwap: afterSwap,
                toggleBeforeUnwind: new AaveQuotePushFlashReceiver.CollateralToggleStep[](0)
            })
        });

        uint256 borrow = 1_000_000;
        receiver.flashQuotePush(address(quote), borrow, p);

        assertEq(pool.toggleCallCount(), 2, "two toggles");
        (address asset0, bool use0) = pool.toggleCallAt(0);
        (address asset1, bool use1) = pool.toggleCallAt(1);
        assertEq(asset0, WETH);
        assertFalse(use0);
        assertEq(asset1, WETH);
        assertTrue(use1);
    }

    function testCollateralSupplyBeforeSwapUsesReceiverBalance() public {
        quote.mint(address(receiver), 500_000);

        AaveQuotePushFlashReceiver.CollateralSupplyStep[] memory supplies =
            new AaveQuotePushFlashReceiver.CollateralSupplyStep[](1);
        supplies[0] = AaveQuotePushFlashReceiver.CollateralSupplyStep({asset: address(quote), amount: 100_000});

        AaveQuotePushFlashReceiver.QuotePushParams memory p = AaveQuotePushFlashReceiver.QuotePushParams({
            integration: address(integration),
            pmmPool: address(0xBEEF),
            baseToken: address(base),
            externalUnwinder: address(unwinder),
            minOutPmm: 1,
            minOutUnwind: 1,
            unwindData: "",
            atomicBridge: AaveQuotePushFlashReceiver.AtomicBridgeParams({
                coordinator: address(0),
                sourceChain: 0,
                destinationChain: 0,
                destinationAsset: address(0),
                bridgeAmount: 0,
                minDestinationAmount: 0,
                destinationRecipient: address(0),
                destinationDeadline: 0,
                routeId: bytes32(0),
                settlementMode: bytes32(0),
                submitCommitment: false
            }),
            collateral: AaveQuotePushFlashReceiver.CollateralParams({
                supplyBeforeSwap: supplies,
                toggleBeforeSwap: new AaveQuotePushFlashReceiver.CollateralToggleStep[](0),
                toggleAfterSwap: new AaveQuotePushFlashReceiver.CollateralToggleStep[](0),
                toggleBeforeUnwind: new AaveQuotePushFlashReceiver.CollateralToggleStep[](0)
            })
        });

        receiver.flashQuotePush(address(quote), 1_000_000, p);
        assertEq(pool.supplyCount(), 1, "supply called once");
    }
}