diff --git a/src/Membership.sol b/src/Membership.sol index 52d316c..c1be2e5 100644 --- a/src/Membership.sol +++ b/src/Membership.sol @@ -6,6 +6,8 @@ import "openzeppelin-contracts/contracts/token/ERC20/IERC20.sol"; import "openzeppelin-contracts/contracts/token/ERC20/utils/SafeERC20.sol"; import "openzeppelin-contracts/contracts/utils/Context.sol"; +import "forge-std/console.sol"; + // The number of periods should be greater than zero error NumberOfPeriodsCantBeZero(); @@ -20,7 +22,7 @@ error InvalidRateLimit(); // It's not possible to acquire the rate limit due to exceeding the expected limits // even after attempting to erase expired memberships -error ExceedMaxRateLimitPerEpoch(); +error ExceedAvailableMaxRateLimitPerEpoch(); // This membership is not in grace period yet error NotInGracePeriod(uint256 idCommitment); @@ -80,7 +82,7 @@ contract Membership { uint256 next; /// @notice amount of the token used to acquire this membership uint256 amount; - /// @notice numPeriods + /// @notice numberOfPeriods uint32 numberOfPeriods; /// @notice timestamp of when the grace period starts for this membership uint256 gracePeriodStartDate; @@ -202,36 +204,44 @@ contract Membership { revert InvalidRateLimit(); } - // Attempt to free expired membership slots - while (totalRateLimitPerEpoch + _rateLimit > maxTotalRateLimitPerEpoch) { - // Determine if there are any available spot in the membership map - // by looking at the oldest membership. If it's expired, we can free it - MembershipInfo memory oldestMembership = members[head]; - - if ( - oldestMembership.holder != address(0) // membership has a holder - && isExpired(oldestMembership.gracePeriodStartDate) - ) { - emit MemberExpired(head, oldestMembership.userMessageLimit, oldestMembership.index); - - // Deduct the expired membership rate limit - totalRateLimitPerEpoch -= oldestMembership.userMessageLimit; - - // Promote the next oldest membership to oldest - uint256 nextOldest = oldestMembership.next; - head = nextOldest; - if (nextOldest != 0) { - members[nextOldest].prev = 0; + // Determine if we exceed the total rate limit + if (totalRateLimitPerEpoch + _rateLimit > maxTotalRateLimitPerEpoch) { + if (head == 0) revert ExceedAvailableMaxRateLimitPerEpoch(); // List is empty + + // Attempt to free expired membership slots + while (totalRateLimitPerEpoch + _rateLimit > maxTotalRateLimitPerEpoch) { + // Determine if there are any available spot in the membership map + // by looking at the oldest membership. If it's expired, we can free it + MembershipInfo memory oldestMembership = members[head]; + if ( + oldestMembership.holder != address(0) // membership has a holder + && _isExpired( + oldestMembership.gracePeriodStartDate, + oldestMembership.gracePeriod, + oldestMembership.numberOfPeriods + ) + ) { + emit MemberExpired(head, oldestMembership.userMessageLimit, oldestMembership.index); + + // Deduct the expired membership rate limit + totalRateLimitPerEpoch -= oldestMembership.userMessageLimit; + + // Promote the next oldest membership to oldest + uint256 nextOldest = oldestMembership.next; + head = nextOldest; + if (nextOldest != 0) { + members[nextOldest].prev = 0; + } + + // Move balance from expired membership to holder balance + balancesToWithdraw[oldestMembership.holder][oldestMembership.token] += oldestMembership.amount; + + availableExpiredIndices.push(oldestMembership.index); + + delete members[head]; + } else { + revert ExceedAvailableMaxRateLimitPerEpoch(); } - - // Move balance from expired membership to holder balance - balancesToWithdraw[oldestMembership.holder][oldestMembership.token] += oldestMembership.amount; - - availableExpiredIndices.push(oldestMembership.index); - - delete members[head]; - } else { - revert ExceedMaxRateLimitPerEpoch(); } } @@ -242,7 +252,6 @@ contract Membership { prev = tail; } else { // First item - // TODO: test adding memberships after the list has been emptied head = _idCommitment; } diff --git a/src/WakuRlnV2.sol b/src/WakuRlnV2.sol index 19c8683..69c23a7 100644 --- a/src/WakuRlnV2.sol +++ b/src/WakuRlnV2.sol @@ -77,6 +77,11 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member public initializer { + require(_maxTotalRateLimitPerEpoch >= maxRateLimitPerMembership); + require(_maxRateLimitPerMembership > minRateLimitPerMembership); + require(_minRateLimitPerMembership > 0); + require(_billingPeriod > 0); + __Ownable_init(); __UUPSUpgradeable_init(); __Membership_init( @@ -244,24 +249,28 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member /// @notice Set the maximum total rate limit of all memberships in the tree /// @param _maxTotalRateLimitPerEpoch new value function setMaxTotalRateLimitPerEpoch(uint32 _maxTotalRateLimitPerEpoch) external onlyOwner { + require(_maxTotalRateLimitPerEpoch >= maxRateLimitPerMembership); maxTotalRateLimitPerEpoch = _maxTotalRateLimitPerEpoch; } /// @notice Set the maximum rate limit of one membership /// @param _maxRateLimitPerMembership new value function setMaxRateLimitPerMembership(uint32 _maxRateLimitPerMembership) external onlyOwner { + require(_maxRateLimitPerMembership >= minRateLimitPerMembership); maxRateLimitPerMembership = _maxRateLimitPerMembership; } /// @notice Set the minimum rate limit of one membership /// @param _minRateLimitPerMembership new value function setMinRateLimitPerMembership(uint32 _minRateLimitPerMembership) external onlyOwner { + require(_minRateLimitPerMembership > 0); minRateLimitPerMembership = _minRateLimitPerMembership; } /// @notice Set the membership billing period /// @param _billingPeriod new value function setBillingPeriod(uint32 _billingPeriod) external onlyOwner { + require(_billingPeriod > 0); billingPeriod = _billingPeriod; } diff --git a/test/WakuRlnV2.t.sol b/test/WakuRlnV2.t.sol index 245826a..86467a4 100644 --- a/test/WakuRlnV2.t.sol +++ b/test/WakuRlnV2.t.sol @@ -363,15 +363,51 @@ contract WakuRlnV2Test is Test { } function test__RegistrationWhenMaxRateLimitIsReached() external { - // TODO: implement - // TODO: validate elements are chained correctly - // TODO: validate reuse of index + vm.pauseGasMetering(); + vm.startPrank(w.owner()); + w.setMinRateLimitPerMembership(1); + w.setMaxRateLimitPerMembership(5); + w.setMaxTotalRateLimitPerEpoch(5); + vm.stopPrank(); + vm.resumeGasMetering(); + + bool isValid = w.isValidUserMessageLimit(6); + assertFalse(isValid); + + // Exceeds the max rate limit per user + uint32 userMessageLimit = 10; + (, uint256 price) = w.priceCalculator().calculate(userMessageLimit, 1); + vm.expectRevert(abi.encodeWithSelector(InvalidRateLimit.selector)); + w.register{ value: price }(1, userMessageLimit, 1); + + // Should register succesfully + userMessageLimit = 4; + (, price) = w.priceCalculator().calculate(userMessageLimit, 1); + w.register{ value: price }(2, userMessageLimit, 1); + + // Exceeds the rate limit + userMessageLimit = 2; + (, price) = w.priceCalculator().calculate(userMessageLimit, 1); + vm.expectRevert(abi.encodeWithSelector(ExceedAvailableMaxRateLimitPerEpoch.selector)); + w.register{ value: price }(3, userMessageLimit, 1); + + // Should register succesfully + userMessageLimit = 1; + (, price) = w.priceCalculator().calculate(userMessageLimit, 1); + w.register{ value: price }(3, userMessageLimit, 1); + + // We ran out of rate limit again + userMessageLimit = 1; + (, price) = w.priceCalculator().calculate(userMessageLimit, 1); + vm.expectRevert(abi.encodeWithSelector(ExceedAvailableMaxRateLimitPerEpoch.selector)); + w.register{ value: price }(4, userMessageLimit, 1); } function test__RegistrationWhenMaxRateLimitIsReachedAndSingleExpiredMemberAvailable() external { // TODO: implement // TODO: validate elements are chained correctly // TODO: validate reuse of index + // TODO: validate that expired event is emitted // TODO: validate balance }