fix: code review

.gas-snapshot

@@ -1,29 +1,28 @@
-WakuRlnV2Test:test__IdCommitmentToMetadata__DoesntExist() (gas: 27535)
-WakuRlnV2Test:test__InsertionNormalOrder(uint32) (runs: 1004, μ: 1093558, ~: 496185)
-WakuRlnV2Test:test__InvalidETHAmount(uint256,uint32) (runs: 1004, μ: 322023, ~: 322023)
-WakuRlnV2Test:test__InvalidPaginationQuery__EndIndexGTcommitmentIndex() (gas: 18378)
-WakuRlnV2Test:test__InvalidPaginationQuery__StartIndexGTEndIndex() (gas: 16225)
-WakuRlnV2Test:test__InvalidRegistration__DuplicateIdCommitment() (gas: 249303)
-WakuRlnV2Test:test__InvalidRegistration__FullTree() (gas: 162941)
-WakuRlnV2Test:test__InvalidRegistration__InvalidIdCommitment__LargerThanField() (gas: 13471)
-WakuRlnV2Test:test__InvalidRegistration__InvalidIdCommitment__Zero() (gas: 12182)
-WakuRlnV2Test:test__InvalidRegistration__InvalidUserMessageLimit__MinMax() (gas: 63423)
-WakuRlnV2Test:test__LinearPriceCalculation(uint32) (runs: 1015, μ: 25964, ~: 25964)
-WakuRlnV2Test:test__RegistrationWhenMaxRateLimitIsReached() (gas: 547839)
-WakuRlnV2Test:test__RegistrationWhenMaxRateLimitIsReachedAndMultipleExpiredMembersAvailable() (gas: 739752)
-WakuRlnV2Test:test__RegistrationWhenMaxRateLimitIsReachedAndSingleExpiredMemberAvailable() (gas: 388429)
-WakuRlnV2Test:test__RegistrationWhenMaxRateLimitReachedAndMultipleExpiredMembersAvailableWithoutEnoughRateLimit() (gas: 748699)
-WakuRlnV2Test:test__RegistrationWithTokens(uint256,uint32) (runs: 1004, μ: 318681, ~: 318681)
-WakuRlnV2Test:test__RemoveAllExpiredMemberships(uint32) (runs: 1004, μ: 6262152, ~: 1315256)
-WakuRlnV2Test:test__RemoveExpiredMemberships(uint32) (runs: 1003, μ: 1075445, ~: 1075445)
-WakuRlnV2Test:test__Upgrade() (gas: 6990368)
-WakuRlnV2Test:test__ValidPaginationQuery(uint32) (runs: 1005, μ: 222425, ~: 53007)
-WakuRlnV2Test:test__ValidPaginationQuery__OneElement() (gas: 262398)
-WakuRlnV2Test:test__ValidRegistration(uint32) (runs: 1003, μ: 267715, ~: 267715)
-WakuRlnV2Test:test__ValidRegistrationExpiry(uint32) (runs: 1003, μ: 1280765, ~: 1280765)
-WakuRlnV2Test:test__ValidRegistrationExtend(uint32) (runs: 1003, μ: 1123986, ~: 1123986)
-WakuRlnV2Test:test__ValidRegistrationExtendSingleMembership(uint32) (runs: 1003, μ: 257339, ~: 257339)
-WakuRlnV2Test:test__ValidRegistration__kats() (gas: 238726)
-WakuRlnV2Test:test__WithdrawETH(uint32) (runs: 1003, μ: 243355, ~: 243356)
-WakuRlnV2Test:test__WithdrawToken(uint32) (runs: 1003, μ: 314604, ~: 314605)
-WakuRlnV2Test:test__indexReuse_eraseMemberships(uint32) (runs: 1004, μ: 2187085, ~: 964428)
+WakuRlnV2Test:test__IdCommitmentToMetadata__DoesntExist() (gas: 27547)
+WakuRlnV2Test:test__InsertionNormalOrder(uint32) (runs: 1005, μ: 1284765, ~: 577056)
+WakuRlnV2Test:test__InvalidPaginationQuery__EndIndexGTnextCommitmentIndex() (gas: 18290)
+WakuRlnV2Test:test__InvalidPaginationQuery__StartIndexGTEndIndex() (gas: 16181)
+WakuRlnV2Test:test__InvalidRegistration__DuplicateIdCommitment() (gas: 322752)
+WakuRlnV2Test:test__InvalidRegistration__FullTree() (gas: 239810)
+WakuRlnV2Test:test__InvalidRegistration__InvalidIdCommitment__LargerThanField() (gas: 36509)
+WakuRlnV2Test:test__InvalidRegistration__InvalidIdCommitment__Zero() (gas: 35220)
+WakuRlnV2Test:test__InvalidRegistration__InvalidUserMessageLimit__MinMax() (gas: 63679)
+WakuRlnV2Test:test__InvalidTokenAmount(uint256,uint32) (runs: 1004, μ: 207863, ~: 207863)
+WakuRlnV2Test:test__LinearPriceCalculation(uint32) (runs: 1015, μ: 26049, ~: 26049)
+WakuRlnV2Test:test__RegistrationWhenMaxRateLimitIsReached() (gas: 638698)
+WakuRlnV2Test:test__RegistrationWhenMaxRateLimitIsReachedAndMultipleExpiredMembersAvailable() (gas: 906196)
+WakuRlnV2Test:test__RegistrationWhenMaxRateLimitIsReachedAndSingleExpiredMemberAvailable() (gas: 461585)
+WakuRlnV2Test:test__RegistrationWhenMaxRateLimitReachedAndMultipleExpiredMembersAvailableWithoutEnoughRateLimit() (gas: 869622)
+WakuRlnV2Test:test__RemoveAllExpiredMemberships(uint32) (runs: 1005, μ: 7359251, ~: 1585880)
+WakuRlnV2Test:test__RemoveExpiredMemberships(uint32) (runs: 1003, μ: 1248285, ~: 1248286)
+WakuRlnV2Test:test__Upgrade() (gas: 7482875)
+WakuRlnV2Test:test__ValidPaginationQuery(uint32) (runs: 1006, μ: 227626, ~: 52991)
+WakuRlnV2Test:test__ValidPaginationQuery__OneElement() (gas: 319333)
+WakuRlnV2Test:test__ValidRegistration(uint32) (runs: 1003, μ: 325571, ~: 325571)
+WakuRlnV2Test:test__ValidRegistrationExpiry(uint32) (runs: 1003, μ: 1456074, ~: 1456074)
+WakuRlnV2Test:test__ValidRegistrationExtend(uint32) (runs: 1003, μ: 1276008, ~: 1276008)
+WakuRlnV2Test:test__ValidRegistrationExtendSingleMembership(uint32) (runs: 1003, μ: 314390, ~: 314390)
+WakuRlnV2Test:test__ValidRegistrationWithEraseList() (gas: 1601871)
+WakuRlnV2Test:test__ValidRegistration__kats() (gas: 295683)
+WakuRlnV2Test:test__WithdrawToken(uint32) (runs: 1003, μ: 301025, ~: 301027)
+WakuRlnV2Test:test__indexReuse_eraseMemberships(uint32) (runs: 1005, μ: 2702739, ~: 1165772)

