Skip to content

Commit

Permalink
refactor: unique register function
Browse files Browse the repository at this point in the history
  • Loading branch information
@richard-ramos
richard-ramos committed Oct 7, 2024
1 parent 66a1a6c commit b78a4c0
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 46 deletions.
1 change: 0 additions & 1 deletion script/Deploy.s.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand Down
15 changes: 0 additions & 15 deletions src/WakuRlnV2.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
62 changes: 32 additions & 30 deletions test/WakuRlnV2.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ contract WakuRlnV2Test is Test {

address internal deployer;

uint256[] noCommitments = new uint256[](0);

function setUp() public virtual {
token = new TestToken();

Expand All @@ -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);
Expand Down Expand Up @@ -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) =
Expand Down Expand Up @@ -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 {
Expand All @@ -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 {
Expand All @@ -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 {
Expand All @@ -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 {
Expand All @@ -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));
Expand Down Expand Up @@ -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);
Expand All @@ -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);

Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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));
}
Expand All @@ -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);
Expand Down Expand Up @@ -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 {
Expand All @@ -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;
Expand All @@ -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
Expand All @@ -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());
Expand All @@ -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);
}
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand All @@ -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]);
Expand All @@ -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();

Expand Down

0 comments on commit b78a4c0

Please sign in to comment.