diff --git a/src/abstract/AbstractRecipientRegistry.sol b/src/abstract/AbstractRecipientRegistry.sol index 5b8051c..66aa391 100644 --- a/src/abstract/AbstractRecipientRegistry.sol +++ b/src/abstract/AbstractRecipientRegistry.sol @@ -24,9 +24,15 @@ abstract contract AbstractRecipientRegistry is IRecipientRegistry, OwnableUpgrad /// @notice Mapping to quickly check if an address is an active recipient /// @dev Maps recipient address to true if active, false otherwise mapping(address => bool) isRecipientMapping; + /// @notice Mapping to quickly check if an address is queued for addition + /// @dev Maps recipient address to true if queued for addition, false otherwise + mapping(address => bool) isQueuedForAdditionMapping; + /// @notice Mapping to quickly check if an address is queued for removal + /// @dev Maps recipient address to true if queued for removal, false otherwise + mapping(address => bool) isQueuedForRemovalMapping; } - // keccak256(abi.encode(uint256(keccak256("crowdstake.storage.AbstractRecipientRegistry")) - 1)) & ~bytes32(uint256(0xff)) + /// keccak256(abi.encode(uint256(keccak256("crowdstake.storage.AbstractRecipientRegistry")) - 1)) & ~bytes32(uint256(0xff)) bytes32 private constant ABSTRACT_RECIPIENT_REGISTRY_STORAGE = 0x347caeef91698b68f09c13de18e96db5bda028445fd11b86dc029946f360f200; @@ -36,6 +42,9 @@ abstract contract AbstractRecipientRegistry is IRecipientRegistry, OwnableUpgrad } } + /// @notice Maximum queue size + uint256 private constant MAX_QUEUE_SIZE = 100; + // ============ Public Getters ============ /// @notice Array of active recipient addresses @@ -89,14 +98,10 @@ abstract contract AbstractRecipientRegistry is IRecipientRegistry, OwnableUpgrad AbstractRecipientRegistryStorage storage $ = _getAbstractRecipientRegistryStorage(); if (recipient == address(0)) revert InvalidRecipient(); if ($.isRecipientMapping[recipient]) revert RecipientAlreadyExists(); + if ($.isQueuedForAdditionMapping[recipient]) revert RecipientAlreadyQueued(); + if ($.queuedRecipientsForAddition.length + 1 > MAX_QUEUE_SIZE) revert MaxQueueSizeReached(); - // Check if already queued to prevent duplicates - for (uint256 i = 0; i < $.queuedRecipientsForAddition.length; i++) { - if ($.queuedRecipientsForAddition[i] == recipient) { - revert RecipientAlreadyQueued(); - } - } - + $.isQueuedForAdditionMapping[recipient] = true; $.queuedRecipientsForAddition.push(recipient); emit RecipientQueued(recipient, true); } @@ -109,15 +114,12 @@ abstract contract AbstractRecipientRegistry is IRecipientRegistry, OwnableUpgrad /// @dev Access control should be implemented in the calling public function function _queueForRemoval(address recipient) internal { AbstractRecipientRegistryStorage storage $ = _getAbstractRecipientRegistryStorage(); + if (recipient == address(0)) revert InvalidRecipient(); if (!$.isRecipientMapping[recipient]) revert RecipientNotFound(); + if ($.isQueuedForRemovalMapping[recipient]) revert RecipientAlreadyQueued(); + if ($.queuedRecipientsForRemoval.length + 1 > MAX_QUEUE_SIZE) revert MaxQueueSizeReached(); - // Check if already queued for removal to prevent duplicates - for (uint256 i = 0; i < $.queuedRecipientsForRemoval.length; i++) { - if ($.queuedRecipientsForRemoval[i] == recipient) { - revert RecipientAlreadyQueued(); - } - } - + $.isQueuedForRemovalMapping[recipient] = true; $.queuedRecipientsForRemoval.push(recipient); emit RecipientQueued(recipient, false); } @@ -144,6 +146,7 @@ abstract contract AbstractRecipientRegistry is IRecipientRegistry, OwnableUpgrad address recipient = addedList[i]; $.recipients.push(recipient); $.isRecipientMapping[recipient] = true; + $.isQueuedForAdditionMapping[recipient] = false; emit RecipientAdded(recipient); } @@ -161,6 +164,7 @@ abstract contract AbstractRecipientRegistry is IRecipientRegistry, OwnableUpgrad if (recipient == removedList[j]) { shouldRemove = true; $.isRecipientMapping[recipient] = false; + $.isQueuedForRemovalMapping[recipient] = false; emit RecipientRemoved(recipient); break; } @@ -184,14 +188,22 @@ abstract contract AbstractRecipientRegistry is IRecipientRegistry, OwnableUpgrad /// @dev Only owner can clear the queue. Use this to cancel all pending additions /// @dev This will remove all addresses from the addition queue without adding them function clearAdditionQueue() external onlyOwner { - delete _getAbstractRecipientRegistryStorage().queuedRecipientsForAddition; + AbstractRecipientRegistryStorage storage $ = _getAbstractRecipientRegistryStorage(); + for (uint256 i = 0; i < $.queuedRecipientsForAddition.length; i++) { + $.isQueuedForAdditionMapping[$.queuedRecipientsForAddition[i]] = false; + } + delete $.queuedRecipientsForAddition; } /// @notice Clear the removal queue without processing /// @dev Only owner can clear the queue. Use this to cancel all pending removals /// @dev This will remove all addresses from the removal queue without removing them function clearRemovalQueue() external onlyOwner { - delete _getAbstractRecipientRegistryStorage().queuedRecipientsForRemoval; + AbstractRecipientRegistryStorage storage $ = _getAbstractRecipientRegistryStorage(); + for (uint256 i = 0; i < $.queuedRecipientsForRemoval.length; i++) { + $.isQueuedForRemovalMapping[$.queuedRecipientsForRemoval[i]] = false; + } + delete $.queuedRecipientsForRemoval; } /// @notice Get all active recipients @@ -226,26 +238,14 @@ abstract contract AbstractRecipientRegistry is IRecipientRegistry, OwnableUpgrad /// @param recipient Address to check in the addition queue /// @return isQueued True if the address is queued for addition, false otherwise function isQueuedForAddition(address recipient) external view returns (bool isQueued) { - AbstractRecipientRegistryStorage storage $ = _getAbstractRecipientRegistryStorage(); - for (uint256 i = 0; i < $.queuedRecipientsForAddition.length; i++) { - if ($.queuedRecipientsForAddition[i] == recipient) { - return true; - } - } - return false; + return _getAbstractRecipientRegistryStorage().isQueuedForAdditionMapping[recipient]; } /// @notice Check if an address is queued for removal /// @param recipient Address to check in the removal queue /// @return isQueued True if the address is queued for removal, false otherwise function isQueuedForRemoval(address recipient) external view returns (bool isQueued) { - AbstractRecipientRegistryStorage storage $ = _getAbstractRecipientRegistryStorage(); - for (uint256 i = 0; i < $.queuedRecipientsForRemoval.length; i++) { - if ($.queuedRecipientsForRemoval[i] == recipient) { - return true; - } - } - return false; + return _getAbstractRecipientRegistryStorage().isQueuedForRemovalMapping[recipient]; } /// @notice Check if an address is currently an active recipient diff --git a/src/interfaces/IRecipientRegistry.sol b/src/interfaces/IRecipientRegistry.sol index a332f85..c872768 100644 --- a/src/interfaces/IRecipientRegistry.sol +++ b/src/interfaces/IRecipientRegistry.sol @@ -40,6 +40,9 @@ interface IRecipientRegistry { /// @notice Thrown when attempting to queue a recipient that is already queued error RecipientAlreadyQueued(); + /// @notice Thrown when attempting to queue a recipient when the queue reached the maximum size + error MaxQueueSizeReached(); + /// @notice Queue a recipient for addition to the registry /// @dev Access control varies by implementation (admin-only vs recipient voting) /// @dev The recipient will be added when the queue is processed diff --git a/test/RecipientRegistry.t.sol b/test/RecipientRegistry.t.sol index 7fb0861..d930b47 100644 --- a/test/RecipientRegistry.t.sol +++ b/test/RecipientRegistry.t.sol @@ -175,6 +175,46 @@ contract RecipientRegistryTest is TestWrapper { assertEq(registry.getQueuedAdditions().length, 0); } + function test_ReQueueRecipientAfterRemoval() public { + registry.queueRecipientAddition(RECIPIENT_1); + registry.processQueue(); + + registry.queueRecipientRemoval(RECIPIENT_1); + registry.processQueue(); + + assertFalse(registry.isRecipient(RECIPIENT_1)); + assertFalse(registry.isQueuedForRemoval(RECIPIENT_1)); + + registry.queueRecipientAddition(RECIPIENT_1); + + assertTrue(registry.isQueuedForAddition(RECIPIENT_1)); + assertEq(registry.getQueuedAdditions().length, 1); + + registry.processQueue(); + + assertTrue(registry.isRecipient(RECIPIENT_1)); + assertFalse(registry.isQueuedForAddition(RECIPIENT_1)); + assertEq(registry.getRecipientCount(), 1); + } + + function test_ReQueueRecipientAfterClearingAdditionQueue() public { + registry.queueRecipientAddition(RECIPIENT_1); + + registry.clearAdditionQueue(); + + assertEq(registry.getQueuedAdditions().length, 0); + assertFalse(registry.isQueuedForAddition(RECIPIENT_1)); + + registry.queueRecipientAddition(RECIPIENT_1); + assertTrue(registry.isQueuedForAddition(RECIPIENT_1)); + + registry.processQueue(); + + assertTrue(registry.isRecipient(RECIPIENT_1)); + assertFalse(registry.isQueuedForAddition(RECIPIENT_1)); + assertEq(registry.getRecipientCount(), 1); + } + function test_ClearRemovalQueue() public { // Add recipients first registry.queueRecipientAddition(RECIPIENT_1); @@ -221,7 +261,7 @@ contract RecipientRegistryTest is TestWrapper { function test_RevertOnDuplicateQueuedAddition() public { registry.queueRecipientAddition(RECIPIENT_1); - vm.expectRevert(); + vm.expectRevert(IRecipientRegistry.RecipientAlreadyQueued.selector); registry.queueRecipientAddition(RECIPIENT_1); } @@ -236,10 +276,29 @@ contract RecipientRegistryTest is TestWrapper { registry.queueRecipientRemoval(RECIPIENT_1); - vm.expectRevert(); + vm.expectRevert(IRecipientRegistry.RecipientAlreadyQueued.selector); registry.queueRecipientRemoval(RECIPIENT_1); } + function test_AdditionAndRemovalQueuesAreIndependent() public { + registry.queueRecipientAddition(RECIPIENT_1); + registry.processQueue(); + + registry.queueRecipientAddition(RECIPIENT_2); + registry.queueRecipientRemoval(RECIPIENT_1); + + assertTrue(registry.isQueuedForAddition(RECIPIENT_2)); + assertTrue(registry.isQueuedForRemoval(RECIPIENT_1)); + assertFalse(registry.isQueuedForAddition(RECIPIENT_1)); + assertFalse(registry.isQueuedForRemoval(RECIPIENT_2)); + + registry.processQueue(); + + assertFalse(registry.isRecipient(RECIPIENT_1)); + assertTrue(registry.isRecipient(RECIPIENT_2)); + assertEq(registry.getRecipientCount(), 1); + } + function test_OnlyOwnerCanQueue() public { vm.prank(address(0xdead)); vm.expectRevert();