2024-09-23 17:14:27 +02:00

299 lines
12 KiB
Solidity

// SPDX-License-Identifier: MIT
pragma solidity 0.8.24;
import { LazyIMT, LazyIMTData } from "@zk-kit/imt.sol/LazyIMT.sol";
import { PoseidonT3 } from "poseidon-solidity/PoseidonT3.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 { MembershipUpgradeable } from "./Membership.sol";
import { IPriceCalculator } from "./IPriceCalculator.sol";
/// The tree is full
error FullTree();
/// Member is already registered
error DuplicateIdCommitment();
/// Invalid idCommitment
error InvalidIdCommitment(uint256 idCommitment);
/// Invalid userMessageLimit
error InvalidUserMessageLimit(uint32 messageLimit);
/// Invalid pagination query
error InvalidPaginationQuery(uint256 startIndex, uint256 endIndex);
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;
/// @notice The depth of the merkle tree
uint8 public constant DEPTH = 20;
/// @notice The size of the merkle tree, i.e 2^depth
uint32 public SET_SIZE;
/// @notice the deployed block number
uint32 public deployedBlockNumber;
/// @notice the stored imt data
LazyIMTData public imtData;
/// Emitted when a new member is added to the set
/// @param rateCommitment the rateCommitment of the member
/// @param index The index of the member in the set
event MemberRegistered(uint256 rateCommitment, uint32 index);
/// @notice the modifier to check if the idCommitment is valid
/// @param idCommitment The idCommitment of the member
modifier onlyValidIdCommitment(uint256 idCommitment) {
if (!isValidCommitment(idCommitment)) revert InvalidIdCommitment(idCommitment);
_;
}
constructor() {
_disableInitializers();
}
/// @dev contract initializer
/// @param _priceCalculator Address of an instance of IPriceCalculator
/// @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 _expirationTerm Membership expiration term
/// @param _gracePeriod Membership grace period
function initialize(
address _priceCalculator,
uint32 _maxTotalRateLimitPerEpoch,
uint32 _minRateLimitPerMembership,
uint32 _maxRateLimitPerMembership,
uint32 _expirationTerm,
uint32 _gracePeriod
)
public
initializer
{
__Ownable_init();
__UUPSUpgradeable_init();
__MembershipUpgradeable_init(
_priceCalculator,
_maxTotalRateLimitPerEpoch,
_minRateLimitPerMembership,
_maxRateLimitPerMembership,
_expirationTerm,
_gracePeriod
);
SET_SIZE = uint32(1 << DEPTH);
deployedBlockNumber = uint32(block.number);
LazyIMT.init(imtData, DEPTH);
nextCommitmentIndex = 0;
}
function _authorizeUpgrade(address newImplementation) internal override onlyOwner { } // solhint-disable-line
/// @notice Checks if a commitment is valid
/// @param idCommitment The idCommitment of the member
/// @return true if the commitment is valid, false otherwise
function isValidCommitment(uint256 idCommitment) public pure returns (bool) {
return idCommitment != 0 && idCommitment < Q;
}
/// @notice Returns the rateCommitment of a member
/// @param index The index of the member
/// @return The rateCommitment of the member
function indexToCommitment(uint32 index) internal view returns (uint256) {
return imtData.elements[LazyIMT.indexForElement(0, index)];
}
/// @notice Returns the metadata of a member
/// @param idCommitment The idCommitment of the member
/// @return The metadata of the member (userMessageLimit, index, rateCommitment)
function idCommitmentToMetadata(uint256 idCommitment) public view returns (uint32, uint32, uint256) {
MembershipInfo memory mdetails = members[idCommitment];
// we cannot call indexToCommitment for 0 index if the member doesn't exist
if (mdetails.userMessageLimit == 0) {
return (0, 0, 0);
}
return (mdetails.userMessageLimit, mdetails.index, indexToCommitment(mdetails.index));
}
/// @notice Checks if a member exists
/// @param idCommitment The idCommitment of the member
/// @return true if the member exists, false otherwise
function memberExists(uint256 idCommitment) public view returns (bool) {
(,, uint256 rateCommitment) = idCommitmentToMetadata(idCommitment);
return rateCommitment != 0;
}
/// @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 onlyValidIdCommitment(idCommitment) {
if (memberExists(idCommitment)) revert DuplicateIdCommitment();
uint32 index;
bool reusedIndex;
(index, reusedIndex) = _acquireMembership(_msgSender(), idCommitment, userMessageLimit);
_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];
if (mdetails.userMessageLimit == 0) revert InvalidIdCommitment(idCommitmentToErase);
_eraseMembership(_msgSender(), idCommitmentToErase, mdetails);
LazyIMT.update(imtData, 0, mdetails.index);
}
uint32 index;
bool reusedIndex;
(index, reusedIndex) = _acquireMembership(_msgSender(), idCommitment, userMessageLimit);
_register(idCommitment, userMessageLimit, index, reusedIndex);
}
/// @dev Registers a member
/// @param idCommitment The idCommitment of the member
/// @param userMessageLimit The message limit of the member
/// @param index Indicates the index in the merkle tree
/// @param reusedIndex indicates whether we're inserting a new element in the merkle tree or updating a existing
/// leaf
function _register(uint256 idCommitment, uint32 userMessageLimit, uint32 index, bool reusedIndex) internal {
if (nextCommitmentIndex >= SET_SIZE) revert FullTree();
uint256 rateCommitment = PoseidonT3.hash([idCommitment, userMessageLimit]);
if (reusedIndex) {
LazyIMT.update(imtData, rateCommitment, index);
} else {
LazyIMT.insert(imtData, rateCommitment);
nextCommitmentIndex += 1;
}
emit MemberRegistered(rateCommitment, index);
}
/// @notice Returns the commitments of a range of members
/// @param startIndex The start index of the range
/// @param endIndex The end index of the range
/// @return The commitments of the members
function getCommitments(uint32 startIndex, uint32 endIndex) public view returns (uint256[] memory) {
if (startIndex > endIndex) revert InvalidPaginationQuery(startIndex, endIndex);
if (endIndex > nextCommitmentIndex) revert InvalidPaginationQuery(startIndex, endIndex);
uint256[] memory commitments = new uint256[](endIndex - startIndex + 1);
for (uint32 i = startIndex; i <= endIndex; i++) {
commitments[i - startIndex] = indexToCommitment(i);
}
return commitments;
}
/// @notice Returns the root of the IMT
/// @return The root of the IMT
function root() external view returns (uint256) {
return LazyIMT.root(imtData, DEPTH);
}
/// @notice Returns the merkle proof elements of a given membership
/// @param index The index of the member
/// @return The merkle proof elements of the member
function merkleProofElements(uint40 index) public view returns (uint256[DEPTH] memory) {
uint256[DEPTH] memory castedProof;
uint256[] memory proof = LazyIMT.merkleProofElements(imtData, index, DEPTH);
for (uint8 i = 0; i < DEPTH; i++) {
castedProof[i] = proof[i];
}
return castedProof;
}
/// @notice Extend a membership expiration date. Memberships must be on grace period
/// @param idCommitments list of idcommitments
function extend(uint256[] calldata idCommitments) external {
for (uint256 i = 0; i < idCommitments.length; i++) {
uint256 idCommitment = idCommitments[i];
_extendMembership(_msgSender(), idCommitment);
}
}
/// @notice Remove expired memberships or owned memberships in grace period.
/// The user can determine offchain which expired memberships slots
/// are available, and proceed to free them.
/// This is also used to erase memberships in grace period if they're
/// held by the sender. The sender can then withdraw the tokens.
/// @param idCommitments list of idcommitments of the memberships
function eraseMemberships(uint256[] calldata idCommitments) external {
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);
}
}
/// @notice Withdraw any available balance in tokens after a membership is erased.
/// @param token The address of the token to withdraw. Use 0x000...000 to withdraw ETH
function withdraw(address token) external {
_withdraw(_msgSender(), token);
}
/// @notice Set the address of the price calculator
/// @param _priceCalculator new price calculator address
function setPriceCalculator(address _priceCalculator) external onlyOwner {
priceCalculator = IPriceCalculator(_priceCalculator);
}
/// @notice Set the maximum total rate limit of all memberships in the tree
/// @param _maxTotalRateLimitPerEpoch new value
function setMaxTotalRateLimitPerEpoch(uint32 _maxTotalRateLimitPerEpoch) external onlyOwner {
require(_maxTotalRateLimitPerEpoch >= maxRateLimitPerMembership);
maxTotalRateLimitPerEpoch = _maxTotalRateLimitPerEpoch;
}
/// @notice Set the maximum rate limit of one membership
/// @param _maxRateLimitPerMembership new value
function setMaxRateLimitPerMembership(uint32 _maxRateLimitPerMembership) external onlyOwner {
require(_maxRateLimitPerMembership >= minRateLimitPerMembership);
maxRateLimitPerMembership = _maxRateLimitPerMembership;
}
/// @notice Set the minimum rate limit of one membership
/// @param _minRateLimitPerMembership new value
function setMinRateLimitPerMembership(uint32 _minRateLimitPerMembership) external onlyOwner {
require(_minRateLimitPerMembership > 0);
minRateLimitPerMembership = _minRateLimitPerMembership;
}
/// @notice Set the membership expiration term
/// @param _expirationTerm new value
function setExpirationTerm(uint32 _expirationTerm) external onlyOwner {
require(_expirationTerm > 0);
expirationTerm = _expirationTerm;
}
/// @notice Set the membership grace period
/// @param _gracePeriod new value
function setGracePeriod(uint32 _gracePeriod) external onlyOwner {
gracePeriod = _gracePeriod;
}
}