From 0201d38f6fd4aa5f89eb594586e477863ed31419 Mon Sep 17 00:00:00 2001 From: Sergei Tikhomirov Date: Wed, 2 Oct 2024 17:57:28 +0200 Subject: [PATCH] minor refactoring, comments --- src/Membership.sol | 102 +++++++++++++++++++-------------------------- src/WakuRlnV2.sol | 52 ++++++++++++----------- 2 files changed, 70 insertions(+), 84 deletions(-) diff --git a/src/Membership.sol b/src/Membership.sol index 5bb78bb..fd28711 100644 --- a/src/Membership.sol +++ b/src/Membership.sol @@ -6,17 +6,17 @@ import { IERC20 } from "openzeppelin-contracts/contracts/token/ERC20/IERC20.sol" import { SafeERC20 } from "openzeppelin-contracts/contracts/token/ERC20/utils/SafeERC20.sol"; import { Initializable } from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol"; -// The specified rate limit was not correct or within the expected limits +// The rate limit is outside the expected limits error InvalidMembershipRateLimit(); -// It's not possible to acquire the rate limit due to exceeding the expected limits +// Cannot acquire the rate limit for a new membership due to exceeding the expected limits // even after attempting to erase expired memberships error ExceededMaxTotalRateLimit(); -// This membership is not in grace period +// This membership is not in its grace period error NotInGracePeriod(uint256 idCommitment); -// The sender is not the holder of the membership +// The sender is not the holder of this membership error AttemptedExtensionByNonHolder(uint256 idCommitment); // This membership cannot be erased @@ -46,8 +46,8 @@ abstract contract MembershipUpgradeable is Initializable { /// @notice Membership grace period duration (G in the spec) uint32 public gracePeriodDurationForNewMemberships; - /// @notice deposits available for withdrawal - /// Note: deposits unavailable for withdrawal are stored in MembershipInfo. + /// @notice Deposits available for withdrawal + /// Note: amount of deposits unavailable for withdrawal are stored in MembershipInfo elements. mapping(address holder => mapping(address token => uint256 balance)) public depositsToWithdraw; /// @notice Current total rate limit of all memberships in the membership set (messages per epoch) @@ -59,7 +59,7 @@ abstract contract MembershipUpgradeable is Initializable { /// @notice The index in the membership set for the next membership to be registered uint32 public nextFreeIndex; - /// @notice Indices of memberships marked as erased but not yet actually erased from the membership set + /// @notice Indices of erased memberships that can be reused for new registrations uint32[] public indicesOfLazilyErasedMemberships; struct MembershipInfo { @@ -75,9 +75,9 @@ abstract contract MembershipUpgradeable is Initializable { uint32 rateLimit; /// @notice the index of the membership in the membership set uint32 index; - /// @notice address of the holder of this membership + /// @notice the address of the holder of this membership address holder; - /// @notice token used to make the deposit to register this membership + /// @notice the token used to make the deposit to register this membership address token; } @@ -161,7 +161,7 @@ abstract contract MembershipUpgradeable is Initializable { } /// @dev acquire a membership and trasnfer the deposit to the contract - /// @param _sender address of the holder of the new membership + /// @param _sender the address of the transaction sender /// @param _idCommitment the idCommitment of the new membership /// @param _rateLimit the membership rate limit /// @return index the index of the new membership in the membership set @@ -188,7 +188,7 @@ abstract contract MembershipUpgradeable is Initializable { (address token, uint256 depositAmount) = priceCalculator.calculate(_rateLimit); - // Possibly reuse an index of an available erased membership + // Possibly reuse an index of an erased membership (index, indexReused) = _getFreeIndex(); memberships[_idCommitment] = MembershipInfo({ @@ -212,13 +212,13 @@ abstract contract MembershipUpgradeable is Initializable { return minMembershipRateLimit <= rateLimit && rateLimit <= maxMembershipRateLimit; } - /// @dev Get a free index (possibly from reusing an index of an erased membership) - /// @return index index to be used for another membership registration - /// @return indexReused indicates whether index was reused from an erased membership + /// @dev Get a free index (possibly reuse an index of an erased membership) + /// @return index index to be used for the new membership registration + /// @return indexReused indicates whether the index was reused from an erased membership function _getFreeIndex() internal returns (uint32 index, bool indexReused) { - // Reuse the last erased membership uint256 numIndices = indicesOfLazilyErasedMemberships.length; if (numIndices != 0) { + // Reuse the index of the latest erased membership index = indicesOfLazilyErasedMemberships[numIndices - 1]; indicesOfLazilyErasedMemberships.pop(); indexReused = true; @@ -227,13 +227,13 @@ abstract contract MembershipUpgradeable is Initializable { } } - /// @dev Extend the expiration date of a grace-period membership - /// @param _sender the address of the holder of the membership + /// @dev Extend a grace-period membership + /// @param _sender the address of the transaction sender /// @param _idCommitment the idCommitment of the membership function _extendMembership(address _sender, uint256 _idCommitment) public { MembershipInfo storage membership = memberships[_idCommitment]; - if (!_isInGracePeriod(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration)) { + if (!_isInPeriod(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration)) { revert NotInGracePeriod(_idCommitment); } @@ -251,16 +251,16 @@ 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 _sender the address of the transaction sender /// @param _idCommitment idCommitment of the membership to erase 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 membershipExpired = _isAfterPeriod(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration); bool membershipIsInGracePeriod = - _isInGracePeriod(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration); + _isInPeriod(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration); bool isHolder = (membership.holder == _sender); if (!membershipExpired && !(membershipIsInGracePeriod && isHolder)) { @@ -284,72 +284,56 @@ abstract contract MembershipUpgradeable is Initializable { } emit MembershipErased(_idCommitment, membership.rateLimit, membership.index); - // The caller uses this index to erase data from the membership set (the Merkle tree) + // This index will be used to erase the data from the Merkle tree that represents the membership set return membership.index; } - /// @notice Determine if a membership is in grace period now + /// @notice Determine if a membership is in its grace period /// @param _idCommitment the idCommitment of the membership function isInGracePeriod(uint256 _idCommitment) public view returns (bool) { MembershipInfo memory membership = memberships[_idCommitment]; - return _isInGracePeriod(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration); + return _isInPeriod(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration); } - /// @dev Determine whether the current timestamp is in a given grace period - /// @param _gracePeriodStartTimestamp timestamp in which the grace period starts - /// @param _gracePeriodDuration duration of the grace period - function _isInGracePeriod( - uint256 _gracePeriodStartTimestamp, - uint32 _gracePeriodDuration - ) - internal - view - returns (bool) - { + /// @dev Determine whether the current timestamp is within a given period + /// @param _start timestamp in which the period starts + /// @param _duration duration of the period + function _isInPeriod(uint256 _start, uint32 _duration) internal view returns (bool) { uint256 timeNow = block.timestamp; - return ( - _gracePeriodStartTimestamp <= timeNow - && timeNow <= _gracePeriodStartTimestamp + uint256(_gracePeriodDuration) - ); + return (_start <= timeNow && timeNow <= _start + uint256(_duration)); } /// @notice Determine if a membership is expired /// @param _idCommitment the idCommitment of the membership function isExpired(uint256 _idCommitment) public view returns (bool) { MembershipInfo memory membership = memberships[_idCommitment]; - return _isExpired(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration); + return _isAfterPeriod(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration); } - /// @dev Determine whether a grace period ends(the membership is expired) - /// @param _gracePeriodStartTimestamp timestamp in which the grace period starts - /// @param _gracePeriodDuration duration of the grace period - function _isExpired(uint256 _gracePeriodStartTimestamp, uint32 _gracePeriodDuration) internal view returns (bool) { - return block.timestamp >= _membershipExpirationTimestamp(_gracePeriodStartTimestamp, _gracePeriodDuration); + /// @dev Determine whether the current timestamp is after a given period + /// @param _start timestamp in which the period starts + /// @param _duration duration of the period + function _isAfterPeriod(uint256 _start, uint32 _duration) internal view returns (bool) { + uint256 timeNow = block.timestamp; + return (_start + uint256(_duration) < timeNow); } /// @notice Returns the timestamp on which a membership can be considered expired (i.e. when its grace period ends) /// @param _idCommitment the idCommitment of the membership function membershipExpirationTimestamp(uint256 _idCommitment) public view returns (uint256) { MembershipInfo memory membership = memberships[_idCommitment]; - return _membershipExpirationTimestamp(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration); + return _timestampAfterPeriod(membership.gracePeriodStartTimestamp, membership.gracePeriodDuration); } - /// @dev Determine when the grace period ends - /// @param _gracePeriodStartTimestamp timestamp in which the grace period starts - /// @param _gracePeriodDuration duration of the grace period - function _membershipExpirationTimestamp( - uint256 _gracePeriodStartTimestamp, - uint32 _gracePeriodDuration - ) - internal - pure - returns (uint256) - { - return _gracePeriodStartTimestamp + uint256(_gracePeriodDuration) + 1; + /// @dev Returns the first timestamp after a specified period + /// @param _start timestamp in which the period starts + /// @param _duration duration of the period + function _timestampAfterPeriod(uint256 _start, uint32 _duration) internal pure returns (uint256) { + return _start + uint256(_duration) + 1; } /// @dev Withdraw any available deposit balance in tokens after a membership is erased. - /// @param _sender the address of the owner of the tokens + /// @param _sender the address of the transaction sender (who withdraws their tokens) /// @param _token the address of the token to withdraw function _withdraw(address _sender, address _token) internal { require(_token != address(0), "ETH is not allowed"); diff --git a/src/WakuRlnV2.sol b/src/WakuRlnV2.sol index ead0834..b8ed7cb 100644 --- a/src/WakuRlnV2.sol +++ b/src/WakuRlnV2.sol @@ -169,10 +169,10 @@ contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, M _register(idCommitment, rateLimit); } - /// @notice Register a membership + /// @notice Register a membership while erasing some expired memberships to reuse their rate limit /// @param idCommitment The idCommitment of the new membership /// @param rateLimit The rate limit of the new membership - /// @param idCommitmentsToErase List of idCommitments of expired memberships to erase + /// @param idCommitmentsToErase The list of idCommitments of expired memberships to erase function register( uint256 idCommitment, uint32 rateLimit, @@ -183,18 +183,16 @@ contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, M noDuplicateMembership(idCommitment) membershipSetNotFull { + // erase memberships without overwriting membership set data to zero (save gas) _eraseMemberships(idCommitmentsToErase, false); _register(idCommitment, rateLimit); } - /// @dev Registers a membership + /// @dev Register a membership (internal function) /// @param idCommitment The idCommitment of the membership /// @param rateLimit The rate limit of the membership function _register(uint256 idCommitment, uint32 rateLimit) internal { - uint32 index; - bool indexReused; - (index, indexReused) = _acquireMembership(_msgSender(), idCommitment, rateLimit); - + (uint32 index, bool indexReused) = _acquireMembership(_msgSender(), idCommitment, rateLimit); uint256 rateCommitment = PoseidonT3.hash([idCommitment, rateLimit]); if (indexReused) { LazyIMT.update(merkleTree, rateCommitment, index); @@ -224,25 +222,23 @@ contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, M return fixedSizeProof; } - /// @notice Extend a grace-period membership + /// @notice Extend a grace-period membership under the same conditions /// @param idCommitments list of idCommitments of memberships to extend function extendMemberships(uint256[] calldata idCommitments) external { for (uint256 i = 0; i < idCommitments.length; i++) { - uint256 idCommitmentToExtend = idCommitments[i]; - _extendMembership(_msgSender(), idCommitmentToExtend); + _extendMembership(_msgSender(), idCommitments[i]); } } /// @notice Erase expired memberships or owned grace-period memberships /// The user can select expired memberships offchain, /// and proceed to erase them. - /// 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 - /// @param eraseFromMembershipSet Indicates whether to erase membership data from the set + /// This function is also used by the holder to erase their own grace-period memberships. + /// The holder can then withdraw the deposited tokens. + /// @param idCommitments The list of idCommitments of the memberships to erase + /// @param eraseFromMembershipSet Indicates whether to erase the membership's rate commitment from the membership + /// 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); } @@ -250,19 +246,25 @@ contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, M /// @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 { + // eraseFromMembershipSet == true means full clean-up. + // Erase memberships from memberships array (free up the rate limit and index), + // and also erase the rate limit from the Merkle tree that represents the membership set (reduce the tree + // size). + // eraseFromMembershipSet == false means lazy erasure. + // Only erase memberships from the membership array (consume less gas). + // Merkle tree data will be overwritten when the correspondind index is reused. for (uint256 i = 0; i < idCommitmentsToErase.length; i++) { - uint256 idCommitmentToErase = idCommitmentsToErase[i]; // 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 + uint32 indexToErase = _eraseMembershipLazily(_msgSender(), idCommitmentsToErase[i]); + // Optionally, also erase the rate commitment data from the membership set. + // This does not affect the total rate limit control, or index reusal for new membership registrations. if (eraseFromMembershipSet) { LazyIMT.update(merkleTree, 0, indexToErase); } } } - /// @notice Withdraw any available deposit balance in tokens after a membership is erased. + /// @notice Withdraw any available deposit balance in tokens after a membership is erased /// @param token The address of the token to withdraw function withdraw(address token) external { _withdraw(_msgSender(), token); @@ -277,14 +279,14 @@ contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, M /// @notice Set the maximum total rate limit of all memberships in the membership set /// @param _maxTotalRateLimit new maximum total rate limit (messages per epoch) function setMaxTotalRateLimit(uint32 _maxTotalRateLimit) external onlyOwner { - require(_maxTotalRateLimit >= maxMembershipRateLimit); + require(maxMembershipRateLimit <= _maxTotalRateLimit); maxTotalRateLimit = _maxTotalRateLimit; } /// @notice Set the maximum rate limit of one membership /// @param _maxMembershipRateLimit new maximum rate limit per membership (messages per epoch) function setMaxMembershipRateLimit(uint32 _maxMembershipRateLimit) external onlyOwner { - require(_maxMembershipRateLimit >= minMembershipRateLimit); + require(minMembershipRateLimit <= _maxMembershipRateLimit); maxMembershipRateLimit = _maxMembershipRateLimit; } @@ -297,13 +299,13 @@ contract WakuRlnV2 is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, M } /// @notice Set the active duration for new memberships (terms of existing memberships don't change) - /// @param _activeDurationForNewMembership new active duration + /// @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) + /// @notice Set the grace period for new memberships (terms of existing memberships don't change) /// @param _gracePeriodDurationForNewMembership new grace period duration function setGracePeriodDuration(uint32 _gracePeriodDurationForNewMembership) external onlyOwner { // Note: grace period duration may be equal to zero