From 3805a53daad458c00859a4265f829653b03ff561 Mon Sep 17 00:00:00 2001 From: stubbsta Date: Wed, 29 Oct 2025 18:40:02 +0200 Subject: [PATCH] Update _eraseMemberships to upate roots --- src/WakuRlnV2.sol | 22 +++++++++++------ test/WakuRlnV2.t.sol | 59 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 8 deletions(-) diff --git a/src/WakuRlnV2.sol b/src/WakuRlnV2.sol index be61bb0..5caf6b8 100644 --- a/src/WakuRlnV2.sol +++ b/src/WakuRlnV2.sol @@ -200,6 +200,7 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member /// @notice Returns the list of recent roots, newest first function getRecentRoots() external view returns (uint256[HISTORY_SIZE] memory ordered) { + // refresh the latest root - needed if memberships were changed/erased but no new root stored yet if (!initialized) { return ordered; // empty array if no roots yet } @@ -221,6 +222,12 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member return recentRoots[index]; } + /// @notice Returns the root of the Merkle tree that stores rate commitments of memberships + /// @return The root of the Merkle tree that stores rate commitments of memberships + function root() public view returns (uint256) { + return LazyIMT.root(merkleTree, MERKLE_TREE_DEPTH); + } + /// @dev Register a membership (internal function) /// @param idCommitment The idCommitment of the membership /// @param rateLimit The rate limit of the membership @@ -234,14 +241,7 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member } // ✅ 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 - /// @return The root of the Merkle tree that stores rate commitments of memberships - function root() external view returns (uint256) { - return LazyIMT.root(merkleTree, MERKLE_TREE_DEPTH); + _storeNewRoot(root()); } /// @notice Returns the Merkle proof that a given membership is in the membership set @@ -294,6 +294,7 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member // eraseFromMembershipSet == false means lazy erasure. // Only erase memberships from the memberships array (consume less gas). // Merkle tree data will be overwritten when the correspondind index is reused. + bool treeModified = false; for (uint256 i = 0; i < idCommitmentsToErase.length; i++) { // Erase the membership from the memberships array in contract storage uint32 indexToErase = _eraseMembershipLazily(_msgSender(), idCommitmentsToErase[i]); @@ -301,8 +302,13 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member // This does not affect the total rate limit control, or index reusal for new membership registrations. if (eraseFromMembershipSet) { LazyIMT.update(merkleTree, 0, indexToErase); + treeModified = true; } } + // Record only the final root once per batch full clean-up call to avoid duplicates. + if (treeModified) { + _storeNewRoot(root()); + } } /// @notice Withdraw any available deposit balance in tokens after a membership is erased diff --git a/test/WakuRlnV2.t.sol b/test/WakuRlnV2.t.sol index fc5cfe1..0ac698d 100644 --- a/test/WakuRlnV2.t.sol +++ b/test/WakuRlnV2.t.sol @@ -1644,4 +1644,63 @@ contract WakuRlnV2Test is Test { assertEq(w.getRootAt(0), root2); assertEq(w.getRootAt(1), root1); } + + function test__RecentRoots_FullCleanupErasureUpdatesHistory() external { + uint32 rate = w.minMembershipRateLimit(); + (, uint256 price) = w.priceCalculator().calculate(rate); + + // Register three memberships (1,2,3), capturing roots after each + token.approve(address(w), price); + w.register(1, rate, noIdCommitmentsToErase); + uint256 r1 = w.root(); + + token.approve(address(w), price); + w.register(2, rate, noIdCommitmentsToErase); + uint256 r2 = w.root(); + + token.approve(address(w), price); + w.register(3, rate, noIdCommitmentsToErase); + uint256 r3 = w.root(); + + // Expire id 1 and erase with full clean-up (tree changes) + (,, uint256 gStart1, uint32 gDur1,,,,) = w.memberships(1); + vm.warp(gStart1 + gDur1 + 1); + uint256[] memory toErase = new uint256[](1); + toErase[0] = 1; + w.eraseMemberships(toErase, true); + uint256 r4 = w.root(); + assertNotEq(r4, r3); + + // Recent roots should be [r4, r3, r2, r1, 0] + { + uint256[HISTORY] memory recent = w.getRecentRoots(); + assertEq(recent[0], r4); + assertEq(recent[1], r3); + assertEq(recent[2], r2); + assertEq(recent[3], r1); + assertEq(recent[4], 0); + } + + // Now expire ids 2 and 3 and erase both in a single full clean-up call + (,, uint256 gStart3, uint32 gDur3,,,,) = w.memberships(3); + vm.warp(gStart3 + gDur3 + 1); + uint256[] memory batchErase = new uint256[](2); + batchErase[0] = 2; + batchErase[1] = 3; + // Capture the root before the batch to verify only one new history entry is pushed + uint256 preBatchRoot = w.root(); + w.eraseMemberships(batchErase, true); + uint256 r5 = w.root(); + assertNotEq(r5, preBatchRoot); + + // Expect exactly one new entry for the batch (final root), so recent = [r5, r4, r3, r2, r1] + { + uint256[HISTORY] memory recent = w.getRecentRoots(); + assertEq(recent[0], r5); + assertEq(recent[1], r4); + assertEq(recent[2], r3); + assertEq(recent[3], r2); + assertEq(recent[4], r1); + } + } }