feat: allow choosing automatic membership erases and specific membership removal on register

This commit is contained in:
Richard Ramos 2024-09-09 17:54:29 -04:00
parent c66f72897b
commit 12e837a1e2
No known key found for this signature in database
GPG Key ID: 1CE87DB518195760
2 changed files with 40 additions and 13 deletions

View File

@ -138,18 +138,21 @@ contract Membership {
/// @param _sender address of the owner of the new membership
/// @param _idCommitment the idcommitment of the new membership
/// @param _rateLimit the user message limit
/// @param _eraseIfNeeded Erase expired memberships if the `_rateLimit` exceeds the available rate limit
/// @return index the index in the merkle tree
/// @return reusedIndex indicates whether a new leaf is being used or if using an existing leaf in the merkle tree
function _acquireMembership(
address _sender,
uint256 _idCommitment,
uint32 _rateLimit
uint32 _rateLimit,
bool _eraseIfNeeded
)
internal
returns (uint32 index, bool reusedIndex)
{
(address token, uint256 amount) = priceCalculator.calculate(_rateLimit);
(index, reusedIndex) = _setupMembershipDetails(_sender, _idCommitment, _rateLimit, token, amount);
(index, reusedIndex) =
_setupMembershipDetails(_sender, _idCommitment, _rateLimit, token, amount, _eraseIfNeeded);
_transferFees(_sender, token, amount);
}
@ -166,6 +169,7 @@ contract Membership {
/// @param _rateLimit User message limit
/// @param _token Address of the token used to acquire the membership
/// @param _amount Amount of the token used to acquire the membership
/// @param _eraseIfNeeded Erase expired memberships if the `_rateLimit` exceeds the available rate limit
/// @return index membership index on the merkle tree
/// @return reusedIndex indicates whether the index returned was a reused slot on the tree or not
function _setupMembershipDetails(
@ -173,7 +177,8 @@ contract Membership {
uint256 _idCommitment,
uint32 _rateLimit,
address _token,
uint256 _amount
uint256 _amount,
bool _eraseIfNeeded
)
internal
returns (uint32 index, bool reusedIndex)
@ -191,7 +196,8 @@ contract Membership {
// Determine if we exceed the total rate limit
if (_totalRateLimitPerEpoch + _rateLimit > _maxTotalRateLimitPerEpoch) {
if (_head == 0) revert ExceedAvailableMaxRateLimitPerEpoch(); // List is empty
if (_head == 0 || !_eraseIfNeeded) revert ExceedAvailableMaxRateLimitPerEpoch(); // List is empty or can't
// erase memberships automatically
// Attempt to free expired membership slots
while (_totalRateLimitPerEpoch + _rateLimit > _maxTotalRateLimitPerEpoch && _head != 0) {

View File

@ -138,19 +138,40 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member
/// @notice Allows a user to register as a member
/// @param idCommitment The idCommitment of the member
/// @param userMessageLimit The message limit of the member
function register(
uint256 idCommitment,
uint32 userMessageLimit
)
external
payable
onlyValidIdCommitment(idCommitment)
{
function register(uint256 idCommitment, uint32 userMessageLimit) external onlyValidIdCommitment(idCommitment) {
if (memberExists(idCommitment)) revert DuplicateIdCommitment();
uint32 index;
bool reusedIndex;
(index, reusedIndex) = _acquireMembership(_msgSender(), idCommitment, userMessageLimit);
(index, reusedIndex) = _acquireMembership(_msgSender(), idCommitment, userMessageLimit, true);
_register(idCommitment, userMessageLimit, index, reusedIndex);
}
/// @notice Allows a user to register as a member
/// @param idCommitment The idCommitment of the member
/// @param userMessageLimit The message limit of the member
/// @param membershipsToErase List of expired idCommitments to erase
function register(
uint256 idCommitment,
uint32 userMessageLimit,
uint256[] calldata membershipsToErase
)
external
onlyValidIdCommitment(idCommitment)
{
if (memberExists(idCommitment)) revert DuplicateIdCommitment();
for (uint256 i = 0; i < membershipsToErase.length; i++) {
uint256 idCommitmentToErase = membershipsToErase[i];
MembershipInfo memory mdetails = members[idCommitmentToErase];
_eraseMembership(_msgSender(), idCommitmentToErase, mdetails);
LazyIMT.update(imtData, 0, mdetails.index);
}
uint32 index;
bool reusedIndex;
(index, reusedIndex) = _acquireMembership(_msgSender(), idCommitment, userMessageLimit, false);
_register(idCommitment, userMessageLimit, index, reusedIndex);
}