diff --git a/src/Membership.sol b/src/Membership.sol index 4dc5a12..5bb78bb 100644 --- a/src/Membership.sol +++ b/src/Membership.sol @@ -22,6 +22,9 @@ error AttemptedExtensionByNonHolder(uint256 idCommitment); // This membership cannot be erased error CannotEraseMembership(uint256 idCommitment); +// This membership does not exist +error MembershipDoesNotExist(uint256 idCommitment); + abstract contract MembershipUpgradeable is Initializable { using SafeERC20 for IERC20; @@ -250,39 +253,39 @@ abstract contract MembershipUpgradeable is Initializable { /// @dev Erase expired memberships or owned grace-period memberships. /// @param _sender address of the sender of transaction (will be used to check memberships in grace period) /// @param _idCommitment idCommitment of the membership to erase - function _eraseMembershipAndSaveSlotToReuse( - address _sender, - uint256 _idCommitment, - MembershipInfo memory _membership - ) - internal - { - bool membershipExpired = _isExpired(_membership.gracePeriodStartTimestamp, _membership.gracePeriodDuration); + function _eraseMembershipLazily(address _sender, uint256 _idCommitment) internal returns (uint32 index) { + MembershipInfo memory membership = memberships[_idCommitment]; + + if (membership.rateLimit == 0) revert MembershipDoesNotExist(_idCommitment); + + bool membershipExpired = _isExpired(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration); bool membershipIsInGracePeriod = - _isInGracePeriod(_membership.gracePeriodStartTimestamp, _membership.gracePeriodDuration); - bool isHolder = (_membership.holder == _sender); + _isInGracePeriod(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration); + bool isHolder = (membership.holder == _sender); if (!membershipExpired && !(membershipIsInGracePeriod && isHolder)) { revert CannotEraseMembership(_idCommitment); } // Move deposit balance from the membership to be erased to holder deposit balance - depositsToWithdraw[_membership.holder][_membership.token] += _membership.depositAmount; + depositsToWithdraw[membership.holder][membership.token] += membership.depositAmount; // Deduct the rate limit of this membership from the total rate limit - currentTotalRateLimit -= _membership.rateLimit; + currentTotalRateLimit -= membership.rateLimit; // Mark this membership as reusable - indicesOfLazilyErasedMemberships.push(_membership.index); + indicesOfLazilyErasedMemberships.push(membership.index); // Erase this membership from the memberships mapping - // Note: the Merkle tree data will be erased when the index is reused delete memberships[_idCommitment]; if (membershipExpired) { - emit MembershipExpired(_idCommitment, _membership.rateLimit, _membership.index); + emit MembershipExpired(_idCommitment, membership.rateLimit, membership.index); } - emit MembershipErased(_idCommitment, _membership.rateLimit, _membership.index); + emit MembershipErased(_idCommitment, membership.rateLimit, membership.index); + + // The caller uses this index to erase data from the membership set (the Merkle tree) + return membership.index; } /// @notice Determine if a membership is in grace period now diff --git a/src/WakuRlnV2.sol b/src/WakuRlnV2.sol index 77da296..2340e91 100644 --- a/src/WakuRlnV2.sol +++ b/src/WakuRlnV2.sol @@ -47,7 +47,7 @@ contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, M /// @notice the modifier to check that the membership with this idCommitment doesn't already exist /// @param idCommitment The idCommitment of the membership modifier noDuplicateMembership(uint256 idCommitment) { - require(!membershipExists(idCommitment), "Duplicate idCommitment: membership already exists"); + require(!isInMembershipSet(idCommitment), "Duplicate idCommitment: membership already exists"); _; } @@ -105,10 +105,10 @@ contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, M return idCommitment != 0 && idCommitment < Q; } - /// @notice Checks if a membership exists + /// @notice Checks if a membership is in the membership set /// @param idCommitment The idCommitment of the membership - /// @return true if the membership exists, false otherwise - function membershipExists(uint256 idCommitment) public view returns (bool) { + /// @return true if the membership is in the membership set, false otherwise + function isInMembershipSet(uint256 idCommitment) public view returns (bool) { (,, uint256 rateCommitment) = getMembershipInfo(idCommitment); return rateCommitment != 0; } @@ -176,7 +176,7 @@ contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, M noDuplicateMembership(idCommitment) membershipSetNotFull { - _eraseMemberships(idCommitmentsToErase); + _eraseMemberships(idCommitmentsToErase, false); _register(idCommitment, rateLimit); } @@ -232,19 +232,26 @@ contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, M /// This function is also used to erase the user's own grace-period memberships. /// The user (i.e. the transaction sender) can then withdraw the deposited tokens. /// @param idCommitments list of idCommitments of the memberships to erase - function eraseMemberships(uint256[] calldata idCommitments) external { - _eraseMemberships(idCommitments); + /// @param eraseFromMembershipSet Indicates whether to erase membership data from the set + function eraseMemberships(uint256[] calldata idCommitments, bool eraseFromMembershipSet) external { + // Clean-up: erase memberships from membership array (free up the rate limit), + // and from the membership set (free up tree indices). + _eraseMemberships(idCommitments, eraseFromMembershipSet); } /// @dev Erase memberships from the list of idCommitments - /// @param idCommitmentsToErase The idCommitments to erase - function _eraseMemberships(uint256[] calldata idCommitmentsToErase) internal { + /// @param idCommitmentsToErase The idCommitments of memberships to erase from storage + /// @param eraseFromMembershipSet Indicates whether to erase membership data from the set + function _eraseMemberships(uint256[] calldata idCommitmentsToErase, bool eraseFromMembershipSet) internal { for (uint256 i = 0; i < idCommitmentsToErase.length; i++) { uint256 idCommitmentToErase = idCommitmentsToErase[i]; - MembershipInfo memory membershipToErase = memberships[idCommitmentToErase]; - if (membershipToErase.rateLimit == 0) revert InvalidIdCommitment(idCommitmentToErase); - _eraseMembershipAndSaveSlotToReuse(_msgSender(), idCommitmentToErase, membershipToErase); - LazyIMT.update(merkleTree, 0, membershipToErase.index); + // Erase the membership from the memberships array in contract storage + uint32 indexToErase = _eraseMembershipLazily(_msgSender(), idCommitmentToErase); + // Optionally, also erase the rate commitment data from the membership set (the Merkle tree) + // This does not affect index reusal for new membership registrations + if (eraseFromMembershipSet) { + LazyIMT.update(merkleTree, 0, indexToErase); + } } } @@ -283,16 +290,16 @@ contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, M } /// @notice Set the active duration for new memberships (terms of existing memberships don't change) - /// @param _activeDuration new active duration - function setActiveDuration(uint32 _activeDuration) external onlyOwner { - require(_activeDuration > 0); - activeDurationForNewMemberships = _activeDuration; + /// @param _activeDurationForNewMembership new active duration + function setActiveDuration(uint32 _activeDurationForNewMembership) external onlyOwner { + require(_activeDurationForNewMembership > 0); + activeDurationForNewMemberships = _activeDurationForNewMembership; } /// @notice Set the grace period for new memberships (grace periods of existing memberships don't change) - /// @param _gracePeriodDuration new grace period duration - function setGracePeriodDuration(uint32 _gracePeriodDuration) external onlyOwner { + /// @param _gracePeriodDurationForNewMembership new grace period duration + function setGracePeriodDuration(uint32 _gracePeriodDurationForNewMembership) external onlyOwner { // Note: grace period duration may be equal to zero - gracePeriodDurationForNewMemberships = _gracePeriodDuration; + gracePeriodDurationForNewMemberships = _gracePeriodDurationForNewMembership; } } diff --git a/test/WakuRlnV2.t.sol b/test/WakuRlnV2.t.sol index 094687a..354585d 100644 --- a/test/WakuRlnV2.t.sol +++ b/test/WakuRlnV2.t.sol @@ -47,7 +47,7 @@ contract WakuRlnV2Test is Test { w.register(idCommitment, membershipRateLimit); vm.pauseGasMetering(); assertEq(w.nextFreeIndex(), 1); - assertEq(w.membershipExists(idCommitment), true); + assertEq(w.isInMembershipSet(idCommitment), true); (,,,, uint32 fetchedMembershipRateLimit, uint32 index, address holder,) = w.memberships(idCommitment); assertEq(fetchedMembershipRateLimit, membershipRateLimit); assertEq(holder, address(this)); @@ -102,7 +102,7 @@ contract WakuRlnV2Test is Test { vm.assume(w.isValidMembershipRateLimit(membershipRateLimit)); vm.resumeGasMetering(); - assertEq(w.membershipExists(idCommitment), false); + assertEq(w.isInMembershipSet(idCommitment), false); token.approve(address(w), price); w.register(idCommitment, membershipRateLimit); uint256 rateCommitment = PoseidonT3.hash([idCommitment, membershipRateLimit]); @@ -434,7 +434,7 @@ contract WakuRlnV2Test is Test { // time travel to the moment we can erase all expired memberships uint256 membershipExpirationTimestamp = w.membershipExpirationTimestamp(idCommitmentsLength); vm.warp(membershipExpirationTimestamp); - w.eraseMemberships(commitmentsToErase); + w.eraseMemberships(commitmentsToErase, false); // Verify that expired indices match what we expect for (uint32 i = 0; i < idCommitmentsLength; i++) { @@ -506,7 +506,7 @@ contract WakuRlnV2Test is Test { emit MembershipUpgradeable.MembershipExpired(commitmentsToErase[0], 0, 0); vm.expectEmit(true, false, false, false); // only check the first parameter of the event (the idCommitment) emit MembershipUpgradeable.MembershipExpired(commitmentsToErase[0], 0, 0); - w.eraseMemberships(commitmentsToErase); + w.eraseMemberships(commitmentsToErase, false); address holder; @@ -523,7 +523,7 @@ contract WakuRlnV2Test is Test { commitmentsToErase[0] = idCommitment; commitmentsToErase[1] = idCommitment + 4; vm.expectRevert(abi.encodeWithSelector(CannotEraseMembership.selector, idCommitment + 4)); - w.eraseMemberships(commitmentsToErase); + w.eraseMemberships(commitmentsToErase, false); } function test__RemoveAllExpiredMemberships(uint32 idCommitmentsLength) external { @@ -554,7 +554,7 @@ contract WakuRlnV2Test is Test { emit MembershipUpgradeable.MembershipExpired(i + 1, 0, 0); } - w.eraseMemberships(commitmentsToErase); + w.eraseMemberships(commitmentsToErase, false); // Erased memberships are gone! for (uint256 i = 0; i < commitmentsToErase.length; i++) { @@ -586,7 +586,7 @@ contract WakuRlnV2Test is Test { uint256[] memory commitmentsToErase = new uint256[](1); commitmentsToErase[0] = idCommitment; - w.eraseMemberships(commitmentsToErase); + w.eraseMemberships(commitmentsToErase, false); uint256 availableBalance = w.depositsToWithdraw(address(this), address(token));