From 266c61e386c7108a4cc312c3d977fb5bf50cc91a Mon Sep 17 00:00:00 2001
From: Junion <69495294+Jun1on@users.noreply.github.com>
Date: Thu, 11 Jul 2024 19:02:56 -0400
Subject: [PATCH] fix bugs

---
 contracts/middleware/MiddlewareRemove.sol | 13 +++++++++----
 test/MiddlewareRemoveFactory.t.sol        |  2 +-
 test/middleware/RemoveOutOfGas.sol        |  1 +
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/contracts/middleware/MiddlewareRemove.sol b/contracts/middleware/MiddlewareRemove.sol
index 7333d2b4..c2a72f01 100644
--- a/contracts/middleware/MiddlewareRemove.sol
+++ b/contracts/middleware/MiddlewareRemove.sol
@@ -13,6 +13,7 @@ import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
 import {CustomRevert} from "@uniswap/v4-core/src/libraries/CustomRevert.sol";
 import {NonZeroDeltaCount} from "@uniswap/v4-core/src/libraries/NonZeroDeltaCount.sol";
 import {IExttload} from "@uniswap/v4-core/src/interfaces/IExttload.sol";
+import {console} from "./../../lib/forge-gas-snapshot/lib/forge-std/src/console.sol";
 
 contract MiddlewareRemove is BaseMiddleware {
     using CustomRevert for bytes4;
@@ -28,7 +29,7 @@ contract MiddlewareRemove is BaseMiddleware {
     }
 
     bytes internal constant ZERO_BYTES = bytes("");
-    uint256 public constant gasLimit = 10000000;
+    uint256 public constant GAS_LIMIT = 10_000_000;
 
     constructor(IPoolManager _manager, address _impl) BaseMiddleware(_manager, _impl) {
         if (IHooks(address(this)).hasPermission(Hooks.AFTER_REMOVE_LIQUIDITY_RETURNS_DELTA_FLAG)) {
@@ -42,7 +43,7 @@ contract MiddlewareRemove is BaseMiddleware {
         IPoolManager.ModifyLiquidityParams calldata,
         bytes calldata
     ) external returns (bytes4) {
-        implementation.delegatecall{gas: gasLimit}(msg.data);
+        implementation.delegatecall{gas: GAS_LIMIT}(msg.data);
         return BaseHook.beforeRemoveLiquidity.selector;
     }
 
@@ -53,14 +54,18 @@ contract MiddlewareRemove is BaseMiddleware {
         BalanceDelta delta,
         bytes calldata hookData
     ) external returns (bytes4, BalanceDelta) {
-        address(this).delegatecall(abi.encodeWithSelector(this._callAndEnsureZeroDeltas.selector, msg.data));
+        (bool success,) = address(this).delegatecall{gas: GAS_LIMIT}(
+            abi.encodeWithSelector(this._callAndEnsureZeroDeltas.selector, msg.data)
+        );
+        console.log("secondary", success);
         return (BaseHook.afterRemoveLiquidity.selector, BalanceDeltaLibrary.ZERO_DELTA);
     }
 
     function _callAndEnsureZeroDeltas(bytes calldata data) external {
         bytes32 slot = bytes32(NonZeroDeltaCount.NONZERO_DELTA_COUNT_SLOT);
         uint256 countBefore = uint256(IExttload(address(manager)).exttload(slot));
-        address(implementation).delegatecall(msg.data);
+        (bool success,) = address(implementation).delegatecall(data);
+        console.log("main", success);
         uint256 countAfter = uint256(IExttload(address(manager)).exttload(slot));
         if (countAfter > countBefore) {
             // purpousely revert to cause the whole hook to reset
diff --git a/test/MiddlewareRemoveFactory.t.sol b/test/MiddlewareRemoveFactory.t.sol
index e27042bc..6aa72543 100644
--- a/test/MiddlewareRemoveFactory.t.sol
+++ b/test/MiddlewareRemoveFactory.t.sol
@@ -143,7 +143,7 @@ contract MiddlewareRemoveFactoryTest is Test, Deployers {
         (address hookAddress, bytes32 salt) = HookMiner.find(
             address(factory), flags, type(MiddlewareRemove).creationCode, abi.encode(address(manager), address(counter))
         );
-        vm.expectRevert(MiddlewareRemove.HookPermissionForbidden.selector);
+        vm.expectRevert(abi.encodePacked(bytes16(MiddlewareRemove.HookPermissionForbidden.selector), hookAddress));
         factory.createMiddleware(address(counter), salt);
     }
 
diff --git a/test/middleware/RemoveOutOfGas.sol b/test/middleware/RemoveOutOfGas.sol
index fa4997c8..79f03ae2 100644
--- a/test/middleware/RemoveOutOfGas.sol
+++ b/test/middleware/RemoveOutOfGas.sol
@@ -7,6 +7,7 @@ import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
 import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
 import {BalanceDelta, BalanceDeltaLibrary} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
 import {BaseHook} from "./../../contracts/BaseHook.sol";
+import {console} from "./../../lib/forge-gas-snapshot/lib/forge-std/src/console.sol";
 
 contract RemoveOutOfGas is BaseHook {
     uint256 counter;