smom-dbis-138/contracts/flash/AaveUniV2CwStableRebalanceFlashReceiver.sol

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

import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import {
    IAavePoolLike,
    IAaveFlashLoanReceiver
} from "./AaveQuotePushFlashReceiver.sol";

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

    function removeLiquidity(
        address tokenA,
        address tokenB,
        uint256 liquidity,
        uint256 amountAMin,
        uint256 amountBMin,
        address to,
        uint256 deadline
    ) external returns (uint256 amountA, uint256 amountB);
}

/**
 * @title AaveUniV2CwStableRebalanceFlashReceiver
 * @notice Same-block Aave V3 USDC flash loan to quote-side rebalance a skewed cW/stable UniV2 pair,
 *         remove owned LP, sell cW* back into the pair for USDC, and repay Aave.
 *
 * @dev LP tokens must be on this contract before `runRebalanceRemove` (transfer or `pullLpFrom`).
 */
contract AaveUniV2CwStableRebalanceFlashReceiver is IAaveFlashLoanReceiver, Ownable {
    using SafeERC20 for IERC20;

    address public immutable pool;

    struct RebalanceRemoveParams {
        address router;
        address pair;
        address cwToken;
        address stableToken;
        uint256 lpAmount;
        uint256 rebalanceStableIn;
        uint256 minCwFromRebalance;
        uint256 minStableFromRemove;
        uint256 minCwFromRemove;
        uint256 cwToSellForRepay;
        uint256 minStableFromRepaySwap;
        address recipient;
    }

    error UntrustedPool();
    error UntrustedInitiator();
    error BadParams();
    error InsufficientToRepay();
    error NothingToSweep();

    event RebalanceRemoveExecuted(
        address indexed pair,
        address indexed stableToken,
        uint256 borrowedAmount,
        uint256 premium,
        uint256 lpRemoved,
        uint256 stableSurplus
    );
    event TokenSwept(address indexed token, address indexed to, uint256 amount);

    constructor(address pool_, address initialOwner) Ownable(initialOwner) {
        if (pool_ == address(0) || initialOwner == address(0)) revert BadParams();
        pool = pool_;
    }

    function pullLpFrom(address pair, address from, uint256 amount) external onlyOwner {
        if (pair == address(0) || from == address(0) || amount == 0) revert BadParams();
        IERC20(pair).safeTransferFrom(from, address(this), amount);
    }

    function runRebalanceRemove(address stableToken, uint256 amount, RebalanceRemoveParams calldata params)
        external
        onlyOwner
    {
        if (stableToken == address(0) || amount == 0) revert BadParams();
        if (params.lpAmount == 0 || params.rebalanceStableIn == 0 || params.recipient == address(0)) {
            revert BadParams();
        }
        if (IERC20(params.pair).balanceOf(address(this)) < params.lpAmount) revert BadParams();

        address[] memory assets = new address[](1);
        uint256[] memory amts = new uint256[](1);
        uint256[] memory modes = new uint256[](1);
        assets[0] = stableToken;
        amts[0] = amount;
        modes[0] = 0;
        IAavePoolLike(pool).flashLoan(
            address(this), assets, amts, modes, address(this), abi.encode(address(this), params), 0
        );
    }

    function executeOperation(
        address[] calldata assets,
        uint256[] calldata amounts,
        uint256[] calldata premiums,
        address initiator,
        bytes calldata params
    ) external override returns (bool) {
        if (msg.sender != pool) revert UntrustedPool();
        if (assets.length != 1 || amounts.length != 1 || premiums.length != 1) revert BadParams();
        _executeRebalanceRemove(assets[0], amounts[0], premiums[0], initiator, params);
        return true;
    }

    function sweepToken(address token, address to, uint256 amount) external onlyOwner {
        if (token == address(0) || to == address(0) || amount == 0) revert BadParams();
        IERC20(token).safeTransfer(to, amount);
        emit TokenSwept(token, to, amount);
    }

    function _executeRebalanceRemove(
        address stableToken,
        uint256 amount,
        uint256 premium,
        address initiator,
        bytes calldata params
    ) internal {
        (address expectedInitiator, RebalanceRemoveParams memory p) =
            abi.decode(params, (address, RebalanceRemoveParams));
        if (initiator != expectedInitiator) revert UntrustedInitiator();
        if (
            p.router == address(0) || p.pair == address(0) || p.cwToken == address(0)
                || p.stableToken != stableToken
        ) revert BadParams();
        if (p.rebalanceStableIn != amount) revert BadParams();

        uint256 deadline = block.timestamp;
        address[] memory stableToCw = _path(p.stableToken, p.cwToken);
        address[] memory cwToStable = _path(p.cwToken, p.stableToken);

        IERC20(p.stableToken).forceApprove(p.router, p.rebalanceStableIn);
        IUniswapV2Router02Minimal(p.router).swapExactTokensForTokens(
            p.rebalanceStableIn, p.minCwFromRebalance, stableToCw, address(this), deadline
        );

        IERC20(p.pair).forceApprove(p.router, p.lpAmount);
        (address tokenA, address tokenB) = _sortTokens(p.cwToken, p.stableToken);
        (uint256 minA, uint256 minB) = p.cwToken < p.stableToken
            ? (p.minCwFromRemove, p.minStableFromRemove)
            : (p.minStableFromRemove, p.minCwFromRemove);
        IUniswapV2Router02Minimal(p.router).removeLiquidity(
            tokenA, tokenB, p.lpAmount, minA, minB, address(this), deadline
        );

        uint256 repayNeed = amount + premium;
        uint256 stableBal = IERC20(p.stableToken).balanceOf(address(this));
        if (stableBal < repayNeed && p.cwToSellForRepay > 0) {
            IERC20(p.cwToken).forceApprove(p.router, p.cwToSellForRepay);
            IUniswapV2Router02Minimal(p.router).swapExactTokensForTokens(
                p.cwToSellForRepay, p.minStableFromRepaySwap, cwToStable, address(this), deadline
            );
            stableBal = IERC20(p.stableToken).balanceOf(address(this));
        }

        if (stableBal < repayNeed) revert InsufficientToRepay();
        uint256 surplus = stableBal - repayNeed;
        IERC20(p.stableToken).forceApprove(pool, repayNeed);
        if (surplus > 0) {
            IERC20(p.stableToken).safeTransfer(p.recipient, surplus);
        }

        uint256 cwLeft = IERC20(p.cwToken).balanceOf(address(this));
        if (cwLeft > 0) {
            IERC20(p.cwToken).safeTransfer(p.recipient, cwLeft);
        }

        emit RebalanceRemoveExecuted(p.pair, p.stableToken, amount, premium, p.lpAmount, surplus);
    }

    function _path(address tokenIn, address tokenOut) internal pure returns (address[] memory path) {
        path = new address[](2);
        path[0] = tokenIn;
        path[1] = tokenOut;
    }

    function _sortTokens(address tokenA, address tokenB) internal pure returns (address, address) {
        return tokenA < tokenB ? (tokenA, tokenB) : (tokenB, tokenA);
    }
}