hyperlane-monorepo/solidity/contracts/libs/RateLimited.sol

// SPDX-License-Identifier: MIT OR Apache-2.0
pragma solidity >=0.8.0;
import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";

/**
 * @title RateLimited
 * @notice A contract used to keep track of an address sender's token amount limits.
 * @dev Implements a modified token bucket algorithm where the bucket is full in the beginning and gradually refills
 * See: https://dev.to/satrobit/rate-limiting-using-the-token-bucket-algorithm-3cjh
 **/
contract RateLimited is OwnableUpgradeable {
    uint256 public constant DURATION = 1 days; // 86400
    uint256 public filledLevel; /// @notice Current filled level
    uint256 public refillRate; /// @notice Tokens per second refill rate
    uint256 public lastUpdated; /// @notice Timestamp of the last time an action has been taken TODO prob can be uint40

    event RateLimitSet(uint256 _oldCapacity, uint256 _newCapacity);

    constructor(uint256 _capacity) {
        _transferOwnership(msg.sender);
        setRefillRate(_capacity);
        filledLevel = _capacity;
    }

    error RateLimitExceeded(uint256 newLimit, uint256 targetLimit);

    /**
     * @return The max capacity where the bucket will no longer refill
     */
    function maxCapacity() public view returns (uint256) {
        return refillRate * DURATION;
    }

    /**
     * Calculates the refilled amount based on time and refill rate
     *
     * Consider an example where there is a 1e18 max token limit per day (86400s)
     * If half of the tokens has been used, and half a day (43200s) has passed,
     * then there should be a refill of 0.5e18
     *
     * To calculate:
     *   Refilled = Elapsed * RefilledRate
     *   Elapsed = timestamp - Limit.lastUpdated
     *   RefilledRate = Capacity / DURATION
     *
     *   If half of the day (43200) has passed, then
     *   (86400 - 43200) * (1e18 / 86400)  = 0.5e18
     */
    function calculateRefilledAmount() internal view returns (uint256) {
        uint256 elapsed = block.timestamp - lastUpdated;
        return elapsed * refillRate;
    }

    /**
     * Calculates the adjusted fill level based on time
     */
    function calculateCurrentLevel() public view returns (uint256) {
        uint256 _capacity = maxCapacity();
        require(_capacity > 0, "RateLimitNotSet");

        if (block.timestamp > lastUpdated + DURATION) {
            // If last update is in the previous window, return the max capacity
            return _capacity;
        } else {
            // If within the window, refill the capacity
            uint256 replenishedLevel = filledLevel + calculateRefilledAmount();

            // Only return _capacity, in the case where newCurrentCapcacity overflows
            return replenishedLevel > _capacity ? _capacity : replenishedLevel;
        }
    }

    /**
     * Sets the refill rate by giving a capacity
     * @param _capacity new maximum capacity to set
     */
    function setRefillRate(
        uint256 _capacity
    ) public onlyOwner returns (uint256) {
        uint256 _oldRefillRate = refillRate;
        uint256 _newRefillRate = _capacity / DURATION;
        refillRate = _newRefillRate;

        emit RateLimitSet(_oldRefillRate, _newRefillRate);

        return _newRefillRate;
    }

    /**
     * Validate an amount and decreases the currentCapacity
     * @param _newAmount The amount to consume the fill level
     * @return The new filled level
     */
    function validateAndConsumeFilledLevel(
        uint256 _newAmount
    ) public returns (uint256) {
        uint256 adjustedFilledLevel = calculateCurrentLevel();
        require(_newAmount <= adjustedFilledLevel, "RateLimitExceeded");

        // Reduce the filledLevel and update lastUpdated
        uint256 _filledLevel = adjustedFilledLevel - _newAmount;
        filledLevel = _filledLevel;
        lastUpdated = block.timestamp;

        return _filledLevel;
    }
}