diff --git a/.changeset/rate-limiter-library.md b/.changeset/rate-limiter-library.md new file mode 100644 index 00000000000..164e63cbbe3 --- /dev/null +++ b/.changeset/rate-limiter-library.md @@ -0,0 +1,5 @@ +--- +'openzeppelin-solidity': minor +--- + +`RateLimiter`: Add a library that provides primitives for limiting the rate at which an action can be performed, with two complementary strategies: a refilling token bucket and a sliding window counter. diff --git a/contracts/mocks/Stateless.sol b/contracts/mocks/Stateless.sol index 0669c675248..da281280b3d 100644 --- a/contracts/mocks/Stateless.sol +++ b/contracts/mocks/Stateless.sol @@ -46,6 +46,7 @@ import {NoncesKeyed} from "../utils/NoncesKeyed.sol"; import {P256} from "../utils/cryptography/P256.sol"; import {Packing} from "../utils/Packing.sol"; import {Panic} from "../utils/Panic.sol"; +import {RateLimiter} from "../utils/RateLimiter.sol"; import {RelayedCall} from "../utils/RelayedCall.sol"; import {RLP} from "../utils/RLP.sol"; import {RSA} from "../utils/cryptography/RSA.sol"; diff --git a/contracts/utils/README.adoc b/contracts/utils/README.adoc index 10cc2c85fd4..ff2e1a04200 100644 --- a/contracts/utils/README.adoc +++ b/contracts/utils/README.adoc @@ -40,6 +40,7 @@ Miscellaneous contracts and libraries containing utility functions you can use t * {Multicall}: Abstract contract with a utility to allow batching together multiple calls in a single transaction. Useful for allowing EOAs to perform multiple operations at once. * {Packing}: A library for packing and unpacking multiple values into bytes32. * {Panic}: A library to revert with https://docs.soliditylang.org/en/v0.8.20/control-structures.html#panic-via-assert-and-error-via-require[Solidity panic codes]. + * {RateLimiter}: A library that provides primitives for limiting the rate at which an action can be performed, using a refilling token bucket or a sliding window counter. * {RelayedCall}: A library for performing calls that use minimal and predictable relayers to hide the sender. * {RLP}: Library for encoding and decoding data in Ethereum's Recursive Length Prefix format. * {ShortStrings}: Library to encode (and decode) short strings into (or from) a single bytes32 slot for optimizing costs. Short strings are limited to 31 characters. @@ -147,6 +148,8 @@ Ethereum contracts have no native concept of an interface, so applications must {{Panic}} +{{RateLimiter}} + {{RelayedCall}} {{RLP}} diff --git a/contracts/utils/RateLimiter.sol b/contracts/utils/RateLimiter.sol new file mode 100644 index 00000000000..4fb08b4cef2 --- /dev/null +++ b/contracts/utils/RateLimiter.sol @@ -0,0 +1,264 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.27; + +import {Math} from "./math/Math.sol"; +import {SafeCast} from "./math/SafeCast.sol"; +import {Checkpoints} from "./structs/Checkpoints.sol"; +import {Time} from "./types/Time.sol"; + +/** + * @dev This library provides primitives for limiting the rate at which an action can be performed. + * + * Two complementary strategies are available, each represented by a storage struct that the consumer keeps in its + * own storage: + * + * - {RefillingBucket}: a token bucket that refills linearly over time. Suitable when the protected resource is + * expected to regenerate continuously and short bursts up to the bucket capacity are acceptable. Storage cost is + * constant regardless of consumption history. + * + * - {SlidingWindow}: a moving-window counter that caps the cumulative consumption over any `window`-second + * interval. Suitable when a strict cap on usage within a rolling window is required. Each successful consumption + * appends a checkpoint, making it a most expensive option with a larger storage footprint. + * + * Both strategies expose the same set of operations ({state}, {used}, {available}, {tryConsume}, {consume} and + * {updateSettings}), distinguished by the storage struct passed as the first argument. + * + * Example usage: + * + * ```solidity + * using RateLimiter for RateLimiter.RefillingBucket; + * + * mapping(address user => RateLimiter.RefillingBucket) private _withdrawLimits; + * + * function withdraw(uint256 amount) external { + * _withdrawLimits[msg.sender].consume(amount); + * // ... + * } + * ``` + */ +library RateLimiter { + using Checkpoints for Checkpoints.Trace208; + + /** + * @dev The requested quantity exceeds the currently available capacity. + */ + error RateLimitExceeded(); + + // ================================================ RefillingBucket ================================================ + /** + * @dev A token bucket that refills linearly over time. + * + * The bucket has a maximum `capacity` and refills at a rate of `capacity / window` per second, so that an empty + * bucket fully refills in `window` seconds. The current state is reconstructed lazily from `lastUsed` and + * `lastTimepoint` on read, keeping storage cost constant (2 packed slots). + */ + struct RefillingBucketItem { + uint208 lastUsed; + uint48 lastTimepoint; + } + struct RefillingBucket { + uint208 capacity; + uint48 window; + mapping(bytes32 key => RefillingBucketItem) items; + } + + /** + * @dev Returns the current `used` and `available` quantities for a {RefillingBucket}, accounting for the + * time-based refill that has accrued since the last update. + */ + function state( + RefillingBucket storage self, + bytes32 key + ) internal view returns (uint256 used_, uint256 available_) { + uint208 cacheCapacity = self.capacity; + uint48 cacheWindow = self.window; + uint208 cacheLastUsed = self.items[key].lastUsed; + uint48 cacheLastTimepoint = self.items[key].lastTimepoint; + + used_ = Math.saturatingSub( + cacheLastUsed, + Math.mulDiv(Time.timestamp() - cacheLastTimepoint, cacheCapacity, Math.max(cacheWindow, 1)) + ); + available_ = Math.saturatingSub(cacheCapacity, used_); + } + + /** + * @dev Returns the currently used quantity. See {state}. + */ + function used(RefillingBucket storage self, bytes32 key) internal view returns (uint256 used_) { + (used_, ) = state(self, key); + } + + /** + * @dev Returns the currently available quantity. See {state}. + */ + function available(RefillingBucket storage self, bytes32 key) internal view returns (uint256 available_) { + (, available_) = state(self, key); + } + + /** + * @dev Attempts to consume `quantity` from the bucket. Returns `true` on success, `false` if the available + * quantity is insufficient. + * + * A `quantity` of 0 is always accepted and does not modify storage. + */ + function tryConsume(RefillingBucket storage self, bytes32 key, uint256 quantity) internal returns (bool) { + (uint256 used_, uint256 available_) = state(self, key); + if (quantity == 0) { + return true; + } else if (quantity <= available_) { + self.items[key].lastTimepoint = Time.timestamp(); + self.items[key].lastUsed = SafeCast.toUint208(used_ + quantity); + return true; + } else { + return false; + } + } + + /** + * @dev Consumes `quantity` from the bucket. Reverts with {RateLimitExceeded} if the available quantity is + * insufficient. See {tryConsume}. + */ + function consume(RefillingBucket storage self, bytes32 key, uint256 quantity) internal { + bool success = tryConsume(self, key, quantity); + require(success, RateLimitExceeded()); + } + + /** + * @dev Resets the bucket to a fully-available state. + * + * The `capacity` and `window` settings are preserved; only the consumed quantity is cleared. + */ + function reset(RefillingBucket storage self, bytes32 key) internal { + self.items[key].lastUsed = 0; + } + + /** + * @dev Updates the `capacity` and `window` of the bucket. + * + * NOTE: The new settings will retroactively affect all the keys. The new replenishing rate (capacity / window) is + * applied from the last update timepoint of each key. Therefore, if the new settings correspond to a faster + * replenishing rate, some quantity may become available immediately. Conversely, if the new settings correspond + * to a slower replenishing rate, some quantity that would otherwise be available immediately may become + * unavailable. This side effect can be mitigated by calling {refresh} on the relevant keys before updating the + * settings. There is no mechanism to automatically refresh all the keys in a single operation. + */ + function updateSettings(RefillingBucket storage self, uint48 newWindow, uint208 newCapacity) internal { + self.capacity = newCapacity; + self.window = newWindow; + } + + /** + * @dev Refreshes the bucket by applying the accrued refill since the last update timepoint to `lastUsed` and + * `lastTimepoint`, effectively moving the timepoint forward to now. This can be used to mitigate the side effect + * of {updateSettings} when the refreshing rate is modified. + */ + function refresh(RefillingBucket storage self, bytes32 key) internal { + self.items[key].lastUsed = uint208(used(self, key)); + self.items[key].lastTimepoint = Time.timestamp(); + } + + // ================================================= SlidingWindow ================================================= + /** + * @dev A moving-window counter that caps cumulative consumption within any `window`-second interval. + * + * Each successful consumption appends a checkpoint to `history` recording the running cumulative total. The + * current `used` quantity is the difference between the cumulative total at `block.timestamp` and the cumulative + * total at `block.timestamp - window`. + * + * NOTE: The cumulative total is stored as a `uint208`. Once it reaches `2²⁰⁸ - 1`, further consumption will + * revert in {SafeCast}. This bound is unreachable for any realistic `limit`, but consumers should be aware of it. + * + * NOTE: Old checkpoints are never pruned. The storage footprint grows with the number of {tryConsume} calls + * that succeed with a non-zero `quantity`. + */ + struct SlidingWindow { + uint208 limit; + uint48 window; + mapping(bytes32 key => Checkpoints.Trace208) items; + } + + /** + * @dev Returns the current `used` and `available` quantities for a {SlidingWindow}, computed as the cumulative + * consumption over the last `window` seconds. + */ + function state(SlidingWindow storage self, bytes32 key) internal view returns (uint256 used_, uint256 available_) { + uint208 cacheLimit = self.limit; + uint48 cacheWindow = self.window; + + used_ = Math.saturatingSub( + self.items[key].latest(), + self.items[key].upperLookupRecent(uint48(Math.saturatingSub(Time.timestamp(), Math.max(cacheWindow, 1)))) + ); + available_ = Math.saturatingSub(cacheLimit, used_); + } + + /** + * @dev Returns the currently used quantity within the rolling window. See {state}. + */ + function used(SlidingWindow storage self, bytes32 key) internal view returns (uint256 used_) { + (used_, ) = state(self, key); + } + + /** + * @dev Returns the currently available quantity within the rolling window. See {state}. + */ + function available(SlidingWindow storage self, bytes32 key) internal view returns (uint256 available_) { + (, available_) = state(self, key); + } + + /** + * @dev Attempts to record a consumption of `quantity`. Returns `true` on success, `false` if the available + * quantity within the current window is insufficient. + * + * A `quantity` of 0 is always accepted and does not modify storage. + */ + function tryConsume(SlidingWindow storage self, bytes32 key, uint256 quantity) internal returns (bool) { + if (quantity == 0) { + return true; + } else if (quantity <= available(self, key)) { + self.items[key].push(Time.timestamp(), SafeCast.toUint208(self.items[key].latest() + quantity)); + return true; + } else { + return false; + } + } + + /** + * @dev Records a consumption of `quantity`. Reverts with {RateLimitExceeded} if the available quantity within + * the current window is insufficient. See {tryConsume}. + */ + function consume(SlidingWindow storage self, bytes32 key, uint256 quantity) internal { + bool success = tryConsume(self, key, quantity); + require(success, RateLimitExceeded()); + } + + /** + * @dev Resets the rolling window to a fully-available state. + * + * The `limit` and `window` settings are preserved; only the consumed quantity is cleared. + * + * NOTE: This will reset the entire history, meaning it can also be used to recover from the cumulative total + * approaching the `uint208` ceiling. The underlying storage slots holding past checkpoints are not zeroed out. + * As a consequence, there is no gas refunded, but future {consume}/{tryConsume} operations are cheaper from + * reusing "dirty" slots. + */ + function reset(SlidingWindow storage self, bytes32 key) internal { + Checkpoints.Checkpoint208[] storage trace = self.items[key]._checkpoints; + assembly ("memory-safe") { + sstore(trace.slot, 0) + } + } + + /** + * @dev Updates the `limit` and `window` of the rate limiter. + * + * NOTE: The history of past consumptions is not modified. Increasing `window` retroactively brings older + * consumptions back into the rolling window until they age out under the new duration; decreasing `window` + * conversely causes older consumptions to drop out sooner. + */ + function updateSettings(SlidingWindow storage self, uint48 newWindow, uint208 newLimit) internal { + self.limit = newLimit; + self.window = newWindow; + } +} diff --git a/test/utils/RateLimiter.test.js b/test/utils/RateLimiter.test.js new file mode 100644 index 00000000000..0ff944e442f --- /dev/null +++ b/test/utils/RateLimiter.test.js @@ -0,0 +1,537 @@ +const { ethers } = require('hardhat'); +const { expect } = require('chai'); +const { loadFixture } = require('@nomicfoundation/hardhat-network-helpers'); + +const { MAX_UINT48 } = require('../helpers/constants'); +const { batchInBlock } = require('../helpers/txpool'); +const time = require('../helpers/time'); + +const WINDOW = 100n; +const CAPACITY = 1000n; +const refillPerSecond = CAPACITY / WINDOW; + +const defaultKey = ethers.ZeroHash; +const key1 = ethers.id('key1'); +const key2 = ethers.id('key2'); + +async function fixture() { + const mock = await ethers.deployContract('$RateLimiter'); + return { mock }; +} + +const wrap = (mock, type) => ({ + state: (k = defaultKey) => mock.getFunction(`$state_RateLimiter_${type}`)(0n, k), + used: (k = defaultKey) => mock.getFunction(`$used_RateLimiter_${type}`)(0n, k), + available: (k = defaultKey) => mock.getFunction(`$available_RateLimiter_${type}`)(0n, k), + tryConsume: (q, k = defaultKey) => mock.getFunction(`$tryConsume_RateLimiter_${type}`)(0n, k, q), + tryConsumeStatic: (q, k = defaultKey) => mock.getFunction(`$tryConsume_RateLimiter_${type}`).staticCall(0n, k, q), + consume: (q, k = defaultKey) => mock.getFunction(`$consume_RateLimiter_${type}`)(0n, k, q), + reset: (k = defaultKey) => mock.getFunction(`$reset_RateLimiter_${type}`)(0n, k), + updateSettings: (window, capacity) => mock.getFunction(`$updateSettings_RateLimiter_${type}`)(0n, window, capacity), + refresh: type == 'RefillingBucket' ? (k = defaultKey) => mock.$refresh(0n, k) : undefined, +}); + +describe('RateLimiter', function () { + beforeEach(async function () { + Object.assign(this, await loadFixture(fixture)); + }); + + describe('RefillingBucket', function () { + beforeEach(async function () { + Object.assign(this.mock, wrap(this.mock, 'RefillingBucket')); + }); + + it('starts empty', async function () { + await expect(this.mock.state()).to.eventually.deep.equal([0n, 0n]); + await expect(this.mock.used()).to.eventually.equal(0n); + await expect(this.mock.available()).to.eventually.equal(0n); + }); + + describe('with some capacity', function () { + beforeEach(async function () { + await this.mock.updateSettings(WINDOW, CAPACITY); + }); + + it('reports full capacity available', async function () { + await expect(this.mock.state()).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used()).to.eventually.equal(0n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY); + }); + + describe('consume', function () { + it('consume reduces available and increases used', async function () { + await this.mock.consume(17n); + + await expect(this.mock.state()).to.eventually.deep.equal([17n, CAPACITY - 17n]); + await expect(this.mock.used()).to.eventually.equal(17n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 17n); + }); + + it('consume one key does not affect another key', async function () { + await expect(this.mock.state(key1)).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used(key1)).to.eventually.equal(0n); + await expect(this.mock.available(key1)).to.eventually.equal(CAPACITY); + await expect(this.mock.state(key2)).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used(key2)).to.eventually.equal(0n); + await expect(this.mock.available(key2)).to.eventually.equal(CAPACITY); + + await this.mock.consume(17n, key1); + + await expect(this.mock.state(key1)).to.eventually.deep.equal([17n, CAPACITY - 17n]); + await expect(this.mock.used(key1)).to.eventually.equal(17n); + await expect(this.mock.available(key1)).to.eventually.equal(CAPACITY - 17n); + await expect(this.mock.state(key2)).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used(key2)).to.eventually.equal(0n); + await expect(this.mock.available(key2)).to.eventually.equal(CAPACITY); + + await this.mock.consume(42n, key2); + + await expect(this.mock.state(key1)).to.eventually.deep.equal([ + 17n - refillPerSecond, + CAPACITY - 17n + refillPerSecond, + ]); + await expect(this.mock.used(key1)).to.eventually.equal(17n - refillPerSecond); + await expect(this.mock.available(key1)).to.eventually.equal(CAPACITY - 17n + refillPerSecond); + await expect(this.mock.state(key2)).to.eventually.deep.equal([42n, CAPACITY - 42n]); + await expect(this.mock.used(key2)).to.eventually.equal(42n); + await expect(this.mock.available(key2)).to.eventually.equal(CAPACITY - 42n); + }); + + it('consume reverts when over capacity', async function () { + await expect(this.mock.consume(CAPACITY + 1n)).to.be.revertedWithCustomError(this.mock, 'RateLimitExceeded'); + }); + + it('consume(0) is a no-op that returns true', async function () { + await this.mock.consume(42n); + + await expect(this.mock.state()).to.eventually.deep.equal([42n, CAPACITY - 42n]); + await expect(this.mock.used()).to.eventually.equal(42n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 42n); + + await this.mock.consume(0n); + + // one second passed so we have to account for the auto refill + await expect(this.mock.state()).to.eventually.deep.equal([ + 42n - refillPerSecond, + CAPACITY - 42n + refillPerSecond, + ]); + await expect(this.mock.used()).to.eventually.equal(42n - refillPerSecond); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 42n + refillPerSecond); + }); + }); + + describe('tryConsume', function () { + it('tryConsume reduces available and increases used', async function () { + // Static + await expect(this.mock.tryConsumeStatic(42n)).to.eventually.be.true; + + // execute (and get event) + await expect(this.mock.tryConsume(42n)) + .to.emit(this.mock, 'return$tryConsume_RateLimiter_RefillingBucket_bytes32_uint256') + .withArgs(true); + + await expect(this.mock.state()).to.eventually.deep.equal([42n, CAPACITY - 42n]); + await expect(this.mock.used()).to.eventually.equal(42n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 42n); + }); + + it('tryConsume one key does not affect another key', async function () { + await expect(this.mock.state(key1)).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used(key1)).to.eventually.equal(0n); + await expect(this.mock.available(key1)).to.eventually.equal(CAPACITY); + await expect(this.mock.state(key2)).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used(key2)).to.eventually.equal(0n); + await expect(this.mock.available(key2)).to.eventually.equal(CAPACITY); + + await expect(this.mock.tryConsume(17n, key1)) + .to.emit(this.mock, 'return$tryConsume_RateLimiter_RefillingBucket_bytes32_uint256') + .withArgs(true); + + await expect(this.mock.state(key1)).to.eventually.deep.equal([17n, CAPACITY - 17n]); + await expect(this.mock.used(key1)).to.eventually.equal(17n); + await expect(this.mock.available(key1)).to.eventually.equal(CAPACITY - 17n); + await expect(this.mock.state(key2)).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used(key2)).to.eventually.equal(0n); + await expect(this.mock.available(key2)).to.eventually.equal(CAPACITY); + + await expect(this.mock.tryConsume(42n, key2)) + .to.emit(this.mock, 'return$tryConsume_RateLimiter_RefillingBucket_bytes32_uint256') + .withArgs(true); + + await expect(this.mock.state(key1)).to.eventually.deep.equal([ + 17n - refillPerSecond, + CAPACITY - 17n + refillPerSecond, + ]); + await expect(this.mock.used(key1)).to.eventually.equal(17n - refillPerSecond); + await expect(this.mock.available(key1)).to.eventually.equal(CAPACITY - 17n + refillPerSecond); + await expect(this.mock.state(key2)).to.eventually.deep.equal([42n, CAPACITY - 42n]); + await expect(this.mock.used(key2)).to.eventually.equal(42n); + await expect(this.mock.available(key2)).to.eventually.equal(CAPACITY - 42n); + }); + + it('tryConsume returns false and does not update state when over capacity', async function () { + // Static + await expect(this.mock.tryConsumeStatic(CAPACITY + 1n)).to.eventually.be.false; + + // execute (and get event) + await expect(this.mock.tryConsume(CAPACITY + 1n)) + .to.emit(this.mock, 'return$tryConsume_RateLimiter_RefillingBucket_bytes32_uint256') + .withArgs(false); + + await expect(this.mock.state()).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used()).to.eventually.equal(0n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY); + }); + + it('consume(0) is a no-op that returns true', async function () { + await this.mock.consume(42n); + + await expect(this.mock.state()).to.eventually.deep.equal([42n, CAPACITY - 42n]); + await expect(this.mock.used()).to.eventually.equal(42n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 42n); + + await expect(this.mock.tryConsume(0n)) + .to.emit(this.mock, 'return$tryConsume_RateLimiter_RefillingBucket_bytes32_uint256') + .withArgs(true); + + await expect(this.mock.state()).to.eventually.deep.equal([ + 42n - refillPerSecond, + CAPACITY - 42n + refillPerSecond, + ]); + await expect(this.mock.used()).to.eventually.equal(42n - refillPerSecond); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 42n + refillPerSecond); + }); + }); + + it('refills linearly over time', async function () { + const wait = 17n; + const refilled = (wait * CAPACITY) / WINDOW; + + await this.mock.consume(CAPACITY); + + // bucket is empty; advance time by half a window + await time.increaseBy.timestamp(wait); + + await expect(this.mock.state()).to.eventually.deep.equal([CAPACITY - refilled, refilled]); + await expect(this.mock.used()).to.eventually.equal(CAPACITY - refilled); + await expect(this.mock.available()).to.eventually.equal(refilled); + }); + + it('fully refills after a window has elapsed', async function () { + await this.mock.consume(CAPACITY); + await time.increaseBy.timestamp(WINDOW); + + await expect(this.mock.state()).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used()).to.eventually.equal(0n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY); + }); + + it('used does not go below zero', async function () { + await this.mock.consume(10n); + await time.increaseBy.timestamp(WINDOW + 1n); + + await expect(this.mock.state()).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used()).to.eventually.equal(0n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY); + }); + + it('updateSettings side effect on usage', async function () { + const d1 = 3n; + const d2 = 4n; + const newRefillPerSecond = (2n * CAPACITY) / WINDOW; + + // consume the whole bucket + await this.mock.consume(CAPACITY); + + await time.increaseBy.timestamp(d1, false); + await this.mock.updateSettings(WINDOW, CAPACITY * 2n); + await time.increaseBy.timestamp(d2); + + await expect(this.mock.state()).to.eventually.deep.equal([ + CAPACITY - (d1 + d2) * newRefillPerSecond, + CAPACITY + (d1 + d2) * newRefillPerSecond, + ]); + await expect(this.mock.used()).to.eventually.equal(CAPACITY - (d1 + d2) * newRefillPerSecond); + await expect(this.mock.available()).to.eventually.equal(CAPACITY + (d1 + d2) * newRefillPerSecond); + }); + + it('using refresh to mitigate updateSettings side effect', async function () { + const d1 = 3n; + const d2 = 4n; + const newRefillPerSecond = (2n * CAPACITY) / WINDOW; + + // consume the whole bucket + await this.mock.consume(CAPACITY); + + await time.increaseBy.timestamp(d1, false); + await batchInBlock([() => this.mock.refresh(), () => this.mock.updateSettings(WINDOW, CAPACITY * 2n)]); + await time.increaseBy.timestamp(d2); + + await expect(this.mock.state()).to.eventually.deep.equal([ + CAPACITY - d1 * refillPerSecond - d2 * newRefillPerSecond, + CAPACITY + d1 * refillPerSecond + d2 * newRefillPerSecond, + ]); + await expect(this.mock.used()).to.eventually.equal(CAPACITY - d1 * refillPerSecond - d2 * newRefillPerSecond); + await expect(this.mock.available()).to.eventually.equal( + CAPACITY + d1 * refillPerSecond + d2 * newRefillPerSecond, + ); + }); + + it('reset clears used', async function () { + await this.mock.consume(CAPACITY / 2n); + + await this.mock.reset(); + + await expect(this.mock.state()).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used()).to.eventually.equal(0n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY); + }); + + it('window saturation prevents underflow when block.timestamp < window', async function () { + // computing state with a window larger than block.timestamp must not revert + await this.mock.updateSettings(MAX_UINT48, CAPACITY); + await expect(this.mock.state()).to.not.be.reverted; + }); + }); + }); + + describe('SlidingWindow', function () { + beforeEach(async function () { + Object.assign(this.mock, wrap(this.mock, 'SlidingWindow')); + }); + + it('starts empty', async function () { + await expect(this.mock.state()).to.eventually.deep.equal([0n, 0n]); + await expect(this.mock.used()).to.eventually.equal(0n); + await expect(this.mock.available()).to.eventually.equal(0n); + }); + + describe('with some capacity', function () { + beforeEach(async function () { + await this.mock.updateSettings(WINDOW, CAPACITY); + }); + + it('reports full limit available', async function () { + await expect(this.mock.state()).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used()).to.eventually.equal(0n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY); + }); + + describe('consume', function () { + it('consume reduces available and increases used', async function () { + await this.mock.consume(17n); + + await expect(this.mock.state()).to.eventually.deep.equal([17n, CAPACITY - 17n]); + await expect(this.mock.used()).to.eventually.equal(17n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 17n); + }); + + it('consume one key does not affect another key', async function () { + const key1 = ethers.id('key1'); + const key2 = ethers.id('key2'); + + await expect(this.mock.state(key1)).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used(key1)).to.eventually.equal(0n); + await expect(this.mock.available(key1)).to.eventually.equal(CAPACITY); + await expect(this.mock.state(key2)).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used(key2)).to.eventually.equal(0n); + await expect(this.mock.available(key2)).to.eventually.equal(CAPACITY); + + await this.mock.consume(17n, key1); + + await expect(this.mock.state(key1)).to.eventually.deep.equal([17n, CAPACITY - 17n]); + await expect(this.mock.used(key1)).to.eventually.equal(17n); + await expect(this.mock.available(key1)).to.eventually.equal(CAPACITY - 17n); + await expect(this.mock.state(key2)).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used(key2)).to.eventually.equal(0n); + await expect(this.mock.available(key2)).to.eventually.equal(CAPACITY); + + await this.mock.consume(42n, key2); + + await expect(this.mock.state(key1)).to.eventually.deep.equal([17n, CAPACITY - 17n]); + await expect(this.mock.used(key1)).to.eventually.equal(17n); + await expect(this.mock.available(key1)).to.eventually.equal(CAPACITY - 17n); + await expect(this.mock.state(key2)).to.eventually.deep.equal([42n, CAPACITY - 42n]); + await expect(this.mock.used(key2)).to.eventually.equal(42n); + await expect(this.mock.available(key2)).to.eventually.equal(CAPACITY - 42n); + }); + + it('multiple consume accumulate within the window', async function () { + await this.mock.consume(100n); + await this.mock.consume(200n); + await this.mock.consume(50n); + + await expect(this.mock.state()).to.eventually.deep.equal([350n, CAPACITY - 350n]); + await expect(this.mock.used()).to.eventually.equal(350n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 350n); + }); + + it('consume reverts when over limit', async function () { + await expect(this.mock.consume(CAPACITY + 1n)).to.be.revertedWithCustomError(this.mock, 'RateLimitExceeded'); + }); + + it('consume(0) is a no-op that returns true', async function () { + await this.mock.consume(42n); + + await expect(this.mock.state()).to.eventually.deep.equal([42n, CAPACITY - 42n]); + await expect(this.mock.used()).to.eventually.equal(42n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 42n); + + await this.mock.consume(0n); + + // one second passed so we have to account for the auto refill + await expect(this.mock.state()).to.eventually.deep.equal([42n, CAPACITY - 42n]); + await expect(this.mock.used()).to.eventually.equal(42n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 42n); + }); + }); + + describe('tryConsume', function () { + it('tryConsume reduces available and increases used', async function () { + // Static + await expect(this.mock.tryConsumeStatic(42n)).to.eventually.be.true; + + // execute (and get event) + await expect(this.mock.tryConsume(42n)) + .to.emit(this.mock, 'return$tryConsume_RateLimiter_SlidingWindow_bytes32_uint256') + .withArgs(true); + + await expect(this.mock.state()).to.eventually.deep.equal([42n, CAPACITY - 42n]); + await expect(this.mock.used()).to.eventually.equal(42n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 42n); + }); + + it('tryConsume one key does not affect another key', async function () { + const key1 = ethers.id('key1'); + const key2 = ethers.id('key2'); + + await expect(this.mock.state(key1)).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used(key1)).to.eventually.equal(0n); + await expect(this.mock.available(key1)).to.eventually.equal(CAPACITY); + await expect(this.mock.state(key2)).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used(key2)).to.eventually.equal(0n); + await expect(this.mock.available(key2)).to.eventually.equal(CAPACITY); + + await expect(this.mock.tryConsume(17n, key1)) + .to.emit(this.mock, 'return$tryConsume_RateLimiter_SlidingWindow_bytes32_uint256') + .withArgs(true); + + await expect(this.mock.state(key1)).to.eventually.deep.equal([17n, CAPACITY - 17n]); + await expect(this.mock.used(key1)).to.eventually.equal(17n); + await expect(this.mock.available(key1)).to.eventually.equal(CAPACITY - 17n); + await expect(this.mock.state(key2)).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used(key2)).to.eventually.equal(0n); + await expect(this.mock.available(key2)).to.eventually.equal(CAPACITY); + + await expect(this.mock.tryConsume(42n, key2)) + .to.emit(this.mock, 'return$tryConsume_RateLimiter_SlidingWindow_bytes32_uint256') + .withArgs(true); + + await expect(this.mock.state(key1)).to.eventually.deep.equal([17n, CAPACITY - 17n]); + await expect(this.mock.used(key1)).to.eventually.equal(17n); + await expect(this.mock.available(key1)).to.eventually.equal(CAPACITY - 17n); + await expect(this.mock.state(key2)).to.eventually.deep.equal([42n, CAPACITY - 42n]); + await expect(this.mock.used(key2)).to.eventually.equal(42n); + await expect(this.mock.available(key2)).to.eventually.equal(CAPACITY - 42n); + }); + + it('multiple tryConsume accumulate within the window', async function () { + await this.mock.tryConsume(100n); + await this.mock.tryConsume(200n); + await this.mock.tryConsume(50n); + + await expect(this.mock.state()).to.eventually.deep.equal([350n, CAPACITY - 350n]); + await expect(this.mock.used()).to.eventually.equal(350n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 350n); + }); + + it('tryConsume returns false and does not update state when over capacity', async function () { + // Static + await expect(this.mock.tryConsumeStatic(CAPACITY + 1n)).to.eventually.be.false; + + // execute (and get event) + await expect(this.mock.tryConsume(CAPACITY + 1n)) + .to.emit(this.mock, 'return$tryConsume_RateLimiter_SlidingWindow_bytes32_uint256') + .withArgs(false); + + await expect(this.mock.state()).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used()).to.eventually.equal(0n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY); + }); + + it('consume(0) is a no-op that returns true', async function () { + await this.mock.consume(42n); + + await expect(this.mock.state()).to.eventually.deep.equal([42n, CAPACITY - 42n]); + await expect(this.mock.used()).to.eventually.equal(42n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 42n); + + await expect(this.mock.tryConsume(0n)) + .to.emit(this.mock, 'return$tryConsume_RateLimiter_SlidingWindow_bytes32_uint256') + .withArgs(true); + + await expect(this.mock.state()).to.eventually.deep.equal([42n, CAPACITY - 42n]); + await expect(this.mock.used()).to.eventually.equal(42n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 42n); + }); + }); + + it('past consumptions drop out after the window elapses', async function () { + await this.mock.consume(CAPACITY / 2n); + + // step past the window + await time.increaseBy.timestamp(WINDOW + 1n); + + await expect(this.mock.state()).to.eventually.deep.equal([0n, CAPACITY]); + await expect(this.mock.used()).to.eventually.equal(0n); + await expect(this.mock.available()).to.eventually.equal(CAPACITY); + }); + + it('only consumptions outside the window drop out', async function () { + // first batch + await this.mock.consume(17n); + + // second batch, half-a-window later + await time.increaseBy.timestamp(WINDOW / 2n); + await this.mock.consume(42n); + + // both still within the window + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 59n); + await expect(this.mock.used()).to.eventually.equal(59n); // 17n + 42n + + // step forward so the first consume falls out but the second is still in-window + await time.increaseBy.timestamp(WINDOW / 2n + 5n); + + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 42n); + await expect(this.mock.used()).to.eventually.equal(42n); + }); + + it('updateSettings does not modify recorded history', async function () { + await this.mock.consume(42n); + + await expect(this.mock.available()).to.eventually.equal(CAPACITY - 42n); + await expect(this.mock.used()).to.eventually.equal(42n); + + await this.mock.updateSettings(WINDOW, CAPACITY * 2n); + + // used reflects existing checkpoints, only the new limit/window applies going forward + await expect(this.mock.available()).to.eventually.equal(2n * CAPACITY - 42n); + await expect(this.mock.used()).to.eventually.equal(42n); + }); + + it('reset clears the history', async function () { + await this.mock.consume(CAPACITY / 2n); + await this.mock.reset(); + + expect(await this.mock.used()).to.equal(0n); + expect(await this.mock.available()).to.equal(CAPACITY); + }); + + it('window saturation prevents underflow when block.timestamp < window', async function () { + // computing state with a window larger than block.timestamp must not revert + await this.mock.updateSettings(MAX_UINT48, CAPACITY); + await expect(this.mock.state()).to.not.be.reverted; + }); + }); + }); +});