From 97a3627e089c9171b776d5d0db347b835700fa51 Mon Sep 17 00:00:00 2001
From: Junion <69495294+Jun1on@users.noreply.github.com>
Date: Mon, 1 Jul 2024 11:29:29 -0400
Subject: [PATCH] update MiddlewareProtect to use tstore

---
 contracts/libraries/ReentrancyState.sol    | 39 ++++++++++++++++++++++
 contracts/middleware/MiddlewareProtect.sol | 32 ++++++++----------
 lib/forge-gas-snapshot                     |  2 +-
 lib/forge-std                              |  2 +-
 lib/openzeppelin-contracts                 |  2 +-
 5 files changed, 56 insertions(+), 21 deletions(-)
 create mode 100644 contracts/libraries/ReentrancyState.sol

diff --git a/contracts/libraries/ReentrancyState.sol b/contracts/libraries/ReentrancyState.sol
new file mode 100644
index 00000000..966d1269
--- /dev/null
+++ b/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;
+    }
+}
diff --git a/contracts/middleware/MiddlewareProtect.sol b/contracts/middleware/MiddlewareProtect.sol
index faeb21ae..0ee8a20d 100644
--- a/contracts/middleware/MiddlewareProtect.sol
+++ b/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);
     }
-}
+}
\ No newline at end of file
diff --git a/lib/forge-gas-snapshot b/lib/forge-gas-snapshot
index 2f884282..9161f7c0 160000
--- a/lib/forge-gas-snapshot
+++ b/lib/forge-gas-snapshot
@@ -1 +1 @@
-Subproject commit 2f884282b4cd067298e798974f5b534288b13bc2
+Subproject commit 9161f7c0b6c6788a89081e2b3b9c67592b71e689
diff --git a/lib/forge-std b/lib/forge-std
index 2b58ecbc..75b3fcf0 160000
--- a/lib/forge-std
+++ b/lib/forge-std
@@ -1 +1 @@
-Subproject commit 2b58ecbcf3dfde7a75959dc7b4eb3d0670278de6
+Subproject commit 75b3fcf052cc7886327e4c2eac3d1a1f36942b41
diff --git a/lib/openzeppelin-contracts b/lib/openzeppelin-contracts
index 337bfd5e..5ae63068 160000
--- a/lib/openzeppelin-contracts
+++ b/lib/openzeppelin-contracts
@@ -1 +1 @@
-Subproject commit 337bfd5ea4df9f7ebc755cd3cb4ecb3bd3d33fc7
+Subproject commit 5ae630684a0f57de400ef69499addab4c32ac8fb