update MiddlewareProtect to use tstore

contracts/libraries/ReentrancyState.sol

@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.20;
+
+library ReentrancyState {
+    // bytes32(uint256(keccak256("ReentrancyState")) - 1)
+    bytes32 constant REENTRANCY_STATE_SLOT = 0xbedc9a60a226d4ae7b727cbc828f66c94c4eead57777428ceab2f04b0efca3a5;
+
+    function unlock() internal {
+        assembly {
+            tstore(REENTRANCY_STATE_SLOT, 0)
+        }
+    }
+
+    function lockSwap() internal {
+        assembly {
+            tstore(REENTRANCY_STATE_SLOT, 1)
+        }
+    }
+
+    function lockSwapRemove() internal {
+        assembly {
+            tstore(REENTRANCY_STATE_SLOT, 2)
+        }
+    }
+
+    function read() internal view returns (uint256 state) {
+        assembly {
+            state := tload(REENTRANCY_STATE_SLOT)
+        }
+    }
+
+    function swapLocked() internal view returns (bool) {
+        return read() == 1 || read() == 2;
+    }
+
+    function removeLocked() internal view returns (bool) {
+        return read() == 2;
+    }
+}

contracts/middleware/MiddlewareProtect.sol

@@ -8,26 +8,24 @@ import {BalanceDelta, BalanceDeltaLibrary} from "@uniswap/v4-core/src/types/Bala
 import {BeforeSwapDelta} from "@uniswap/v4-core/src/types/BeforeSwapDelta.sol";
 import {console} from "../../lib/forge-std/src/console.sol";
 import {BaseHook} from "./../BaseHook.sol";
+import {ReentrancyState} from "./../libraries/ReentrancyState.sol";
 
 contract MiddlewareProtect is BaseMiddleware {
-    bool private swapBlocked;
-    bool private removeBlocked;
-
     uint256 public constant gasLimit = 100000;
 
     error ActionBetweenHook();
 
     constructor(IPoolManager _poolManager, address _impl) BaseMiddleware(_poolManager, _impl) {}
 
-    modifier swapNotBlocked() {
-        if (swapBlocked) {
+    modifier swapNotLocked() {
+        if (ReentrancyState.swapLocked()) {
             revert ActionBetweenHook();
         }
         _;
     }
 
-    modifier removeNotBlocked() {
-        if (removeBlocked) {
+    modifier removeNotLocked() {
+        if (ReentrancyState.removeLocked()) {
             revert ActionBetweenHook();
         }
         _;
@@ -36,16 +34,14 @@ contract MiddlewareProtect is BaseMiddleware {
     // block swaps and removes
     function beforeSwap(address, PoolKey calldata, IPoolManager.SwapParams calldata, bytes calldata)
         external
-        swapNotBlocked
+        swapNotLocked
         returns (bytes4, BeforeSwapDelta, uint24)
     {
-        swapBlocked = true;
-        removeBlocked = true;
+        ReentrancyState.lockSwapRemove();
         console.log("beforeSwap middleware");
         (bool success, bytes memory returnData) = implementation.delegatecall{gas: gasLimit}(msg.data);
         require(success);
-        swapBlocked = false;
-        removeBlocked = false;
+        ReentrancyState.unlock();
         return abi.decode(returnData, (bytes4, BeforeSwapDelta, uint24));
     }
 
@@ -56,11 +52,11 @@ contract MiddlewareProtect is BaseMiddleware {
         external
         returns (bytes4)
     {
-        swapBlocked = true;
+        ReentrancyState.lockSwap();
         console.log("beforeAddLiquidity middleware");
         (bool success, bytes memory returnData) = implementation.delegatecall{gas: gasLimit}(msg.data);
         require(success);
-        swapBlocked = false;
+        ReentrancyState.unlock();
         return abi.decode(returnData, (bytes4));
     }
 
@@ -72,11 +68,11 @@ contract MiddlewareProtect is BaseMiddleware {
         PoolKey calldata,
         IPoolManager.ModifyLiquidityParams calldata,
         bytes calldata
-    ) external removeNotBlocked returns (bytes4) {
-        swapBlocked = true;
+    ) external removeNotLocked returns (bytes4) {
+        ReentrancyState.lockSwap();
         console.log("beforeRemoveLiquidity middleware");
         implementation.delegatecall{gas: gasLimit}(msg.data);
-        swapBlocked = false;
+        ReentrancyState.unlock();
         return BaseHook.beforeRemoveLiquidity.selector;
     }
 
@@ -91,4 +87,4 @@ contract MiddlewareProtect is BaseMiddleware {
         implementation.delegatecall{gas: gasLimit}(msg.data);
         return (BaseHook.afterRemoveLiquidity.selector, BalanceDeltaLibrary.ZERO_DELTA);
     }
-}
+}