From b78a4c02cb13b4a597e7022e0359d85d1bfb32b5 Mon Sep 17 00:00:00 2001 From: Richard Ramos Date: Mon, 7 Oct 2024 16:39:13 -0400 Subject: [PATCH] refactor: unique register function --- script/Deploy.s.sol | 1 - src/WakuRlnV2.sol | 15 ----------- test/WakuRlnV2.t.sol | 62 +++++++++++++++++++++++--------------------- 3 files changed, 32 insertions(+), 46 deletions(-) diff --git a/script/Deploy.s.sol b/script/Deploy.s.sol index ed9c70e..cc8b1ba 100644 --- a/script/Deploy.s.sol +++ b/script/Deploy.s.sol @@ -6,7 +6,6 @@ import { LinearPriceCalculator } from "../src/LinearPriceCalculator.sol"; import { PoseidonT3 } from "poseidon-solidity/PoseidonT3.sol"; import { LazyIMT } from "@zk-kit/imt.sol/LazyIMT.sol"; import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; -import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol"; import { BaseScript } from "./Base.s.sol"; diff --git a/src/WakuRlnV2.sol b/src/WakuRlnV2.sol index efaed97..965bbaa 100644 --- a/src/WakuRlnV2.sol +++ b/src/WakuRlnV2.sol @@ -154,21 +154,6 @@ contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, M return merkleTree.elements[LazyIMT.indexForElement(0, index)]; } - /// @notice Register a membership - /// @param idCommitment The idCommitment of the new membership - /// @param rateLimit The rate limit of the new membership - function register( - uint256 idCommitment, - uint32 rateLimit - ) - external - onlyValidIdCommitment(idCommitment) - noDuplicateMembership(idCommitment) - membershipSetNotFull - { - _register(idCommitment, rateLimit); - } - /// @notice Register a membership while erasing some expired memberships to reuse their rate limit /// @param idCommitment The idCommitment of the new membership /// @param rateLimit The rate limit of the new membership diff --git a/test/WakuRlnV2.t.sol b/test/WakuRlnV2.t.sol index 1f9e6f5..cacd7a8 100644 --- a/test/WakuRlnV2.t.sol +++ b/test/WakuRlnV2.t.sol @@ -22,6 +22,8 @@ contract WakuRlnV2Test is Test { address internal deployer; + uint256[] noCommitments = new uint256[](0); + function setUp() public virtual { token = new TestToken(); @@ -44,7 +46,7 @@ contract WakuRlnV2Test is Test { (, uint256 price) = w.priceCalculator().calculate(membershipRateLimit); vm.resumeGasMetering(); token.approve(address(w), price); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); vm.pauseGasMetering(); assertEq(w.nextFreeIndex(), 1); assertEq(w.isInMembershipSet(idCommitment), true); @@ -104,7 +106,7 @@ contract WakuRlnV2Test is Test { assertEq(w.isInMembershipSet(idCommitment), false); token.approve(address(w), price); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); uint256 rateCommitment = PoseidonT3.hash([idCommitment, membershipRateLimit]); (uint32 fetchedMembershipRateLimit, uint32 index, uint256 fetchedRateCommitment) = @@ -137,7 +139,7 @@ contract WakuRlnV2Test is Test { token.approve(address(w), price - 1); vm.expectRevert(bytes("ERC20: insufficient allowance")); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); } function test__IdCommitmentToMetadata__DoesntExist() external view { @@ -157,7 +159,7 @@ contract WakuRlnV2Test is Test { token.approve(address(w), price); vm.expectRevert(abi.encodeWithSelector(InvalidIdCommitment.selector, 0)); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); } function test__InvalidRegistration__InvalidIdCommitment__LargerThanField() external { @@ -169,7 +171,7 @@ contract WakuRlnV2Test is Test { uint256 idCommitment = w.Q() + 1; token.approve(address(w), price); vm.expectRevert(abi.encodeWithSelector(InvalidIdCommitment.selector, idCommitment)); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); } function test__InvalidRegistration__InvalidMembershipRateLimit__MinMax() external { @@ -179,10 +181,10 @@ contract WakuRlnV2Test is Test { uint32 invalidMax = w.maxMembershipRateLimit() + 1; vm.expectRevert(abi.encodeWithSelector(InvalidMembershipRateLimit.selector)); - w.register(idCommitment, invalidMin); + w.register(idCommitment, invalidMin, noCommitments); vm.expectRevert(abi.encodeWithSelector(InvalidMembershipRateLimit.selector)); - w.register(idCommitment, invalidMax); + w.register(idCommitment, invalidMax, noCommitments); } function test__ValidRegistrationExtend(uint32 membershipRateLimit) external { @@ -196,7 +198,7 @@ contract WakuRlnV2Test is Test { vm.resumeGasMetering(); token.approve(address(w), price); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); (,, uint256 gracePeriodStartTimestamp,,,,,) = w.memberships(idCommitment); assertFalse(w.isInGracePeriod(idCommitment)); @@ -236,7 +238,7 @@ contract WakuRlnV2Test is Test { // Attempt to extend a non grace period membership token.approve(address(w), price); - w.register(idCommitment + 1, membershipRateLimit); + w.register(idCommitment + 1, membershipRateLimit, noCommitments); commitmentsToExtend[0] = idCommitment + 1; vm.expectRevert(abi.encodeWithSelector(CannotExtendActiveMembership.selector, commitmentsToExtend[0])); w.extendMemberships(commitmentsToExtend); @@ -258,7 +260,7 @@ contract WakuRlnV2Test is Test { vm.resumeGasMetering(); token.approve(address(w), price); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); (,, uint256 gracePeriodStartTimestamp, uint32 gracePeriodDuration,,,,) = w.memberships(idCommitment); @@ -289,7 +291,7 @@ contract WakuRlnV2Test is Test { vm.resumeGasMetering(); token.approve(address(w), price); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); uint256 ogExpirationTimestamp = w.membershipExpirationTimestamp(idCommitment); (,, uint256 gracePeriodStartTimestamp,,,,,) = w.memberships(idCommitment); @@ -321,7 +323,7 @@ contract WakuRlnV2Test is Test { vm.resumeGasMetering(); token.approve(address(w), price); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); (,, uint256 fetchedgracePeriodStartTimestamp, uint32 fetchedGracePeriod,,,,) = w.memberships(idCommitment); @@ -349,7 +351,7 @@ contract WakuRlnV2Test is Test { for (uint256 i = 1; i <= 5; i++) { token.approve(address(w), priceA); - w.register(i, 20); + w.register(i, 20, noCommitments); // Make sure they're expired vm.warp(w.membershipExpirationTimestamp(i)); } @@ -368,7 +370,7 @@ contract WakuRlnV2Test is Test { // Should fail. There's not enough free rate limit vm.expectRevert(abi.encodeWithSelector(CannotExceedMaxTotalRateLimit.selector)); - w.register(6, 60); + w.register(6, 60, noCommitments); // Attempt to erase 3 memberships including one that can't be erased (the last one) uint256[] memory commitmentsToErase = new uint256[](3); @@ -426,33 +428,33 @@ contract WakuRlnV2Test is Test { (, uint256 price) = w.priceCalculator().calculate(membershipRateLimit); token.approve(address(w), price); vm.expectRevert(abi.encodeWithSelector(InvalidMembershipRateLimit.selector)); - w.register(1, membershipRateLimit); + w.register(1, membershipRateLimit, noCommitments); // Should register succesfully membershipRateLimit = 4; (, price) = w.priceCalculator().calculate(membershipRateLimit); token.approve(address(w), price); - w.register(2, membershipRateLimit); + w.register(2, membershipRateLimit, noCommitments); // Exceeds the rate limit membershipRateLimit = 2; (, price) = w.priceCalculator().calculate(membershipRateLimit); token.approve(address(w), price); vm.expectRevert(abi.encodeWithSelector(CannotExceedMaxTotalRateLimit.selector)); - w.register(3, membershipRateLimit); + w.register(3, membershipRateLimit, noCommitments); // Should register succesfully membershipRateLimit = 1; (, price) = w.priceCalculator().calculate(membershipRateLimit); token.approve(address(w), price); - w.register(3, membershipRateLimit); + w.register(3, membershipRateLimit, noCommitments); // We ran out of rate limit again membershipRateLimit = 1; (, price) = w.priceCalculator().calculate(membershipRateLimit); token.approve(address(w), price); vm.expectRevert(abi.encodeWithSelector(CannotExceedMaxTotalRateLimit.selector)); - w.register(4, membershipRateLimit); + w.register(4, membershipRateLimit, noCommitments); } function test__indexReuse_eraseMemberships(uint32 idCommitmentsLength) external { @@ -463,7 +465,7 @@ contract WakuRlnV2Test is Test { uint256[] memory commitmentsToErase = new uint256[](idCommitmentsLength); for (uint256 i = 1; i <= idCommitmentsLength; i++) { token.approve(address(w), price); - w.register(i, 20); + w.register(i, 20, noCommitments); (,,,,, index,,) = w.memberships(i); assertEq(index, w.nextFreeIndex() - 1); commitmentsToErase[i - 1] = i; @@ -485,7 +487,7 @@ contract WakuRlnV2Test is Test { uint256 expectedindexReusedPos = idCommitmentsLength - i; uint32 expectedReusedIndex = w.indicesOfLazilyErasedMemberships(expectedindexReusedPos); token.approve(address(w), price); - w.register(idCommitment, 20); + w.register(idCommitment, 20, noCommitments); (,,,,, index,,) = w.memberships(idCommitment); assertEq(expectedReusedIndex, index); // Should have been removed from the list @@ -501,7 +503,7 @@ contract WakuRlnV2Test is Test { // Should use a new index since we got rid of all reusable indexes token.approve(address(w), price); - w.register(100, 20); + w.register(100, 20, noCommitments); (,,,,, index,,) = w.memberships(100); assertEq(index, expectedNextFreeIndex); assertEq(expectedNextFreeIndex + 1, w.nextFreeIndex()); @@ -520,7 +522,7 @@ contract WakuRlnV2Test is Test { uint256 time = block.timestamp; for (uint256 i = 0; i < 5; i++) { token.approve(address(w), price); - w.register(idCommitment + i, membershipRateLimit); + w.register(idCommitment + i, membershipRateLimit, noCommitments); time += 100; vm.warp(time); } @@ -574,7 +576,7 @@ contract WakuRlnV2Test is Test { uint256 time = block.timestamp; for (uint256 i = 1; i <= idCommitmentsLength; i++) { token.approve(address(w), price); - w.register(i, membershipRateLimit); + w.register(i, membershipRateLimit, noCommitments); time += 100; vm.warp(time); } @@ -616,7 +618,7 @@ contract WakuRlnV2Test is Test { vm.resumeGasMetering(); token.approve(address(w), price); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); (,, uint256 gracePeriodStartTimestamp,,,,,) = w.memberships(idCommitment); @@ -651,11 +653,11 @@ contract WakuRlnV2Test is Test { vm.resumeGasMetering(); token.approve(address(w), price); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); token.approve(address(w), price); vm.expectRevert(bytes("Duplicate idCommitment: membership already exists")); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); } function test__InvalidRegistration__FullTree() external { @@ -690,7 +692,7 @@ contract WakuRlnV2Test is Test { vm.store(address(w), bytes32(uint256(256)), 0x0000000000000000000000000000000000000000000000000000000000100000); token.approve(address(w), price); vm.expectRevert(bytes("Membership set is full")); - w.register(1, membershipRateLimit); + w.register(1, membershipRateLimit, noCommitments); } function test__InvalidPaginationQuery__StartIndexGTEndIndex() external { @@ -711,7 +713,7 @@ contract WakuRlnV2Test is Test { vm.resumeGasMetering(); token.approve(address(w), price); - w.register(idCommitment, membershipRateLimit); + w.register(idCommitment, membershipRateLimit, noCommitments); uint256[] memory commitments = w.getRateCommitmentsInRangeBoundsInclusive(0, 0); assertEq(commitments.length, 1); uint256 rateCommitment = PoseidonT3.hash([idCommitment, membershipRateLimit]); @@ -726,7 +728,7 @@ contract WakuRlnV2Test is Test { for (uint256 i = 0; i <= idCommitmentsLength; i++) { token.approve(address(w), price); - w.register(i + 1, membershipRateLimit); + w.register(i + 1, membershipRateLimit, noCommitments); } vm.resumeGasMetering();