DexAggregator

// SPDX-License-Identifier: MIT
pragma solidity 0.8.28;

import {Ownable2Step} from "openzeppelin-solidity/contracts/access/Ownable2Step.sol";
import {Pausable} from "openzeppelin-solidity/contracts/security/Pausable.sol";
import {Address} from "openzeppelin-solidity/contracts/utils/Address.sol";
import {Math} from "openzeppelin-solidity/contracts/utils/math/Math.sol";
import {IDexAggregator} from "./interfaces/IDexAggregator.sol";
import {LibAsset} from "./libraries/LibAsset.sol";
import {LibRouter, SwapData} from "./libraries/LibRouter.sol";

error InvalidCall();
error InvalidCaller();
error InvalidFundManager();
error InsufficientAmountOut();
error InvalidNativeAmount();

contract DexAggregator is IDexAggregator, Ownable2Step, Pausable {
    using LibAsset for address;

    mapping(address => bool) public internalCallers;
    mapping(address => bool) public fundsManagers;
    address public coreAddress;
    address public weth;

    modifier onlyInternalCaller() {
        if (!internalCallers[msg.sender]) {
            revert InvalidCaller();
        }
        _;
    }

    modifier onlyFundManager() {
        if (!fundsManagers[msg.sender]) {
            revert InvalidFundManager();
        }
        _;
    }

    /// @dev See {IDexAggregator-updateInternalCaller}
    function updateInternalCaller(address caller, bool value) external onlyOwner {
        internalCallers[caller] = value;

        emit UpdateInternalCaller(msg.sender, caller, value);
    }

    /// @dev See {IDexAggregator-updateFundsManager}
    function updateFundsManager(address caller, bool value) external onlyOwner {
        fundsManagers[caller] = value;

        emit UpdateFundsManager(msg.sender, caller, value);
    }

    /// @dev See {IDexAggregator-updateCoreAddress}
    function updateCoreAddress(address value) external onlyOwner {
        coreAddress = value;
    }

    /// @dev See {IDexAggregator-updateWeth}
    function updateWeth(address value) external onlyOwner {
        weth = value;
    }

    /// @dev See {IDexAggregator-updateAmountIn}
    function updateAmountIn(bytes memory data, uint256 amountIn) external pure returns (bytes memory) {
        assembly {
            mstore(add(data, 164), amountIn) // 32 + 132
        }

        return data;
    }

    /// @dev See {IDexAggregator-pause}
    function pause() public onlyOwner whenNotPaused {
        _pause();
    }

    /// @dev See {IDexAggregator-unpause}
    function unpause() public onlyOwner whenPaused {
        _unpause();
    }

    /// @dev See {IDexAggregator-transfer}
    function transfer(TransferParam calldata transferParam) external onlyFundManager {
        transferParam.assetAddress.transfer(transferParam.recipient, transferParam.amount);
    }

    /// @dev See {IDexAggregator-transferMany}
    function transferMany(TransferParam[] calldata transferParams) external onlyFundManager {
        for (uint256 i = 0; i < transferParams.length; ) {
            transferParams[i].assetAddress.transfer(transferParams[i].recipient, transferParams[i].amount);
            unchecked {
                i++;
            }
        }
    }

    /// @dev See {IDexAggregator-multicall}
    function multicall(bytes[] calldata data) external onlyOwner returns (bytes[] memory results) {
        results = new bytes[](data.length);
        for (uint256 i = 0; i < data.length; i++) {
            results[i] = Address.functionDelegateCall(address(this), data[i]);
        }
        return results;
    }

    /// @dev Verifies the backend signature for a swap operation.
    /// @param swapData The data structure containing information about the swap.
    /// @return fromAddress The address to be used for the swap operation.
    function verifyBackendSignature(SwapData memory swapData) private view returns (address fromAddress) {
        bool hasSwapFee = swapData.swapFeeReceivers.length > 0;
        uint256 messagePtr;
        uint256 messageLength = hasSwapFee ? 480 : 384;
        assembly {
            messagePtr := mload(0x40)
            mstore(0x40, add(messagePtr, messageLength))
            switch hasSwapFee
            case 1 {
                // keccak256("Swap(address router,address sender,address recipient,address fromAsset,address toAsset,uint256 deadline,uint256 amountOutMin,uint256 expectedAmountOut,bytes32 consumerId,uint256 maxRetentionBps,bool transferFromRouter,address[] swapFeeAssets,address[] swapFeeReceivers,uint256[] swapFeeAmounts)")
                mstore(messagePtr, 0x5044c48dcff133da673c91c7d371f756e221fd81c7b9f02129fcfd724b840e49)
            }
            default {
                // keccak256("Swap(address router,address sender,address recipient,address fromAsset,address toAsset,uint256 deadline,uint256 amountOutMin,uint256 expectedAmountOut,bytes32 consumerId,uint256 maxRetentionBps,bool transferFromRouter)")
                mstore(messagePtr, 0x088852e96193cda3374c3e36c04a8bd173e0c18999b51fb3f5c5ab52d7da1fe0)
            }
        }
        fromAddress = LibRouter.verifySignature(
            // keccak256(bytes("Dex Aggregator")),
            0x19994fa387947b8529d5d57c0c84137061b0be98f1cef929a4710ff4f36aed49,
            // keccak256(bytes("1")),
            0xc89efdaa54c0f20c7adf612882df0950f5a951637e0307cdcb4c672f298b8bc6,
            swapData,
            messagePtr,
            messageLength,
            2
        );
    }

    function emitSwap(EmitSwapParams memory p) private {
        emit Swap(
            p.fromAddress,
            p.toAddress,
            p.fromAssetAddress,
            p.toAssetAddress,
            p.amountIn - (p.amountIn > p.transferredAmountIn ? p.amountIn - p.transferredAmountIn : 0),
            p.amountOut,
            p.expectedAmountOut,
            p.amountIn > p.transferredAmountIn ? p.amountIn - p.transferredAmountIn : 0,
            p.transferredAmountOut > p.amountOut ? p.transferredAmountOut - p.amountOut : 0,
            p.consumerId,
            p.swapFeeAssetAddresses,
            p.swapFeeReceivers,
            p.swapFeeAmounts
        );
    }

    /// @dev Swaps tokens based on the provided swap data.
    /// @param swapData The data structure containing information about the swap operation.
    /// @param fromAddress The address initiating the swap. This address is responsible for the input assets.
    /// @return amountOut The amount of tokens or assets received after the swap.
    /// @return gasUsed The amount of gas consumed by the recorded operation.
    function swap(
        SwapData memory swapData,
        address fromAddress,
        address fromAssetAddress
    ) private returns (uint256 amountOut, uint256 gasUsed) {
        amountOut = swapData.toAssetAddress.getBalanceOf(swapData.toAddress);
        uint256 transferredAmountOut = swapData.toAssetAddress.getBalance();
        uint256 transferredAmountIn = swapData.amountIn;

        adjustFromAssetAddress(swapData);

        if (fromAssetAddress.isNative()) {
            // Already transferred and wrapped
            if (swapData.transferFromRouter) {
                swapData.fromAssetAddress.approve(coreAddress, swapData.amountIn);
            } else {
                swapData.fromAssetAddress.transfer(coreAddress, swapData.amountIn);
            }
        } else {
            if (swapData.transferFromRouter) {
                swapData.fromAssetAddress.transferFrom(fromAddress, address(this), swapData.amountIn);
                swapData.fromAssetAddress.approve(coreAddress, swapData.amountIn);
            } else {
                swapData.fromAssetAddress.transferFrom(fromAddress, coreAddress, swapData.amountIn);
            }
        }

        address currentCoreAddress = coreAddress;
        assembly {
            let success := 0
            let inputPtr := mload(0x40)
            let inputLength := add(calldataload(36), 68)
            let payloadLength := sub(inputLength, 68)
            mstore(0x40, add(inputPtr, inputLength))
            mstore(inputPtr, 0x158f689400000000000000000000000000000000000000000000000000000000) // swapWithoutSignature
            mstore(add(inputPtr, 4), 32)
            mstore(add(inputPtr, 36), payloadLength)
            let outputPtr := mload(0x40)
            mstore(0x40, add(outputPtr, 64))
            calldatacopy(add(inputPtr, 68), 68, payloadLength)

            // Override amountIn
            mstore(add(inputPtr, 132), mload(add(swapData, 96)))

            // Override fromAssetAddress
            let afterFromAssetAddress := mload(add(inputPtr, 112))
            mstore(add(inputPtr, 92), shl(96, mload(add(swapData, 32))))
            mstore(add(inputPtr, 112), afterFromAssetAddress)

            success := call(gas(), currentCoreAddress, 0, inputPtr, inputLength, outputPtr, 64)

            if eq(success, 1) {
                gasUsed := mload(outputPtr)
            }

            if eq(success, 0) {
                returndatacopy(0, 0, returndatasize())
                revert(0, returndatasize())
            }
        }

        transferredAmountOut = swapData.toAssetAddress.getBalance() - transferredAmountOut;

        if (swapData.transferFromRouter) {
            // Slippage in fromAsset
            swapData.toAssetAddress.transfer(swapData.toAddress, transferredAmountOut);
            uint256 approval = swapData.fromAssetAddress.allowance(address(this), currentCoreAddress);
            transferredAmountIn = swapData.amountIn - approval;
            if (approval > 0) {
                // Reset approval
                swapData.fromAssetAddress.approve(currentCoreAddress, 0);
            }
        } else {
            // Slippage in toAsset
            swapData.toAssetAddress.transfer(
                swapData.toAddress,
                (transferredAmountOut <= swapData.expectedAmountOut || swapData.expectedAmountOut == 0)
                    ? transferredAmountOut
                    : Math.max(
                        swapData.expectedAmountOut,
                        transferredAmountOut - ((transferredAmountOut * swapData.maxRetentionBps) / 10000)
                    )
            );
        }

        amountOut = swapData.toAssetAddress.getBalanceOf(swapData.toAddress) - amountOut;

        if (amountOut < swapData.amountOutMin) {
            revert InsufficientAmountOut();
        }

        emitSwap(
            EmitSwapParams({
                fromAddress: fromAddress,
                toAddress: swapData.toAddress,
                fromAssetAddress: fromAssetAddress,
                toAssetAddress: swapData.toAssetAddress,
                transferredAmountIn: transferredAmountIn,
                transferredAmountOut: transferredAmountOut,
                amountIn: swapData.amountIn,
                amountOut: amountOut,
                expectedAmountOut: swapData.expectedAmountOut,
                consumerId: swapData.consumerId,
                swapFeeAssetAddresses: swapData.swapFeeAssetAddresses,
                swapFeeReceivers: swapData.swapFeeReceivers,
                swapFeeAmounts: swapData.swapFeeAmounts
            })
        );
    }

    function adjustFromAssetAddress(SwapData memory swapData) private {
        if (swapData.fromAssetAddress.isNative()) {
            swapData.fromAssetAddress = weth;
            swapData.fromAssetAddress.wrap(swapData.amountIn);
        }
    }

    function adjustAmountIn(SwapData memory swapData, uint256 transferredNativeFeeAmount) private {
        // Auto
        if (swapData.amountIn == 0) {
            if (swapData.fromAssetAddress.isNative()) {
                swapData.amountIn = msg.value - transferredNativeFeeAmount;
            } else {
                uint256 allowance = swapData.fromAssetAddress.allowance(msg.sender, address(this));
                uint256 balance = swapData.fromAssetAddress.getBalanceOf(msg.sender);
                swapData.amountIn = allowance < balance ? allowance : balance;
            }
        }

        if (swapData.fromAssetAddress.isNative()) {
            if (transferredNativeFeeAmount + swapData.amountIn != msg.value) {
                revert InvalidNativeAmount();
            }
        } else {
            if (transferredNativeFeeAmount != msg.value) {
                revert InvalidNativeAmount();
            }
        }
    }

    /// @dev See {IDexAggregator-estimateSwapGas}
    function estimateSwapGas(
        bytes calldata
    ) external payable whenNotPaused returns (uint256 amountOut, uint256 gasUsed) {
        SwapData memory swapData = LibRouter.getData();
        address fromAssetAddress = swapData.fromAssetAddress;
        address fromAddress = verifyBackendSignature(swapData);
        uint256 transferredNativeFeeAmount = LibRouter.transferFees(swapData, fromAddress);
        adjustAmountIn(swapData, transferredNativeFeeAmount);
        (amountOut, gasUsed) = swap(swapData, fromAddress, fromAssetAddress);
    }

    /// @dev See {IDexAggregator-swapWithBackendSignature}
    function swapWithBackendSignature(bytes calldata) external payable whenNotPaused returns (uint256 amountOut) {
        SwapData memory swapData = LibRouter.getData();
        address fromAssetAddress = swapData.fromAssetAddress;
        address fromAddress = verifyBackendSignature(swapData);
        uint256 transferredNativeFeeAmount = LibRouter.transferFees(swapData, fromAddress);
        adjustAmountIn(swapData, transferredNativeFeeAmount);
        (amountOut, ) = swap(swapData, fromAddress, fromAssetAddress);
    }

    /// @dev Used to receive ethers
    receive() external payable {}
}