optional lazy erasure from membership set

This commit is contained in:
Sergei Tikhomirov 2024-10-02 11:04:58 +02:00
parent bf70df2911
commit fb9e27975f
No known key found for this signature in database
GPG Key ID: 6A1F8ED9D6538027
3 changed files with 53 additions and 43 deletions

View File

@ -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

View File

@ -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;
}
}

View File

@ -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));