Angstrom Hooks with tests

contracts/AngstromRouter.sol → contracts/AngstromRouterAndHook.sol

@@ -4,8 +4,9 @@ pragma solidity ^0.8.24;
 
 import { IPermit2 } from "permit2/src/interfaces/IPermit2.sol";
 
+import { IBatchRouterQueries } from "@balancer-labs/v3-interfaces/contracts/vault/IBatchRouterQueries.sol";
 import { IWETH } from "@balancer-labs/v3-interfaces/contracts/solidity-utils/misc/IWETH.sol";
-import { IRouter } from "@balancer-labs/v3-interfaces/contracts/vault/IRouter.sol";
+import { IBatchRouter } from "@balancer-labs/v3-interfaces/contracts/vault/IBatchRouter.sol";
 import { IVault } from "@balancer-labs/v3-interfaces/contracts/vault/IVault.sol";
 import {
     SwapPathExactAmountIn,
@@ -18,8 +19,7 @@ import { SingletonAuthentication } from "@balancer-labs/v3-vault/contracts/Singl
 import { BatchRouterHooks } from "@balancer-labs/v3-vault/contracts/BatchRouterHooks.sol";
 import { RouterCommon } from "@balancer-labs/v3-vault/contracts/RouterCommon.sol";
 
-// TODO interface
-contract AngstromRouter is BatchRouterHooks, SingletonAuthentication {
+contract AngstromRouterAndHook is IBatchRouter, BatchRouterHooks, SingletonAuthentication {
     uint256 internal _lastUnlockBlockNumber;
 
     error OnlyOncePerBlock();
@@ -39,7 +39,7 @@ contract AngstromRouter is BatchRouterHooks, SingletonAuthentication {
     /***************************************************************************
                                        Swaps
     ***************************************************************************/
-
+    /// @inheritdoc IBatchRouter
     function swapExactIn(
         SwapPathExactAmountIn[] memory paths,
         uint256 deadline,
@@ -51,8 +51,8 @@ contract AngstromRouter is BatchRouterHooks, SingletonAuthentication {
         saveSender(msg.sender)
         returns (uint256[] memory pathAmountsOut, address[] memory tokensOut, uint256[] memory amountsOut)
     {
-        // Unlocks the router in this block. If the router is already unlocked, reverts.
-        _unlockRouter();
+        // Unlocks the Angstrom network in this block. If the Angstrom network is already unlocked, reverts.
+        _unlockAngstrom();
 
         return
             abi.decode(
@@ -72,6 +72,7 @@ contract AngstromRouter is BatchRouterHooks, SingletonAuthentication {
             );
     }
 
+    /// @inheritdoc IBatchRouter
     function swapExactOut(
         SwapPathExactAmountOut[] memory paths,
         uint256 deadline,
@@ -83,8 +84,8 @@ contract AngstromRouter is BatchRouterHooks, SingletonAuthentication {
         saveSender(msg.sender)
         returns (uint256[] memory pathAmountsIn, address[] memory tokensIn, uint256[] memory amountsIn)
     {
-        // Unlocks the router in this block. If the router is already unlocked, reverts.
-        _unlockRouter();
+        // Unlocks the Angstrom Network in this block. If the Angstrom Network is already unlocked, reverts.
+        _unlockAngstrom();
 
         return
             abi.decode(
@@ -108,6 +109,7 @@ contract AngstromRouter is BatchRouterHooks, SingletonAuthentication {
                                      Queries
     ***************************************************************************/
 
+    /// @inheritdoc IBatchRouterQueries
     function querySwapExactIn(
         SwapPathExactAmountIn[] memory paths,
         address sender,
@@ -139,6 +141,7 @@ contract AngstromRouter is BatchRouterHooks, SingletonAuthentication {
             );
     }
 
+    /// @inheritdoc IBatchRouterQueries
     function querySwapExactOut(
         SwapPathExactAmountOut[] memory paths,
         address sender,
@@ -174,11 +177,10 @@ contract AngstromRouter is BatchRouterHooks, SingletonAuthentication {
                                 Nodes Management
     ***************************************************************************/
 
-    function addNode(address node) external authenticate {
-        _nodes[node] = true;
-    }
-    function removeNode(address node) external authenticate {
-        _nodes[node] = false;
+    function toggleNodes(address[] memory nodes) external authenticate {
+        for (uint256 i = 0; i < nodes.length; i++) {
+            _nodes[nodes[i]] = !_nodes[nodes[i]];
+        }
     }
 
     /***************************************************************************
@@ -193,9 +195,9 @@ contract AngstromRouter is BatchRouterHooks, SingletonAuthentication {
                                 Private Functions
     ***************************************************************************/
 
-    function _unlockRouter() internal {
+    function _unlockAngstrom() internal {
         // The router can only be unlocked once per block.
-        if (_isRouterUnlocked()) {
+        if (_isAngstromUnlocked()) {
             revert OnlyOncePerBlock();
         }
 
@@ -211,7 +213,7 @@ contract AngstromRouter is BatchRouterHooks, SingletonAuthentication {
         return _nodes[account];
     }
 
-    function _isRouterUnlocked() internal view returns (bool) {
+    function _isAngstromUnlocked() internal view returns (bool) {
         return _lastUnlockBlockNumber == block.number;
     }
 }

contracts/test/AngstromRouterMock.sol → contracts/test/AngstromRouterAndHookMock.sol

@@ -7,20 +7,20 @@ import { IPermit2 } from "permit2/src/interfaces/IPermit2.sol";
 import { IWETH } from "@balancer-labs/v3-interfaces/contracts/solidity-utils/misc/IWETH.sol";
 import { IVault } from "@balancer-labs/v3-interfaces/contracts/vault/IVault.sol";
 
-import { AngstromRouter } from "../AngstromRouter.sol";
+import { AngstromRouterAndHook } from "../AngstromRouterAndHook.sol";
 
-contract AngstromRouterMock is AngstromRouter {
+contract AngstromRouterAndHookMock is AngstromRouterAndHook {
     constructor(
         IVault vault,
         IWETH weth,
         IPermit2 permit2,
         string memory routerVersion
-    ) AngstromRouter(vault, weth, permit2, routerVersion) {
+    ) AngstromRouterAndHook(vault, weth, permit2, routerVersion) {
         // solhint-disable-previous-line no-empty-blocks
     }
 
-    function manualUnlockRouter() external {
-        _unlockRouter();
+    function manualUnlockAngstrom() external {
+        _unlockAngstrom();
     }
 
     function getLastUnlockBlockNumber() external view returns (uint256) {

test/foundry/AngstromRouterAndHook.t.sol

@@ -0,0 +1,96 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+pragma solidity ^0.8.24;
+
+import "forge-std/Test.sol";
+
+import { IAuthentication } from "@balancer-labs/v3-interfaces/contracts/solidity-utils/helpers/IAuthentication.sol";
+import {
+    SwapPathExactAmountIn,
+    SwapPathExactAmountOut
+} from "@balancer-labs/v3-interfaces/contracts/vault/BatchRouterTypes.sol";
+
+import { ArrayHelpers } from "@balancer-labs/v3-solidity-utils/contracts/test/ArrayHelpers.sol";
+import { BaseVaultTest } from "@balancer-labs/v3-vault/test/foundry/utils/BaseVaultTest.sol";
+
+import { AngstromRouterAndHookMock } from "../../contracts/test/AngstromRouterAndHookMock.sol";
+import { AngstromRouterAndHook } from "../../contracts/AngstromRouterAndHook.sol";
+
+contract AngstromRouterAndHookTest is BaseVaultTest {
+    using ArrayHelpers for *;
+
+    AngstromRouterAndHookMock private _angstromRouterAndHook;
+
+    function setUp() public virtual override {
+        super.setUp();
+
+        _angstromRouterAndHook = new AngstromRouterAndHookMock(vault, weth, permit2, "AngstromRouterAndHook Mock v1");
+
+        authorizer.grantRole(_angstromRouterAndHook.getActionId(AngstromRouterAndHook.toggleNodes.selector), admin);
+    }
+
+    function testSwapExactInNotNode() public {
+        SwapPathExactAmountIn[] memory paths;
+        vm.expectRevert(AngstromRouterAndHook.NotNode.selector);
+        _angstromRouterAndHook.swapExactIn(paths, MAX_UINT256, false, bytes(""));
+    }
+
+    function testSwapExactInAlreadyUnlocked() public {
+        vm.prank(admin);
+        _angstromRouterAndHook.toggleNodes([bob].toMemoryArray());
+
+        SwapPathExactAmountIn[] memory paths;
+
+        vm.startPrank(bob);
+        _angstromRouterAndHook.manualUnlockAngstrom();
+        vm.expectRevert(AngstromRouterAndHook.OnlyOncePerBlock.selector);
+        _angstromRouterAndHook.swapExactIn(paths, MAX_UINT256, false, bytes(""));
+        vm.stopPrank();
+    }
+
+    function testSwapExactInUnlocksAngstrom() public {
+        vm.prank(admin);
+        _angstromRouterAndHook.toggleNodes([bob].toMemoryArray());
+
+        SwapPathExactAmountIn[] memory paths;
+        vm.prank(bob);
+        _angstromRouterAndHook.swapExactIn(paths, MAX_UINT256, false, bytes(""));
+        assertEq(
+            _angstromRouterAndHook.getLastUnlockBlockNumber(),
+            block.number,
+            "Last unlock block number is not the current block number"
+        );
+    }
+
+    function testSwapExactOutNotNode() public {
+        SwapPathExactAmountOut[] memory paths;
+        vm.expectRevert(AngstromRouterAndHook.NotNode.selector);
+        _angstromRouterAndHook.swapExactOut(paths, MAX_UINT256, false, bytes(""));
+    }
+
+    function testSwapExactOutAlreadyUnlocked() public {
+        vm.prank(admin);
+        _angstromRouterAndHook.toggleNodes([bob].toMemoryArray());
+
+        SwapPathExactAmountOut[] memory paths;
+        vm.startPrank(bob);
+        _angstromRouterAndHook.manualUnlockAngstrom();
+        vm.expectRevert(AngstromRouterAndHook.OnlyOncePerBlock.selector);
+        _angstromRouterAndHook.swapExactOut(paths, MAX_UINT256, false, bytes(""));
+        vm.stopPrank();
+    }
+
+    function testSwapExactOutUnlocksRouter() public {
+        vm.prank(admin);
+        _angstromRouterAndHook.toggleNodes([bob].toMemoryArray());
+
+        SwapPathExactAmountOut[] memory paths;
+        vm.prank(bob);
+        _angstromRouterAndHook.swapExactOut(paths, MAX_UINT256, false, bytes(""));
+        assertEq(
+            _angstromRouterAndHook.getLastUnlockBlockNumber(),
+            block.number,
+            "Last unlock block number is not the current block number"
+        );
+    }
+}

test/foundry/AngstromRouterAndHookUnit.t.sol

@@ -0,0 +1,84 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+pragma solidity ^0.8.24;
+
+import "forge-std/Test.sol";
+
+import { IAuthentication } from "@balancer-labs/v3-interfaces/contracts/solidity-utils/helpers/IAuthentication.sol";
+
+import { ArrayHelpers } from "@balancer-labs/v3-solidity-utils/contracts/test/ArrayHelpers.sol";
+import { BaseVaultTest } from "@balancer-labs/v3-vault/test/foundry/utils/BaseVaultTest.sol";
+
+import { AngstromRouterAndHookMock } from "../../contracts/test/AngstromRouterAndHookMock.sol";
+import { AngstromRouterAndHook } from "../../contracts/AngstromRouterAndHook.sol";
+
+contract AngstromRouterAndHookUnitTest is BaseVaultTest {
+    using ArrayHelpers for *;
+
+    AngstromRouterAndHookMock private _angstromRouterAndHook;
+
+    function setUp() public virtual override {
+        super.setUp();
+
+        _angstromRouterAndHook = new AngstromRouterAndHookMock(vault, weth, permit2, "AngstromRouterAndHook Mock v1");
+
+        authorizer.grantRole(_angstromRouterAndHook.getActionId(AngstromRouterAndHook.toggleNodes.selector), admin);
+    }
+
+    function testToggleNodesIsAuthenticated() public {
+        vm.expectRevert(IAuthentication.SenderNotAllowed.selector);
+        _angstromRouterAndHook.toggleNodes([bob].toMemoryArray());
+    }
+
+    function testToggleNodes() public {
+        vm.prank(admin);
+        _angstromRouterAndHook.toggleNodes([bob].toMemoryArray());
+        assertTrue(_angstromRouterAndHook.isNode(bob), "Bob is not a node");
+    }
+
+    function testAddAndRemoveNodes() public {
+        vm.startPrank(admin);
+        _angstromRouterAndHook.toggleNodes([bob, alice].toMemoryArray());
+        assertTrue(_angstromRouterAndHook.isNode(bob), "Bob is not a node");
+        assertTrue(_angstromRouterAndHook.isNode(alice), "Alice is not a node");
+        _angstromRouterAndHook.toggleNodes([bob].toMemoryArray());
+        vm.stopPrank();
+        assertFalse(_angstromRouterAndHook.isNode(bob), "Bob is still a node");
+        assertTrue(_angstromRouterAndHook.isNode(alice), "Alice is not a node after bob was removed");
+    }
+
+    function testUnlockAngstromNotNode() public {
+        vm.expectRevert(AngstromRouterAndHook.NotNode.selector);
+        _angstromRouterAndHook.manualUnlockAngstrom();
+    }
+
+    function testUnlockAngstromTwice() public {
+        vm.prank(admin);
+        _angstromRouterAndHook.toggleNodes([bob].toMemoryArray());
+
+        vm.startPrank(bob);
+        _angstromRouterAndHook.manualUnlockAngstrom();
+        vm.expectRevert(AngstromRouterAndHook.OnlyOncePerBlock.selector);
+        _angstromRouterAndHook.manualUnlockAngstrom();
+        vm.stopPrank();
+    }
+
+    function testUnlockAngstromSetsLastUnlockBlockNumber() public {
+        vm.prank(admin);
+        _angstromRouterAndHook.toggleNodes([bob].toMemoryArray());
+
+        assertEq(_angstromRouterAndHook.getLastUnlockBlockNumber(), 0, "Last unlock block number is not 0");
+
+        vm.prank(bob);
+        _angstromRouterAndHook.manualUnlockAngstrom();
+        assertEq(
+            _angstromRouterAndHook.getLastUnlockBlockNumber(),
+            block.number,
+            "Last unlock block number is not the current block number"
+        );
+    }
+
+    function testGetVault() public {
+        assertEq(address(_angstromRouterAndHook.getVault()), address(vault), "Wrong vault address");
+    }
+}