Change toggle to add/remove

Author: joaobrunoah

contracts/AngstromBalancer.sol

@@ -376,19 +376,21 @@ contract AngstromBalancer is IBatchRouter, BatchRouterHooks, OwnableAuthenticati
     ***************************************************************************/
 
     /**
-     * @notice Register/unregister nodes that are allowed to unlock the system.
-     * @param nodes The nodes to toggle (register/unregister)
+     * @notice Register a node that is allowed to unlock the system.
+     * @param node The node to register
      */
-    function toggleNodes(address[] memory nodes) external authenticate {
-        for (uint256 i = 0; i < nodes.length; i++) {
-            address node = nodes[i];
-            _angstromValidatorNodes[node] = !_angstromValidatorNodes[node];
-            if (_angstromValidatorNodes[node]) {
-                emit NodeRegistered(node);
-            } else {
-                emit NodeUnregistered(node);
-            }
-        }
+    function addNode(address node) external authenticate {
+        _angstromValidatorNodes[node] = true;
+        emit NodeRegistered(node);
+    }
+
+    /**
+     * @notice Unregister a node that is no longer allowed to unlock the system.
+     * @param node The node to unregister
+     */
+    function removeNode(address node) external authenticate {
+        _angstromValidatorNodes[node] = false;
+        emit NodeUnregistered(node);
     }
 
     /***************************************************************************

test/foundry/AngstromBalancerUnit.t.sol

@@ -15,31 +15,29 @@ import { BaseAngstromTest } from "./utils/BaseAngstromTest.sol";
 contract AngstromBalancerUnitTest is BaseAngstromTest {
     using ArrayHelpers for *;
 
-    function testToggleNodesIsAuthenticated() public {
+    function testAddNodeIsAuthenticated() public {
         vm.prank(alice);
         vm.expectRevert(IAuthentication.SenderNotAllowed.selector);
-        angstromBalancer.toggleNodes([bob].toMemoryArray());
+        angstromBalancer.addNode(bob);
     }
 
-    function testToggleNodes() public {
-        makeAngstromNode(bob);
-        assertTrue(angstromBalancer.isRegisteredNode(bob), "Bob is not a node");
+    function testRemoveNodeIsAuthenticated() public {
+        vm.prank(alice);
+        vm.expectRevert(IAuthentication.SenderNotAllowed.selector);
+        angstromBalancer.removeNode(bob);
     }
 
     function testAddAndRemoveNodes() public {
-        makeAngstromNodes([bob, alice].toMemoryArray());
-        assertTrue(angstromBalancer.isRegisteredNode(bob), "Bob is not a node");
-        assertTrue(angstromBalancer.isRegisteredNode(alice), "Alice is not a node");
-
-        vm.prank(admin);
-        angstromBalancer.toggleNodes([bob].toMemoryArray());
+        addAngstromNode(bob);
+        addAngstromNode(alice);
+        removeAngstromNode(bob);
 
         assertFalse(angstromBalancer.isRegisteredNode(bob), "Bob is still a node");
         assertTrue(angstromBalancer.isRegisteredNode(alice), "Alice is not a node after bob was removed");
     }
 
     function testUnlockAngstromSetsLastUnlockBlockNumber() public {
-        makeAngstromNode(bob);
+        addAngstromNode(bob);
 
         assertEq(angstromBalancer.getLastUnlockBlockNumber(), 0, "Last unlock block number is not 0");
 
@@ -51,8 +49,4 @@ contract AngstromBalancerUnitTest is BaseAngstromTest {
             "Last unlock block number is not the current block number"
         );
     }
-
-    function testGetVault() public view {
-        assertEq(address(angstromBalancer.getVault()), address(vault), "Wrong vault address");
-    }
 }

test/foundry/AngstromHook.t.sol

@@ -51,16 +51,15 @@ contract AngstromHookTest is BaseAngstromTest {
         // If the signature is invalid, the hook reverts (in this case, signer and key do not match).
         (, bytes memory userData) = generateSignatureAndUserData(bob, aliceKey);
 
-        vm.prank(admin);
-        angstromBalancer.toggleNodes([bob].toMemoryArray());
+        addAngstromNode(bob);
 
         vm.expectRevert(AngstromBalancer.InvalidSignature.selector);
         vm.prank(bob);
         router.swapSingleTokenExactIn(pool, dai, usdc, 1e18, 0, MAX_UINT256, false, userData);
     }
 
     function testOnBeforeSwapSucceedsAndSetBlockNumber() public {
-        makeAngstromNode(bob);
+        addAngstromNode(bob);
 
         vm.prank(bob);
         router.swapSingleTokenExactIn(pool, dai, usdc, 1e18, 0, MAX_UINT256, false, bobUserData);
@@ -72,7 +71,7 @@ contract AngstromHookTest is BaseAngstromTest {
     }
 
     function testOnlyOncePerBlock() public {
-        makeAngstromNode(bob);
+        addAngstromNode(bob);
 
         vm.prank(bob);
         angstromBalancer.unlockWithEmptyAttestation(bob, bobSignature);
@@ -119,15 +118,15 @@ contract AngstromHookTest is BaseAngstromTest {
     function testOnBeforeAddLiquidityUnbalancedInvalidSignature() public {
         (, bytes memory userData) = generateSignatureAndUserData(alice, bobKey);
 
-        makeAngstromNode(alice);
+        addAngstromNode(alice);
 
         vm.expectRevert(AngstromBalancer.InvalidSignature.selector);
         vm.prank(alice);
         router.addLiquidityUnbalanced(pool, [FixedPoint.ONE, FixedPoint.ONE].toMemoryArray(), 1e18, false, userData);
     }
 
     function testOnBeforeAddLiquidityUnbalancedSucceedsAndSetBlockNumber() public {
-        makeAngstromNode(alice);
+        addAngstromNode(alice);
 
         vm.prank(alice);
         uint256 expectedBptAmountOut = router.addLiquidityUnbalanced(
@@ -181,15 +180,15 @@ contract AngstromHookTest is BaseAngstromTest {
     function testOnBeforeRemoveLiquidityUnbalancedInvalidSignature() public {
         (, bytes memory userData) = generateSignatureAndUserData(lp, bobKey);
 
-        makeAngstromNode(lp);
+        addAngstromNode(lp);
 
         vm.expectRevert(AngstromBalancer.InvalidSignature.selector);
         vm.prank(lp);
         router.removeLiquiditySingleTokenExactIn(pool, 1e18, dai, 0.1e18, false, userData);
     }
 
     function testOnBeforeRemoveLiquidityUnbalancedSucceedsAndSetBlockNumber() public {
-        makeAngstromNode(lp);
+        addAngstromNode(lp);
 
         Balances memory balancesBefore = getBalances(lp);
 

test/foundry/AngstromRouter.t.sol

@@ -36,7 +36,7 @@ contract AngstromRouterTest is BaseAngstromTest {
     }
 
     function testSwapExactInAlreadyUnlocked() public {
-        makeAngstromNode(bob);
+        addAngstromNode(bob);
 
         SwapPathExactAmountIn[] memory paths;
 
@@ -48,7 +48,7 @@ contract AngstromRouterTest is BaseAngstromTest {
     }
 
     function testSwapExactInUnlocksAngstrom() public {
-        makeAngstromNode(bob);
+        addAngstromNode(bob);
 
         SwapPathExactAmountIn[] memory paths;
 
@@ -76,7 +76,7 @@ contract AngstromRouterTest is BaseAngstromTest {
         ) = angstromBalancer.querySwapExactIn(paths, bob, bytes(""));
         vm.revertTo(snapId);
 
-        makeAngstromNode(bob);
+        addAngstromNode(bob);
 
         vm.prank(bob);
         (uint256[] memory pathAmountsOut, address[] memory tokensOut, uint256[] memory amountsOut) = angstromBalancer
@@ -100,7 +100,7 @@ contract AngstromRouterTest is BaseAngstromTest {
     }
 
     function testSwapExactOutAlreadyUnlocked() public {
-        makeAngstromNode(bob);
+        addAngstromNode(bob);
 
         SwapPathExactAmountOut[] memory paths;
 
@@ -112,7 +112,7 @@ contract AngstromRouterTest is BaseAngstromTest {
     }
 
     function testSwapExactOutUnlocksRouter() public {
-        makeAngstromNode(bob);
+        addAngstromNode(bob);
 
         SwapPathExactAmountOut[] memory paths;
         vm.prank(bob);
@@ -144,7 +144,7 @@ contract AngstromRouterTest is BaseAngstromTest {
         ) = angstromBalancer.querySwapExactOut(paths, bob, bytes(""));
         vm.revertTo(snapId);
 
-        makeAngstromNode(bob);
+        addAngstromNode(bob);
 
         vm.prank(bob);
         (uint256[] memory pathAmountsIn, address[] memory tokensIn, uint256[] memory amountsIn) = angstromBalancer

test/foundry/utils/BaseAngstromTest.sol

@@ -35,6 +35,7 @@ contract BaseAngstromTest is BaseVaultTest {
     }
 
     function createHook() internal override returns (address) {
+        vm.startPrank(admin);
         if (reusingArtifacts) {
             angstromBalancer = AngstromBalancerMock(
                 payable(
@@ -47,27 +48,27 @@ contract BaseAngstromTest is BaseVaultTest {
         } else {
             angstromBalancer = new AngstromBalancerMock(vault, weth, permit2, "AngstromBalancer Mock v1");
         }
-        authorizer.grantRole(angstromBalancer.getActionId(AngstromBalancer.toggleNodes.selector), admin);
+        vm.stopPrank();
 
         return address(angstromBalancer);
     }
 
-    function makeAngstromNode(address account) internal {
-        _ensureEventTogglingNode(account);
+    function addAngstromNode(address account) internal {
+        vm.expectEmit();
+        emit AngstromBalancer.NodeRegistered(account);
 
         vm.prank(admin);
-        angstromBalancer.toggleNodes([account].toMemoryArray());
+        angstromBalancer.addNode(account);
         assertTrue(angstromBalancer.isRegisteredNode(account), "Node registration failed");
     }
 
-    function makeAngstromNodes(address[] memory accounts) internal {
-        for (uint256 i = 0; i < accounts.length; i++) {
-            address account = accounts[i];
-            _ensureEventTogglingNode(account);
-        }
+    function removeAngstromNode(address account) internal {
+        vm.expectEmit();
+        emit AngstromBalancer.NodeUnregistered(account);
 
         vm.prank(admin);
-        angstromBalancer.toggleNodes(accounts);
+        angstromBalancer.removeNode(account);
+        assertFalse(angstromBalancer.isRegisteredNode(account), "Node registration failed");
     }
 
     function generateSignatureAndUserData(