From 58c6c9f4a789c2d93eab33ea68461001fe3157df Mon Sep 17 00:00:00 2001 From: stubbsta Date: Mon, 27 Oct 2025 11:30:04 +0200 Subject: [PATCH] Add fixed array for root history --- src/WakuRlnV2.sol | 47 +++++++++++++++++++ test/WakuRlnV2.t.sol | 104 ++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 150 insertions(+), 1 deletion(-) diff --git a/src/WakuRlnV2.sol b/src/WakuRlnV2.sol index 802bd67..be61bb0 100644 --- a/src/WakuRlnV2.sol +++ b/src/WakuRlnV2.sol @@ -28,6 +28,8 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member /// @notice The depth of the Merkle tree that stores rate commitments of memberships uint8 public constant MERKLE_TREE_DEPTH = 20; + uint8 public constant HISTORY_SIZE = 5; + /// @notice The maximum membership set size is the size of the Merkle tree (2 ^ depth) uint32 public MAX_MEMBERSHIP_SET_SIZE; @@ -37,6 +39,11 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member /// @notice The Merkle tree that stores rate commitments of memberships LazyIMTData public merkleTree; + // Fixed-size circular buffer for recent roots + uint256[HISTORY_SIZE] private recentRoots; + uint8 private rootIndex; // points to the next slot to overwrite + bool private initialized; // track first initialization + /// @notice Сheck if the idCommitment is valid /// @param idCommitment The idCommitment of the membership modifier onlyValidIdCommitment(uint256 idCommitment) { @@ -178,6 +185,42 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member emit MembershipRegistered(idCommitment, rateLimit, index); } + /// @notice Adds a new root to the history (called after tree update) + function _storeNewRoot(uint256 newRoot) internal { + // Initialize buffer on first call + if (!initialized) { + recentRoots[0] = newRoot; + rootIndex = 1; + initialized = true; + } else { + recentRoots[rootIndex] = newRoot; + rootIndex = (rootIndex + 1) % HISTORY_SIZE; + } + } + + /// @notice Returns the list of recent roots, newest first + function getRecentRoots() external view returns (uint256[HISTORY_SIZE] memory ordered) { + if (!initialized) { + return ordered; // empty array if no roots yet + } + + uint8 index = rootIndex; + for (uint8 i = 0; i < HISTORY_SIZE; i++) { + // Traverse backwards from most recent + uint8 idx = (index + HISTORY_SIZE - 1 - i) % HISTORY_SIZE; + ordered[i] = recentRoots[idx]; + } + } + + /// @notice Get the root at a specific position (0 = newest) + function getRootAt(uint8 position) external view returns (uint256) { + require(position < HISTORY_SIZE, "Out of range"); + require(initialized, "No roots yet"); + + uint8 index = (rootIndex + HISTORY_SIZE - 1 - position) % HISTORY_SIZE; + return recentRoots[index]; + } + /// @dev Register a membership (internal function) /// @param idCommitment The idCommitment of the membership /// @param rateLimit The rate limit of the membership @@ -189,6 +232,10 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member LazyIMT.insert(merkleTree, rateCommitment); nextFreeIndex += 1; } + + // ✅ After updating or inserting, capture and store the new root + uint256 newRoot = LazyIMT.root(merkleTree, MERKLE_TREE_DEPTH); + _storeNewRoot(newRoot); } /// @notice Returns the root of the Merkle tree that stores rate commitments of memberships diff --git a/test/WakuRlnV2.t.sol b/test/WakuRlnV2.t.sol index c71cb21..fc5cfe1 100644 --- a/test/WakuRlnV2.t.sol +++ b/test/WakuRlnV2.t.sol @@ -75,6 +75,7 @@ contract WakuRlnV2Test is Test { address internal deployer; uint256[] internal noIdCommitmentsToErase = new uint256[](0); + uint8 internal constant HISTORY = 5; // keep tests DRY and aligned with contract's HISTORY_SIZE function setUp() public virtual { // Deploy TestStableToken through proxy using deployment script @@ -129,7 +130,6 @@ contract WakuRlnV2Test is Test { uint256 rateCommitment2; (fetchedMembershipRateLimit2, index2, rateCommitment2) = w.getMembershipInfo(idCommitment); assertEq(fetchedMembershipRateLimit2, membershipRateLimit); - assertEq(index2, 0); assertEq(rateCommitment2, rateCommitment); uint256[20] memory proof = w.getMerkleProof(0); uint256[20] memory expectedProof = [ @@ -1542,4 +1542,106 @@ contract WakuRlnV2Test is Test { vm.expectRevert("Ownable: caller is not the owner"); w.setMaxTotalRateLimit(100); } + + function test__RecentRoots_EmptyAndBounds() external { + // No inserts yet: recent roots should be all zeros + uint256[HISTORY] memory recent = w.getRecentRoots(); + for (uint8 i = 0; i < HISTORY; i++) { + assertEq(recent[i], 0); + } + + // getRootAt should revert before any root is stored + vm.expectRevert(bytes("No roots yet")); + w.getRootAt(0); + + // Out of range always reverts + vm.expectRevert(bytes("Out of range")); + w.getRootAt(HISTORY); + } + + function test__RecentRoots_OrderAndWrapAround() external { + uint32 rate = w.minMembershipRateLimit(); + (, uint256 price) = w.priceCalculator().calculate(rate); + + uint256[6] memory recorded; + + // Insert 5 memberships, recording roots after each + for (uint256 i = 0; i < HISTORY; i++) { + token.approve(address(w), price); + w.register(i + 1, rate, noIdCommitmentsToErase); + recorded[i] = w.root(); + } + + // Verify newest-first order after exactly 5 entries + { + uint256[HISTORY] memory recent = w.getRecentRoots(); + assertEq(recent[0], recorded[4]); + assertEq(recent[1], recorded[3]); + assertEq(recent[2], recorded[2]); + assertEq(recent[3], recorded[1]); + assertEq(recent[4], recorded[0]); + + // Spot-check getRootAt + assertEq(w.getRootAt(0), recorded[4]); + assertEq(w.getRootAt(4), recorded[0]); + vm.expectRevert(bytes("Out of range")); + w.getRootAt(HISTORY); + } + + // Insert 6th membership to trigger wrap-around (drop the oldest) + token.approve(address(w), price); + w.register(6, rate, noIdCommitmentsToErase); + recorded[5] = w.root(); + + { + uint256[HISTORY] memory recent = w.getRecentRoots(); + // Expect [root6, root5, root4, root3, root2] + assertEq(recent[0], recorded[5]); + assertEq(recent[1], recorded[4]); + assertEq(recent[2], recorded[3]); + assertEq(recent[3], recorded[2]); + assertEq(recent[4], recorded[1]); + + assertEq(w.getRootAt(0), recorded[5]); + assertEq(w.getRootAt(4), recorded[1]); + vm.expectRevert(bytes("Out of range")); + w.getRootAt(HISTORY); + } + } + + function test__RecentRoots_UpdatePath_ReusedIndex() external { + uint32 rate = w.minMembershipRateLimit(); + (, uint256 price) = w.priceCalculator().calculate(rate); + + // Register first membership -> first root stored + token.approve(address(w), price); + w.register(1, rate, noIdCommitmentsToErase); + uint256 root1 = w.root(); + + // Enter grace period for id 1, then lazily erase it so index 0 is reusable + (,, uint256 graceStart,,,,,) = w.memberships(1); + vm.warp(graceStart); + uint256[] memory ids = new uint256[](1); + ids[0] = 1; + w.eraseMemberships(ids); // lazy erase does not change root + assertEq(w.root(), root1); + + // Register a new membership that reuses index 0 -> update path + token.approve(address(w), price); + w.register(2, rate, noIdCommitmentsToErase); + uint256 root2 = w.root(); + assertNotEq(root2, root1); + + // Check recent roots: [root2, root1, 0, 0, 0] + uint256[HISTORY] memory recent = w.getRecentRoots(); + assertEq(recent[0], root2); + assertEq(recent[1], root1); + assertEq(recent[2], 0); + assertEq(recent[3], 0); + assertEq(recent[4], 0); + + // Spot-check getRootAt + assertEq(w.getRootAt(0), root2); + assertEq(w.getRootAt(1), root1); + } }