fix(contracts): `RateLimit` minor changes (#4575)

### Description

- Added a check for invalid capacity and event for token level change

### Drive-by changes

None

### Related issues

- fixes https://github.com/chainlight-io/2024-08-hyperlane/issues/14

### Backward compatibility

Yes

### Testing

Unit tests
pull/4724/head
Kunal Arora 1 month ago committed by GitHub
parent 3e1ab75643
commit c9bd7c3c52
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 42
      solidity/contracts/libs/RateLimited.sol
  2. 47
      solidity/test/lib/RateLimited.t.sol

@ -1,5 +1,19 @@
// SPDX-License-Identifier: MIT OR Apache-2.0
pragma solidity >=0.8.0;
/*@@@@@@@ @@@@@@@@@
@@@@@@@@@ @@@@@@@@@
@@@@@@@@@ @@@@@@@@@
@@@@@@@@@ @@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@ HYPERLANE @@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@ @@@@@@@@@
@@@@@@@@@ @@@@@@@@@
@@@@@@@@@ @@@@@@@@@
@@@@@@@@@ @@@@@@@@*/
// ============ External Imports ============
import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
/**
@ -7,16 +21,26 @@ import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/Own
* @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
/// @notice Current filled level
uint256 public filledLevel;
/// @notice Tokens per second refill rate
uint256 public refillRate;
/// @notice Timestamp of the last time an action has been taken
uint256 public lastUpdated;
event RateLimitSet(uint256 _oldCapacity, uint256 _newCapacity);
event ConsumedFilledLevel(uint256 filledLevel, uint256 lastUpdated);
constructor(uint256 _capacity) {
require(
_capacity >= DURATION,
"Capacity must be greater than DURATION"
);
_transferOwnership(msg.sender);
setRefillRate(_capacity);
filledLevel = _capacity;
@ -88,20 +112,22 @@ contract RateLimited is OwnableUpgradeable {
/**
* Validate an amount and decreases the currentCapacity
* @param _newAmount The amount to consume the fill level
* @param _consumedAmount The amount to consume the fill level
* @return The new filled level
*/
function validateAndConsumeFilledLevel(
uint256 _newAmount
uint256 _consumedAmount
) public returns (uint256) {
uint256 adjustedFilledLevel = calculateCurrentLevel();
require(_newAmount <= adjustedFilledLevel, "RateLimitExceeded");
require(_consumedAmount <= adjustedFilledLevel, "RateLimitExceeded");
// Reduce the filledLevel and update lastUpdated
uint256 _filledLevel = adjustedFilledLevel - _newAmount;
uint256 _filledLevel = adjustedFilledLevel - _consumedAmount;
filledLevel = _filledLevel;
lastUpdated = block.timestamp;
emit ConsumedFilledLevel(filledLevel, lastUpdated);
return _filledLevel;
}
}

@ -1,5 +1,6 @@
// SPDX-License-Identifier: MIT or Apache-2.0
pragma solidity ^0.8.13;
import {Test} from "forge-std/Test.sol";
import {RateLimited} from "../../contracts/libs/RateLimited.sol";
@ -13,8 +14,13 @@ contract RateLimitLibTest is Test {
rateLimited = new RateLimited(MAX_CAPACITY);
}
function testConstructor_revertsWhen_lowCapacity() public {
vm.expectRevert("Capacity must be greater than DURATION");
new RateLimited(1 days - 1);
}
function testRateLimited_setsNewLimit() external {
rateLimited.setRefillRate(2 ether);
assert(rateLimited.setRefillRate(2 ether) > 0);
assertApproxEqRel(rateLimited.maxCapacity(), 2 ether, ONE_PERCENT);
assertEq(rateLimited.refillRate(), uint256(2 ether) / 1 days); // 2 ether / 1 day
}
@ -45,6 +51,25 @@ contract RateLimitLibTest is Test {
rateLimited.setRefillRate(1 ether);
}
function testConsumedFilledLevelEvent() public {
uint256 consumeAmount = 0.5 ether;
vm.expectEmit(true, true, false, true);
emit RateLimited.ConsumedFilledLevel(
499999999999993600,
block.timestamp
); // precision loss
rateLimited.validateAndConsumeFilledLevel(consumeAmount);
assertApproxEqRelDecimal(
rateLimited.filledLevel(),
MAX_CAPACITY - consumeAmount,
1e14,
0
);
assertEq(rateLimited.lastUpdated(), block.timestamp);
}
function testRateLimited_neverReturnsGtMaxLimit(
uint256 _newAmount,
uint40 _newTime
@ -104,4 +129,24 @@ contract RateLimitLibTest is Test {
currentTargetLimit = rateLimited.calculateCurrentLevel();
assertApproxEqRel(currentTargetLimit, MAX_CAPACITY, ONE_PERCENT);
}
function testCalculateCurrentLevel_revertsWhenCapacityIsZero() public {
rateLimited.setRefillRate(0);
vm.expectRevert("RateLimitNotSet");
rateLimited.calculateCurrentLevel();
}
function testValidateAndConsumeFilledLevel_revertsWhenExceedingLimit()
public
{
vm.warp(1 days);
uint256 initialLevel = rateLimited.calculateCurrentLevel();
uint256 excessAmount = initialLevel + 1 ether;
vm.expectRevert("RateLimitExceeded");
rateLimited.validateAndConsumeFilledLevel(excessAmount);
assertEq(rateLimited.calculateCurrentLevel(), initialLevel);
}
}

Loading…
Cancel
Save