Implement swapper different from node.

Author: joaobrunoah

contracts/AngstromBalancer.sol

@@ -68,6 +68,14 @@ contract AngstromBalancer is IBatchRouter, BatchRouterHooks, OwnableAuthenticati
     uint256 internal constant _ATTEST_EMPTY_BLOCK_TYPE_HASH =
         0x3f25e551746414ff93f076a7dd83828ff53735b39366c74015637e004fcb0223;
 
+    /// @dev `keccak256("SwapExactIn(SwapPathExactAmountIn[] paths, uint64 block_number)")`.
+    uint256 internal constant _SWAP_EXACT_IN_TYPE_HASH =
+        0xc3810e534961c3152a90c6e5342d0ef3cdd6517c38c42e01b7640beffc3e41c2;
+
+    /// @dev `keccak256("SwapExactOut(SwapPathExactAmountOut[] paths, uint64 block_number)")`.
+    uint256 internal constant _SWAP_EXACT_OUT_TYPE_HASH =
+        0xb26cc9223a5f7a414a15401ce11a9ef78c5c9daf4bc40c11e20d3f53cb9d79a5;
+
     /// @dev Set of active Angstrom validator nodes, authorized to unlock this contract for operations.
     mapping(address node => bool isActive) internal _angstromValidatorNodes;
 
