-
Notifications
You must be signed in to change notification settings - Fork 12.4k
Add RateLimiter library #6490
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Add RateLimiter library #6490
Changes from all commits
60657c9
6f86b94
1ea03ca
79a8cfe
c78be90
83d639e
88c8736
17a103e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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. |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; | ||
|
Comment on lines
+186
to
+187
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This is necessary to ensure just one sload is used?
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. that is the goal. |
||
|
|
||
| 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. | ||
| * | ||
|
Amxx marked this conversation as resolved.
|
||
| * 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; | ||
| } | ||
| } | ||
|
Comment on lines
+39
to
+264
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🛠️ Refactor suggestion | 🟠 Major | 🏗️ Heavy lift Add targeted tests for boundary semantics before merge. Given the amount of new stateful logic, please add tests for: zero-window behavior, boundary-at- 🤖 Prompt for AI Agents |
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Should we have a create bucket function? I feel that its not obvious what
lastUsedandlastTimepointshould be. Maybe there should be a way to specify if a bucket is full or empty at creation.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The
updateSettingsis what you should use to setup intial config. Maybe the name is not very clear.