hyperlane-monorepo/solidity/test/hooks/RateLimitedHook.t.sol

// SPDX-License-Identifier: MIT or Apache-2.0
pragma solidity ^0.8.13;
import {Test} from "forge-std/Test.sol";

import {Message} from "contracts/libs/Message.sol";
import {TokenMessage} from "contracts/token/libs/TokenMessage.sol";
import {TypeCasts} from "contracts/libs/TypeCasts.sol";
import {RateLimited} from "contracts/libs/RateLimited.sol";

import {RateLimitedHook} from "contracts/hooks/warp-route/RateLimitedHook.sol";
import {HypERC20Collateral} from "contracts/token/HypERC20Collateral.sol";
import {HypERC20} from "contracts/token/HypERC20.sol";

import {TestMailbox} from "contracts/test/TestMailbox.sol";
import {ERC20Test} from "contracts/test/ERC20Test.sol";
import {TestPostDispatchHook} from "../../contracts/test/TestPostDispatchHook.sol";

contract RateLimitedHookTest is Test {
    using Message for bytes;
    using TypeCasts for address;

    uint32 constant ORIGIN = 11;
    uint32 constant DESTINATION = 12;
    uint256 constant MAX_CAPACITY = 1 ether;
    uint256 constant ONE_PERCENT = 0.01 ether;
    uint8 internal constant DECIMALS = 18;
    address constant BOB = address(0x2);

    TestMailbox localMailbox;
    TestMailbox remoteMailbox;
    ERC20Test token;
    TestPostDispatchHook internal noopHook;

    RateLimitedHook rateLimitedHook;
    HypERC20Collateral warpRouteLocal;
    HypERC20 warpRouteRemote;

    function _mintAndApprove(uint256 amount) internal {
        token.mint(amount);
        token.approve(address(warpRouteLocal), amount);
    }

    function setUp() external {
        localMailbox = new TestMailbox(ORIGIN);
        remoteMailbox = new TestMailbox(DESTINATION);

        token = new ERC20Test("Test", "Test", 100 ether, 18);
        noopHook = new TestPostDispatchHook();
        rateLimitedHook = new RateLimitedHook(
            address(localMailbox),
            MAX_CAPACITY
        );

        localMailbox.setDefaultHook(address(noopHook));
        localMailbox.setRequiredHook(address(noopHook));

        warpRouteLocal = new HypERC20Collateral(
            address(token),
            address(localMailbox)
        );

        warpRouteLocal.initialize(
            address(rateLimitedHook),
            address(0),
            address(this)
        );

        warpRouteRemote = new HypERC20(DECIMALS, address(remoteMailbox));

        warpRouteLocal.enrollRemoteRouter(
            DESTINATION,
            address(warpRouteRemote).addressToBytes32()
        );

        warpRouteRemote.enrollRemoteRouter(
            ORIGIN,
            address(warpRouteLocal).addressToBytes32()
        );
    }

    function testRateLimitedHook_revertsIfCalledByNonMailbox(
        bytes calldata _message
    ) external {
        vm.expectRevert("InvalidDispatchedMessage");
        rateLimitedHook.postDispatch(bytes(""), _message);
    }

    function testRateLimitedHook_revertsTransfer_ifExceedsFilledLevel(
        uint128 _amount,
        uint128 _time
    ) external {
        // Warp to a random time, get it's filled level, and try to transfer more than the target max
        vm.warp(_time);
        uint256 filledLevelBefore = rateLimitedHook.calculateCurrentLevel();
        vm.assume(_amount > filledLevelBefore);
        _mintAndApprove(_amount);

        vm.expectRevert("RateLimitExceeded");
        warpRouteLocal.transferRemote{value: 1}(
            DESTINATION,
            BOB.addressToBytes32(),
            _amount
        );
    }

    function testRateLimitedHook_allowsTransfer_ifUnderLimit(
        uint128 _amount,
        uint128 _time
    ) external {
        // Warp to a random time, get it's filled level, and try to transfer less than the target max
        vm.warp(_time);
        uint256 filledLevelBefore = rateLimitedHook.calculateCurrentLevel();
        vm.assume(_amount <= filledLevelBefore);

        _mintAndApprove(_amount);
        warpRouteLocal.transferRemote(
            DESTINATION,
            BOB.addressToBytes32(),
            _amount
        );
        uint256 limitAfter = rateLimitedHook.calculateCurrentLevel();
        assertApproxEqRel(limitAfter, filledLevelBefore - _amount, ONE_PERCENT);
    }

    function testRateLimitedHook_preventsDuplicateMessageFromValidating(
        uint128 _amount
    ) public {
        // Warp to a random time, get it's filled level, and try to transfer less than the target max
        vm.warp(1 days);
        uint256 filledLevelBefore = rateLimitedHook.calculateCurrentLevel();
        vm.assume(_amount <= filledLevelBefore);

        _mintAndApprove(_amount);

        // Generate an outbound message that will be the same as the one created in transferRemote
        bytes memory tokenMessage = TokenMessage.format(
            BOB.addressToBytes32(),
            _amount,
            bytes("")
        );
        vm.prank(address(warpRouteLocal));
        bytes memory message = localMailbox.buildOutboundMessage(
            DESTINATION,
            address(warpRouteRemote).addressToBytes32(),
            tokenMessage
        );

        bytes32 messageId = warpRouteLocal.transferRemote(
            DESTINATION,
            BOB.addressToBytes32(),
            _amount
        );

        assertEq(message.id(), messageId);

        vm.expectRevert("MessageAlreadyValidated");
        rateLimitedHook.postDispatch(bytes(""), message);
    }
}