fix: code review

waku-org · Sep 12, 2024 · 9da6d2c · 9da6d2c
1 parent 3fbcbfe
commit 9da6d2c
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 37 deletions.
diff --git a/src/LinearPriceCalculator.sol b/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;
     }

diff --git a/src/Membership.sol b/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;
 }
diff --git a/src/WakuRlnV2.sol b/src/WakuRlnV2.sol
@@ -8,7 +8,7 @@ import { OwnableUpgradeable } from "@openzeppelin/contracts-upgradeable/access/O
 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, OwnableUpgradeable, 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;
     }

diff --git a/test/WakuRlnV2.t.sol b/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);