From 8f9dd75b5fbf1ff34c8557d43b4872f825872407 Mon Sep 17 00:00:00 2001 From: Filip-L Date: Mon, 23 Jun 2025 12:30:18 +0200 Subject: [PATCH 1/2] Remove allocator if allowance equal 0 --- src/Allocator.sol | 6 +++++- test/Allocator.t.sol | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/src/Allocator.sol b/src/Allocator.sol index efd4134..9b4f74e 100644 --- a/src/Allocator.sol +++ b/src/Allocator.sol @@ -112,8 +112,12 @@ contract Allocator is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, I revert AlreadyZero(); } else if (allowanceBefore < amount) { amount = allowanceBefore; + _allocators.remove(allocator); + } else if (allowanceBefore == amount) { + _allocators.remove(allocator); + } else { + _allocators.set(allocator, allowanceBefore - amount); } - _allocators.set(allocator, allowanceBefore - amount); emit AllowanceChanged(allocator, allowanceBefore, allowance(allocator)); } diff --git a/test/Allocator.t.sol b/test/Allocator.t.sol index a56cc87..61ba0f9 100644 --- a/test/Allocator.t.sol +++ b/test/Allocator.t.sol @@ -301,4 +301,20 @@ contract AllocatorTest is Test { emit IAllocator.AllowanceChanged(vm.addr(1), 100, 50); allocator.decreaseAllowance(vm.addr(1), 50); } + + function testDecreaseAllowanceRemoveAllocator() public { + allocator.setAllowance(vm.addr(1), 100); + allocator.setAllowance(vm.addr(2), 100); + address[] memory allocators = allocator.getAllocators(); + assertEq(allocators[0], vm.addr(1)); + assertEq(allocators[1], vm.addr(2)); + assertEq(allocators.length, 2); + allocator.decreaseAllowance(vm.addr(1), 150); + allocators = allocator.getAllocators(); + assertEq(allocators[0], vm.addr(2)); + assertEq(allocators.length, 1); + allocator.decreaseAllowance(vm.addr(2), 100); + allocators = allocator.getAllocators(); + assertEq(allocators.length, 0); + } } From f63857ad0c518fb5f6b5cd74d3d1d7d70f8486cb Mon Sep 17 00:00:00 2001 From: Filip-L Date: Mon, 30 Jun 2025 12:30:57 +0200 Subject: [PATCH 2/2] Changes after CR --- src/Allocator.sol | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/Allocator.sol b/src/Allocator.sol index 9b4f74e..62efe54 100644 --- a/src/Allocator.sol +++ b/src/Allocator.sol @@ -110,10 +110,7 @@ contract Allocator is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, I uint256 allowanceBefore = allowance(allocator); if (allowanceBefore == 0) { revert AlreadyZero(); - } else if (allowanceBefore < amount) { - amount = allowanceBefore; - _allocators.remove(allocator); - } else if (allowanceBefore == amount) { + } else if (allowanceBefore <= amount) { _allocators.remove(allocator); } else { _allocators.set(allocator, allowanceBefore - amount);