src/LinearPriceCalculator.sol

@@ -1,19 +1,22 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.24;
 
-import { Ownable } from "openzeppelin-contracts/contracts/access/Ownable.sol";
+import { Ownable2Step } from "openzeppelin-contracts/contracts/access/Ownable2Step.sol";
 import { IPriceCalculator } from "./IPriceCalculator.sol";
 
+/// Address 0x0000...0000 was used instead of an ERC20 token address
+error OnlyERC20TokensAllowed();
+
 /// @title Linear Price Calculator to determine the price to acquire a membership
-contract LinearPriceCalculator is IPriceCalculator, Ownable {
-    /// @notice Address of the ERC20 token accepted by this contract. Address(0) represents ETH
+contract LinearPriceCalculator is IPriceCalculator, Ownable2Step {
+    /// @notice Address of the ERC20 token accepted by this contract.
     address public token;
 
     /// @notice The price per message per epoch
     uint256 public pricePerMessagePerEpoch;
 
-    constructor(address _token, uint256 _pricePerMessagePerEpoch) Ownable() {
-        require(_token != address(0), "only tokens can be used");
+    constructor(address _token, uint256 _pricePerMessagePerEpoch) Ownable2Step() {
+        if (_token == address(0)) revert OnlyERC20TokensAllowed();
         token = _token;
         pricePerMessagePerEpoch = _pricePerMessagePerEpoch;
     }

src/Membership.sol

@@ -4,9 +4,7 @@ pragma solidity 0.8.24;
 import { IPriceCalculator } from "./IPriceCalculator.sol";
 import { IERC20 } from "openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
 import { SafeERC20 } from "openzeppelin-contracts/contracts/token/ERC20/utils/SafeERC20.sol";
-
-// An eth value was assigned in the transaction and only tokens were expected
-error OnlyTokensAccepted();
+import { Initializable } from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 
 // The specified rate limit was not correct or within the expected limits
 error InvalidRateLimit();
@@ -24,7 +22,7 @@ error NotHolder(uint256 idCommitment);
 // This membership cannot be erased (either it is not expired or not in grace period and/or not the owner)
 error CantEraseMembership(uint256 idCommitment);
 
-contract Membership {
+abstract contract MembershipUpgradeable is Initializable {
     using SafeERC20 for IERC20;
 
     /// @notice Address of the Price Calculator used to calculate the price of a new membership
@@ -108,7 +106,28 @@ contract Membership {
     /// @param _maxRateLimitPerMembership Maximum rate limit of one membership
     /// @param _expirationTerm Membership expiration term
     /// @param _gracePeriod Membership grace period
-    function __Membership_init(
+    function __MembershipUpgradeable_init(
+        address _priceCalculator,
+        uint32 _maxTotalRateLimitPerEpoch,
+        uint32 _minRateLimitPerMembership,
+        uint32 _maxRateLimitPerMembership,
+        uint32 _expirationTerm,
+        uint32 _gracePeriod
+    )
+        internal
+        onlyInitializing
+    {
+        __MembershipUpgradeable_init_unchained(
+            _priceCalculator,
+            _maxTotalRateLimitPerEpoch,
+            _minRateLimitPerMembership,
+            _maxRateLimitPerMembership,
+            _expirationTerm,
+            _gracePeriod
+        );
+    }
+
+    function __MembershipUpgradeable_init_unchained(
         address _priceCalculator,
         uint32 _maxTotalRateLimitPerEpoch,
         uint32 _minRateLimitPerMembership,
@@ -117,7 +136,13 @@ contract Membership {
         uint32 _gracePeriod
     )
         internal
+        onlyInitializing
     {
+        require(_maxTotalRateLimitPerEpoch >= maxRateLimitPerMembership);
+        require(_maxRateLimitPerMembership > minRateLimitPerMembership);
+        require(_minRateLimitPerMembership > 0);
+        require(_expirationTerm > 0);
+
         priceCalculator = IPriceCalculator(_priceCalculator);
         maxTotalRateLimitPerEpoch = _maxTotalRateLimitPerEpoch;
         maxRateLimitPerMembership = _maxRateLimitPerMembership;
@@ -157,7 +182,6 @@ contract Membership {
     }
 
     function _transferFees(address _from, address _token, uint256 _amount) internal {
-        if (msg.value != 0) revert OnlyTokensAccepted();
         IERC20(_token).safeTransferFrom(_from, address(this), _amount);
     }
 
@@ -377,7 +401,7 @@ contract Membership {
 
         if (!membershipExpired && !isGracePeriodAndOwned) revert CantEraseMembership(_idCommitment);
 
-        emit MemberExpired(head, _mdetails.userMessageLimit, _mdetails.index);
+        emit MemberExpired(_idCommitment, _mdetails.userMessageLimit, _mdetails.index);
 
         // Move balance from expired membership to holder balance
         balancesToWithdraw[_mdetails.holder][_mdetails.token] += _mdetails.amount;
@@ -415,4 +439,11 @@ contract Membership {
         balancesToWithdraw[_sender][_token] = 0;
         IERC20(_token).safeTransfer(_sender, amount);
     }
+
+    /**
+     * @dev This empty reserved space is put in place to allow future versions to add new
+     * variables without shifting down storage in the inheritance chain.
+     * See https://docs.openzeppelin.com/contracts/4.x/upgradeable#storage_gaps
+     */
+    uint256[50] private __gap;
 }

src/WakuRlnV2.sol

@@ -4,11 +4,11 @@ pragma solidity 0.8.24;
 import { LazyIMT, LazyIMTData } from "@zk-kit/imt.sol/LazyIMT.sol";
 import { PoseidonT3 } from "poseidon-solidity/PoseidonT3.sol";
 
-import { OwnableUpgradeable } from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import { Ownable2StepUpgradeable } from "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
 import { Initializable } from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 
-import { Membership } from "./Membership.sol";
+import { MembershipUpgradeable } from "./Membership.sol";
 import { IPriceCalculator } from "./IPriceCalculator.sol";
 
 /// The tree is full
@@ -26,7 +26,7 @@ error InvalidUserMessageLimit(uint32 messageLimit);
 /// Invalid pagination query
 error InvalidPaginationQuery(uint256 startIndex, uint256 endIndex);
 
-contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Membership {
+contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, MembershipUpgradeable {
     /// @notice The Field
     uint256 public constant Q =
         21_888_242_871_839_275_222_246_405_745_257_275_088_548_364_400_416_034_343_698_204_186_575_808_495_617;
@@ -64,32 +64,27 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member
     /// @param _maxTotalRateLimitPerEpoch Maximum total rate limit of all memberships in the tree
     /// @param _minRateLimitPerMembership Minimum rate limit of one membership
     /// @param _maxRateLimitPerMembership Maximum rate limit of one membership
-    /// @param _billingPeriod Membership billing period
+    /// @param _expirationTerm Membership expiration term
     /// @param _gracePeriod Membership grace period
     function initialize(
         address _priceCalculator,
         uint32 _maxTotalRateLimitPerEpoch,
         uint32 _minRateLimitPerMembership,
         uint32 _maxRateLimitPerMembership,
-        uint32 _billingPeriod,
+        uint32 _expirationTerm,
         uint32 _gracePeriod
     )
         public
         initializer
     {
-        require(_maxTotalRateLimitPerEpoch >= maxRateLimitPerMembership);
-        require(_maxRateLimitPerMembership > minRateLimitPerMembership);
-        require(_minRateLimitPerMembership > 0);
-        require(_billingPeriod > 0);
-
         __Ownable_init();
         __UUPSUpgradeable_init();
-        __Membership_init(
+        __MembershipUpgradeable_init(
             _priceCalculator,
             _maxTotalRateLimitPerEpoch,
             _minRateLimitPerMembership,
             _maxRateLimitPerMembership,
-            _billingPeriod,
+            _expirationTerm,
             _gracePeriod
         );
 
@@ -120,8 +115,8 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member
     /// @return The metadata of the member (userMessageLimit, index, rateCommitment)
     function idCommitmentToMetadata(uint256 idCommitment) public view returns (uint32, uint32, uint256) {
         MembershipInfo memory member = members[idCommitment];
-        // we cannot call indexToCommitment if the member doesn't exist
-        if (member.holder == address(0)) {
+        // we cannot call indexToCommitment for 0 index if the member doesn't exist
+        if (member.userMessageLimit == 0) {
             return (0, 0, 0);
         }
         return (member.userMessageLimit, member.index, indexToCommitment(member.index));
@@ -165,6 +160,7 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member
         for (uint256 i = 0; i < membershipsToErase.length; i++) {
             uint256 idCommitmentToErase = membershipsToErase[i];
             MembershipInfo memory mdetails = members[idCommitmentToErase];
+            if (mdetails.userMessageLimit == 0) revert InvalidIdCommitment(idCommitmentToErase);
             _eraseMembership(_msgSender(), idCommitmentToErase, mdetails);
             LazyIMT.update(imtData, 0, mdetails.index);
         }
@@ -248,6 +244,7 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member
         for (uint256 i = 0; i < idCommitments.length; i++) {
             uint256 idCommitment = idCommitments[i];
             MembershipInfo memory mdetails = members[idCommitment];
+            if (mdetails.userMessageLimit == 0) revert InvalidIdCommitment(idCommitment);
             _eraseMembership(_msgSender(), idCommitment, mdetails);
             LazyIMT.update(imtData, 0, mdetails.index);
         }
@@ -288,7 +285,7 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member
 
     /// @notice Set the membership expiration term
     /// @param _expirationTerm  new value
-    function setBillingPeriod(uint32 _expirationTerm) external onlyOwner {
+    function setExpirationTerm(uint32 _expirationTerm) external onlyOwner {
         require(_expirationTerm > 0);
         expirationTerm = _expirationTerm;
     }

test/WakuRlnV2.t.sol

@@ -262,7 +262,7 @@ contract WakuRlnV2Test is Test {
 
         // Attempt to extend the membership (but now we are the owner)
         vm.expectEmit(true, false, false, false); // only check the first parameter of the event (the idCommitment)
-        emit Membership.MemberExtended(idCommitment, 0, 0, 0);
+        emit MembershipUpgradeable.MemberExtended(idCommitment, 0, 0, 0);
         w.extend(commitmentsToExtend);
 
         (,,, uint256 newGracePeriodStartDate,,,,,) = w.members(idCommitment);
@@ -325,7 +325,7 @@ contract WakuRlnV2Test is Test {
 
         // Extend the membership
         vm.expectEmit(true, false, false, false); // only check the first parameter of the event (the idCommitment)
-        emit Membership.MemberExtended(idCommitment, 0, 0, 0);
+        emit MembershipUpgradeable.MemberExtended(idCommitment, 0, 0, 0);
         w.extend(commitmentsToExtend);
 
         // Verify list order is correct
@@ -413,11 +413,11 @@ contract WakuRlnV2Test is Test {
         // Attempt to expire 3 commitments that can be erased
         commitmentsToErase[2] = 4;
         vm.expectEmit(true, false, false, false);
-        emit Membership.MemberExpired(1, 0, 0);
+        emit MembershipUpgradeable.MemberExpired(1, 0, 0);
         vm.expectEmit(true, false, false, false);
-        emit Membership.MemberExpired(2, 0, 0);
+        emit MembershipUpgradeable.MemberExpired(2, 0, 0);
         vm.expectEmit(true, false, false, false);
-        emit Membership.MemberExpired(4, 0, 0);
+        emit MembershipUpgradeable.MemberExpired(4, 0, 0);
         w.register(6, 60, commitmentsToErase);
 
         // Ensure that the chosen memberships were erased and others unaffected
@@ -518,7 +518,7 @@ contract WakuRlnV2Test is Test {
 
         // It should succeed now
         vm.expectEmit();
-        emit Membership.MemberExpired(1, userMessageLimitA, indexA);
+        emit MembershipUpgradeable.MemberExpired(1, userMessageLimitA, indexA);
         w.register(2, userMessageLimitB);
 
         // The previous expired membership should have been erased
@@ -579,9 +579,9 @@ contract WakuRlnV2Test is Test {
         (, uint256 priceB) = w.priceCalculator().calculate(4);
         token.approve(address(w), priceB);
         vm.expectEmit(true, false, false, false);
-        emit Membership.MemberExpired(1, 0, 0);
+        emit MembershipUpgradeable.MemberExpired(1, 0, 0);
         vm.expectEmit(true, false, false, false);
-        emit Membership.MemberExpired(2, 0, 0);
+        emit MembershipUpgradeable.MemberExpired(2, 0, 0);
         w.register(4, 4);
 
         // idCommitment4 will use the last removed index available (since we push to an array)
@@ -740,9 +740,9 @@ contract WakuRlnV2Test is Test {
         commitmentsToErase[1] = idCommitment + 2;
 
         vm.expectEmit(true, false, false, false); // only check the first parameter of the event (the idCommitment)
-        emit Membership.MemberExpired(commitmentsToErase[0], 0, 0);
+        emit MembershipUpgradeable.MemberExpired(commitmentsToErase[0], 0, 0);
         vm.expectEmit(true, false, false, false); // only check the first parameter of the event (the idCommitment)
-        emit Membership.MemberExpired(commitmentsToErase[0], 0, 0);
+        emit MembershipUpgradeable.MemberExpired(commitmentsToErase[0], 0, 0);
         w.eraseMemberships(commitmentsToErase);
 
         address holder;
@@ -803,7 +803,7 @@ contract WakuRlnV2Test is Test {
         for (uint256 i = 0; i < idCommitmentsLength; i++) {
             commitmentsToErase[i] = i + 1;
             vm.expectEmit(true, false, false, false); // only check the first parameter of the event (the idCommitment)
-            emit Membership.MemberExpired(i + 1, 0, 0);
+            emit MembershipUpgradeable.MemberExpired(i + 1, 0, 0);
         }
 
         w.eraseMemberships(commitmentsToErase);
@@ -898,7 +898,7 @@ contract WakuRlnV2Test is Test {
 
         /*| Name                | Type                                                | Slot | Offset | Bytes |
           |---------------------|-----------------------------------------------------|------|--------|-------|
-        | nextCommitmentIndex   | uint32                                              | 206  | 0      | 4     | */
+          | nextCommitmentIndex | uint32                                              | 256  | 0      | 4     | */
 
         /*
         Pro tip: to easily find the storage slot of a variable, without having to calculate the storage layout
@@ -917,7 +917,7 @@ contract WakuRlnV2Test is Test {
         */
 
         // we set nextCommitmentIndex to 4294967295 (1 << 20) = 0x00100000
-        vm.store(address(w), bytes32(uint256(206)), 0x0000000000000000000000000000000000000000000000000000000000100000);
+        vm.store(address(w), bytes32(uint256(256)), 0x0000000000000000000000000000000000000000000000000000000000100000);
         token.approve(address(w), price);
         vm.expectRevert(FullTree.selector);
         w.register(1, userMessageLimit);