feat: allow choosing automatic membership erases and specific members…

…hip removal on register
waku-org · Sep 9, 2024 · 06b5162 · 06b5162
1 parent c66f728
commit 06b5162
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 6 deletions.
diff --git a/src/Membership.sol b/src/Membership.sol
@@ -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,14 +169,16 @@ 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(
         address _sender,
         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) {

diff --git a/src/WakuRlnV2.sol b/src/WakuRlnV2.sol
@@ -143,14 +143,41 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member
         uint32 userMessageLimit
     )
         external
-        payable
         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);
     }