Merge pull request #112 from lidofinance/fix/tiebreaker-sealable-bloc…

…kers-fix Tiebreaker add/remove sealable now revertable
lidofinance · Sep 9, 2024 · accfe4f · accfe4f
2 parents 315f43e + c169036
commit accfe4f
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 10 deletions.
diff --git a/contracts/libraries/Tiebreaker.sol b/contracts/libraries/Tiebreaker.sol
@@ -25,6 +25,8 @@ library Tiebreaker {
     error InvalidTiebreakerActivationTimeout(Duration timeout);
     error CallerIsNotTiebreakerCommittee(address caller);
     error SealableWithdrawalBlockersLimitReached();
+    error SealableWithdrawalBlockerNotFound(address sealable);
+    error SealableWithdrawalBlockerAlreadyAdded(address sealable);
 
     event SealableWithdrawalBlockerAdded(address sealable);
     event SealableWithdrawalBlockerRemoved(address sealable);
@@ -67,20 +69,20 @@ library Tiebreaker {
             revert InvalidSealable(sealableWithdrawalBlocker);
         }
 
-        bool isSuccessfullyAdded = self.sealableWithdrawalBlockers.add(sealableWithdrawalBlocker);
-        if (isSuccessfullyAdded) {
-            emit SealableWithdrawalBlockerAdded(sealableWithdrawalBlocker);
+        if (!self.sealableWithdrawalBlockers.add(sealableWithdrawalBlocker)) {
+            revert SealableWithdrawalBlockerAlreadyAdded(sealableWithdrawalBlocker);
         }
+        emit SealableWithdrawalBlockerAdded(sealableWithdrawalBlocker);
     }
 
     /// @notice Removes a sealable withdrawal blocker.
     /// @param self The context storage.
     /// @param sealableWithdrawalBlocker The address of the sealable withdrawal blocker to remove.
     function removeSealableWithdrawalBlocker(Context storage self, address sealableWithdrawalBlocker) internal {
-        bool isSuccessfullyRemoved = self.sealableWithdrawalBlockers.remove(sealableWithdrawalBlocker);
-        if (isSuccessfullyRemoved) {
-            emit SealableWithdrawalBlockerRemoved(sealableWithdrawalBlocker);
+        if (!self.sealableWithdrawalBlockers.remove(sealableWithdrawalBlocker)) {
+            revert SealableWithdrawalBlockerNotFound(sealableWithdrawalBlocker);
         }
+        emit SealableWithdrawalBlockerRemoved(sealableWithdrawalBlocker);
     }
 
     /// @notice Sets the tiebreaker committee.

diff --git a/test/unit/libraries/Tiebreaker.t.sol b/test/unit/libraries/Tiebreaker.t.sol
@@ -39,16 +39,25 @@ contract TiebreakerTest is UnitTest {
         Tiebreaker.addSealableWithdrawalBlocker(context, address(mockSealable2), 1);
     }
 
+    function test_addSealableWithdrawalBlocker_RevertOn_AlreadyAdded() external {
+        Tiebreaker.addSealableWithdrawalBlocker(context, address(mockSealable1), 2);
+
+        vm.expectRevert(
+            abi.encodeWithSelector(Tiebreaker.SealableWithdrawalBlockerAlreadyAdded.selector, address(mockSealable1))
+        );
+        this.external__addSealableWithdrawalBlocker(address(mockSealable1), 2);
+    }
+
     function test_addSealableWithdrawalBlocker_RevertOn_InvalidSealable() external {
         mockSealable1.setShouldRevertIsPaused(true);
 
         vm.expectRevert(abi.encodeWithSelector(Tiebreaker.InvalidSealable.selector, address(mockSealable1)));
         // external call should be used to intercept the revert
-        this.external__addSealableWithdrawalBlocker(address(mockSealable1));
+        this.external__addSealableWithdrawalBlocker(address(mockSealable1), 2);
 
         vm.expectRevert();
         // external call should be used to intercept the revert
-        this.external__addSealableWithdrawalBlocker(address(0x123));
+        this.external__addSealableWithdrawalBlocker(address(0x123), 2);
     }
 
     function test_removeSealableWithdrawalBlocker_HappyPath() external {
@@ -62,6 +71,13 @@ contract TiebreakerTest is UnitTest {
         assertFalse(context.sealableWithdrawalBlockers.contains(address(mockSealable1)));
     }
 
+    function test_removeSealableWithdrawalBlocker_RevertOn_NotFound() external {
+        vm.expectRevert(
+            abi.encodeWithSelector(Tiebreaker.SealableWithdrawalBlockerNotFound.selector, address(mockSealable1))
+        );
+        this.external__removeSealableWithdrawalBlocker(address(mockSealable1));
+    }
+
     function test_setTiebreakerCommittee_HappyPath() external {
         address newCommittee = address(0x123);
 
@@ -188,7 +204,11 @@ contract TiebreakerTest is UnitTest {
         Tiebreaker.checkCallerIsTiebreakerCommittee(context);
     }
 
-    function external__addSealableWithdrawalBlocker(address sealable) external {
-        Tiebreaker.addSealableWithdrawalBlocker(context, sealable, 1);
+    function external__addSealableWithdrawalBlocker(address sealable, uint256 count) external {
+        Tiebreaker.addSealableWithdrawalBlocker(context, sealable, count);
+    }
+
+    function external__removeSealableWithdrawalBlocker(address sealable) external {
+        Tiebreaker.removeSealableWithdrawalBlocker(context, sealable);
     }
 }