@@ -159,7 +167,7 @@ contract AngstromBalancer is IBatchRouter, BatchRouterHooks, OwnableAuthenticati
             abi.decode(
                 _vault.unlock(
                     abi.encodeCall(
-                        BatchRouterHooks.swapExactInHook,
+                        AngstromBalancer.swapExactInHookAngstrom,
                         SwapExactInHookParams({
                             sender: msg.sender,
                             paths: paths,
@@ -173,6 +181,32 @@ contract AngstromBalancer is IBatchRouter, BatchRouterHooks, OwnableAuthenticati
             );
     }
 
+    function swapExactInHookAngstrom(
+        SwapExactInHookParams calldata params
+    )
+        external
+        nonReentrant
+        onlyVault
+        returns (uint256[] memory pathAmountsOut, address[] memory tokensOut, uint256[] memory amountsOut)
+    {
+        // validate signature and sender.
+        if (params.userData.length < 20) {
+            revert InvalidSignature();
+        }
+
+        (address payer, bytes memory signature) = _splitUserData(params.userData);
+        // The signature looks well-formed.
+        bytes32 digest = _computeDigestSwapExactIn(params.paths);
+
+        if (SignatureCheckerLib.isValidSignatureNow(payer, digest, signature) == false) {
+            revert InvalidSignature();
+        }
+
+        (pathAmountsOut, tokensOut, amountsOut) = _swapExactInHook(params);
+
+        _settlePaths(payer, params.wethIsEth);
+    }
+
     /// @inheritdoc IBatchRouter
     function swapExactOut(
         SwapPathExactAmountOut[] memory paths,
@@ -193,7 +227,7 @@ contract AngstromBalancer is IBatchRouter, BatchRouterHooks, OwnableAuthenticati
             abi.decode(
                 _vault.unlock(
                     abi.encodeCall(
-                        BatchRouterHooks.swapExactOutHook,
+                        AngstromBalancer.swapExactOutHookAngstrom,
                         SwapExactOutHookParams({
                             sender: msg.sender,
                             paths: paths,
@@ -207,6 +241,32 @@ contract AngstromBalancer is IBatchRouter, BatchRouterHooks, OwnableAuthenticati
             );
     }
 
+    function swapExactOutHookAngstrom(
+        SwapExactOutHookParams calldata params
+    )
+        external
+        nonReentrant
+        onlyVault
+        returns (uint256[] memory pathAmountsIn, address[] memory tokensIn, uint256[] memory amountsIn)
+    {
+        // validate signature and sender.
+        if (params.userData.length < 20) {
+            revert InvalidSignature();
+        }
+
+        (address payer, bytes memory signature) = _splitUserData(params.userData);
+        // The signature looks well-formed.
+        bytes32 digest = _computeDigestSwapExactOut(params.paths);
+
+        if (SignatureCheckerLib.isValidSignatureNow(payer, digest, signature) == false) {
+            revert InvalidSignature();
+        }
+
+        (pathAmountsIn, tokensIn, amountsIn) = _swapExactOutHook(params);
+
+        _settlePaths(payer, params.wethIsEth);
+    }
+
     /***************************************************************************
                                      Queries
     ***************************************************************************/
@@ -447,6 +507,88 @@ contract AngstromBalancer is IBatchRouter, BatchRouterHooks, OwnableAuthenticati
         }
     }
 
+    function _computeDigestSwapExactIn(SwapPathExactAmountIn[] memory paths) internal view returns (bytes32) {
+        // First, hash the paths array according to EIP-712
+        bytes32 pathsHash = _hashSwapExactInPathArray(paths);
+
+        bytes32 swapExactInStructHash;
+        // solhint-disable-next-line no-inline-assembly
+        assembly ("memory-safe") {
+            let ptr := mload(0x40)
+            mstore(ptr, _SWAP_EXACT_IN_TYPE_HASH)
+            mstore(add(ptr, 0x20), pathsHash)
+            mstore(add(ptr, 0x40), number())
+            swapExactInStructHash := keccak256(ptr, 0x60)
+        }
+        return _hashTypedData(swapExactInStructHash);
+    }
+
+    // Helper function to hash the SwapPathExactAmountIn array
+    function _hashSwapExactInPathArray(SwapPathExactAmountIn[] memory paths) internal pure returns (bytes32) {
+        bytes32[] memory pathHashes = new bytes32[](paths.length);
+
+        for (uint256 i = 0; i < paths.length; i++) {
+            pathHashes[i] = _hashSwapExactInPath(paths[i]);
+        }
+
+        return keccak256(abi.encodePacked(pathHashes));
+    }
+
+    // Helper function to hash a single SwapPathExactAmountIn
+    function _hashSwapExactInPath(SwapPathExactAmountIn memory path) internal pure returns (bytes32) {
+        // You'll need to define the type hash for SwapPathExactAmountIn
+        // For now, using a simple encoding (adjust based on your EIP-712 schema)
+        bytes32[] memory stepHashes = new bytes32[](path.steps.length);
+
+        for (uint256 i = 0; i < path.steps.length; i++) {
+            stepHashes[i] = keccak256(abi.encode(path.steps[i].pool, path.steps[i].tokenOut, path.steps[i].isBuffer));
+        }
+
+        bytes32 stepsHash = keccak256(abi.encodePacked(stepHashes));
+
+        return keccak256(abi.encode(path.tokenIn, stepsHash, path.exactAmountIn, path.minAmountOut));
+    }
+
+    function _computeDigestSwapExactOut(SwapPathExactAmountOut[] memory paths) internal view returns (bytes32) {
+        // First, hash the paths array according to EIP-712
+        bytes32 pathsHash = _hashSwapExactOutPathArray(paths);
+
+        bytes32 swapExactOutStructHash;
+        // solhint-disable-next-line no-inline-assembly
+        assembly ("memory-safe") {
+            let ptr := mload(0x40)
+            mstore(ptr, _SWAP_EXACT_OUT_TYPE_HASH)
+            mstore(add(ptr, 0x20), pathsHash)
+            mstore(add(ptr, 0x40), number())
+            swapExactOutStructHash := keccak256(ptr, 0x60)
+        }
+        return _hashTypedData(swapExactOutStructHash);
+    }
+
+    // Helper function to hash the SwapPathExactAmountOut array
+    function _hashSwapExactOutPathArray(SwapPathExactAmountOut[] memory paths) internal pure returns (bytes32) {
+        bytes32[] memory pathHashes = new bytes32[](paths.length);
+
+        for (uint256 i = 0; i < paths.length; i++) {
+            pathHashes[i] = _hashSwapExactOutPath(paths[i]);
+        }
+
+        return keccak256(abi.encodePacked(pathHashes));
+    }
+
+    // Helper function to hash a single SwapPathExactAmountOut
+    function _hashSwapExactOutPath(SwapPathExactAmountOut memory path) internal pure returns (bytes32) {
+        bytes32[] memory stepHashes = new bytes32[](path.steps.length);
+
+        for (uint256 i = 0; i < path.steps.length; i++) {
+            stepHashes[i] = keccak256(abi.encode(path.steps[i].pool, path.steps[i].tokenOut, path.steps[i].isBuffer));
+        }
+
+        bytes32 stepsHash = keccak256(abi.encodePacked(stepHashes));
+
+        return keccak256(abi.encode(path.tokenIn, stepsHash, path.maxAmountIn, path.exactAmountOut));
+    }
+
     function _getDigest() internal view returns (bytes32) {
         bytes32 attestationStructHash;
         // solhint-disable-next-line no-inline-assembly

contracts/test/AngstromBalancerMock.sol

@@ -6,6 +6,7 @@ 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 "@balancer-labs/v3-interfaces/contracts/vault/BatchRouterTypes.sol";
 
 import { AngstromBalancer } from "../AngstromBalancer.sol";
 
@@ -23,6 +24,14 @@ contract AngstromBalancerMock is AngstromBalancer {
         _unlockAngstrom();
     }
 
+    function computeDigestSwapExactIn(SwapPathExactAmountIn[] memory paths) external view returns (bytes32) {
+        return _computeDigestSwapExactIn(paths);
+    }
+
+    function computeDigestSwapExactOut(SwapPathExactAmountOut[] memory paths) external view returns (bytes32) {
+        return _computeDigestSwapExactOut(paths);
+    }
+
     function getDigest() external view returns (bytes32) {
         return _getDigest();
     }

test/foundry/AngstromHook.t.sol

@@ -42,7 +42,7 @@ contract AngstromHookTest is BaseAngstromTest {
 
     function testOnBeforeSwapInvalidSignature() public {
         // If the signature is invalid, the hook reverts (in this case, signer and key do not match).
-        (, bytes memory userData) = generateSignatureAndUserData(bob, aliceKey);
+        (, bytes memory userData) = generateSignatureAndUserDataEmptyAttestation(bob, aliceKey);
 
         registerAngstromNode(bob);
 
@@ -109,7 +109,7 @@ contract AngstromHookTest is BaseAngstromTest {
     }
 
     function testOnBeforeAddLiquidityUnbalancedInvalidSignature() public {
-        (, bytes memory userData) = generateSignatureAndUserData(alice, bobKey);
+        (, bytes memory userData) = generateSignatureAndUserDataEmptyAttestation(alice, bobKey);
 
         registerAngstromNode(alice);
 
@@ -171,7 +171,7 @@ contract AngstromHookTest is BaseAngstromTest {
     }
 
     function testOnBeforeRemoveLiquidityUnbalancedInvalidSignature() public {
-        (, bytes memory userData) = generateSignatureAndUserData(lp, bobKey);
+        (, bytes memory userData) = generateSignatureAndUserDataEmptyAttestation(lp, bobKey);
 
         registerAngstromNode(lp);