Reentrancy guard transient

.github/workflows/forge.yml

@@ -56,7 +56,7 @@ jobs:
 
       - name: Run Forge tests
         run: |
-          forge test -vvv
+          forge test --isolate -vvv
         id: test
         env:
           FORK_URL: https://mainnet.infura.io/v3/${{ secrets.INFURA_API_KEY }}

.github/workflows/lint.yml

@@ -40,4 +40,4 @@ jobs:
         id: build
 
       - name: Run linter
-        run: yarn run prettier
+        run: yarn run lint:check

.github/workflows/test.yml

@@ -47,6 +47,6 @@ jobs:
         run: yarn compile
 
       - name: Run Integration tests
-        run: yarn test
+        run: yarn test:hardhat
         env:
           INFURA_API_KEY: ${{ secrets.INFURA_API_KEY }}

contracts/base/Dispatcher.sol

@@ -13,6 +13,7 @@ import {LockAndMsgSender} from './LockAndMsgSender.sol';
 import {ERC20} from 'solmate/src/tokens/ERC20.sol';
 import {IAllowanceTransfer} from 'permit2/src/interfaces/IAllowanceTransfer.sol';
 import {IERC721Permit} from '@uniswap/v3-periphery/contracts/interfaces/IERC721Permit.sol';
+import {Constants} from '../libraries/Constants.sol';
 
 /// @title Decodes and Executes Commands
 /// @notice Called by the UniversalRouter contract to efficiently decode and execute a singular command
@@ -55,7 +56,7 @@ abstract contract Dispatcher is Payments, V2SwapRouter, V3SwapRouter, V3ToV4Migr
                             payerIsUser := calldataload(add(inputs.offset, 0x80))
                         }
                         bytes calldata path = inputs.toBytes(3);
-                        address payer = payerIsUser ? lockedBy : address(this);
+                        address payer = payerIsUser ? _msgSender() : address(this);
                         v3SwapExactInput(map(recipient), amountIn, amountOutMin, path, payer);
                     } else if (command == Commands.V3_SWAP_EXACT_OUT) {
                         // equivalent: abi.decode(inputs, (address, uint256, uint256, bytes, bool))
@@ -71,7 +72,7 @@ abstract contract Dispatcher is Payments, V2SwapRouter, V3SwapRouter, V3ToV4Migr
                             payerIsUser := calldataload(add(inputs.offset, 0x80))
                         }
                         bytes calldata path = inputs.toBytes(3);
-                        address payer = payerIsUser ? lockedBy : address(this);
+                        address payer = payerIsUser ? _msgSender() : address(this);
                         v3SwapExactOutput(map(recipient), amountOut, amountInMax, path, payer);
                     } else if (command == Commands.PERMIT2_TRANSFER_FROM) {
                         // equivalent: abi.decode(inputs, (address, address, uint160))
@@ -83,12 +84,12 @@ abstract contract Dispatcher is Payments, V2SwapRouter, V3SwapRouter, V3ToV4Migr
                             recipient := calldataload(add(inputs.offset, 0x20))
                             amount := calldataload(add(inputs.offset, 0x40))
                         }
-                        permit2TransferFrom(token, lockedBy, map(recipient), amount);
+                        permit2TransferFrom(token, _msgSender(), map(recipient), amount);
                     } else if (command == Commands.PERMIT2_PERMIT_BATCH) {
                         (IAllowanceTransfer.PermitBatch memory permitBatch,) =
                             abi.decode(inputs, (IAllowanceTransfer.PermitBatch, bytes));
                         bytes calldata data = inputs.toBytes(1);
-                        PERMIT2.permit(lockedBy, permitBatch, data);
+                        PERMIT2.permit(_msgSender(), permitBatch, data);
                     } else if (command == Commands.SWEEP) {
                         // equivalent:  abi.decode(inputs, (address, address, uint256))
                         address token;
@@ -142,7 +143,7 @@ abstract contract Dispatcher is Payments, V2SwapRouter, V3SwapRouter, V3ToV4Migr
                             payerIsUser := calldataload(add(inputs.offset, 0x80))
                         }
                         address[] calldata path = inputs.toAddressArray(3);
-                        address payer = payerIsUser ? lockedBy : address(this);
+                        address payer = payerIsUser ? _msgSender() : address(this);
                         v2SwapExactInput(map(recipient), amountIn, amountOutMin, path, payer);
                     } else if (command == Commands.V2_SWAP_EXACT_OUT) {
                         // equivalent: abi.decode(inputs, (address, uint256, uint256, bytes, bool))
@@ -158,7 +159,7 @@ abstract contract Dispatcher is Payments, V2SwapRouter, V3SwapRouter, V3ToV4Migr
                             payerIsUser := calldataload(add(inputs.offset, 0x80))
                         }
                         address[] calldata path = inputs.toAddressArray(3);
-                        address payer = payerIsUser ? lockedBy : address(this);
+                        address payer = payerIsUser ? _msgSender() : address(this);
                         v2SwapExactOutput(map(recipient), amountOut, amountInMax, path, payer);
                     } else if (command == Commands.PERMIT2_PERMIT) {
                         // equivalent: abi.decode(inputs, (IAllowanceTransfer.PermitSingle, bytes))
@@ -167,7 +168,7 @@ abstract contract Dispatcher is Payments, V2SwapRouter, V3SwapRouter, V3ToV4Migr
                             permitSingle := inputs.offset
                         }
                         bytes calldata data = inputs.toBytes(6); // PermitSingle takes first 6 slots (0..5)
-                        PERMIT2.permit(lockedBy, permitSingle, data);
+                        PERMIT2.permit(_msgSender(), permitSingle, data);
                     } else if (command == Commands.WRAP_ETH) {
                         // equivalent: abi.decode(inputs, (address, uint256))
                         address recipient;
@@ -189,7 +190,7 @@ abstract contract Dispatcher is Payments, V2SwapRouter, V3SwapRouter, V3ToV4Migr
                     } else if (command == Commands.PERMIT2_TRANSFER_FROM_BATCH) {
                         (IAllowanceTransfer.AllowanceTransferDetails[] memory batchDetails) =
                             abi.decode(inputs, (IAllowanceTransfer.AllowanceTransferDetails[]));
-                        permit2TransferFrom(batchDetails, lockedBy);
+                        permit2TransferFrom(batchDetails, _msgSender());
                     } else if (command == Commands.BALANCE_CHECK_ERC20) {
                         // equivalent: abi.decode(inputs, (address, address, uint256))
                         address owner;
@@ -236,7 +237,7 @@ abstract contract Dispatcher is Payments, V2SwapRouter, V3SwapRouter, V3ToV4Migr
                     // This can be done in 2 ways:
                     //    1. This contract is permitted for the specific token and the caller is approved for ALL of the owner's tokens
                     //    2. This contract is permitted for ALL of the owner's tokens and the caller is permitted for the specific token
-                    if (!isAuthorizedForToken(lockedBy, tokenId)) {
+                    if (!isAuthorizedForToken(_msgSender(), tokenId)) {
                         revert NotAuthorizedForToken(tokenId);
                     }
 
@@ -259,6 +260,19 @@ abstract contract Dispatcher is Payments, V2SwapRouter, V3SwapRouter, V3ToV4Migr
         }
     }
 
+    /// @notice Calculates the recipient address for a command
+    /// @param recipient The recipient or recipient-flag for the command
+    /// @return output The resultant recipient for the command
+    function map(address recipient) internal view returns (address) {
+        if (recipient == Constants.MSG_SENDER) {
+            return _msgSender();
+        } else if (recipient == Constants.ADDRESS_THIS) {
+            return address(this);
+        } else {
+            return recipient;
+        }
+    }
+
     /// @notice Executes encoded commands along with provided inputs.
     /// @param commands A set of concatenated commands, each 1 byte in length
     /// @param inputs An array of byte strings containing abi encoded inputs for each command

contracts/base/LockAndMsgSender.sol

@@ -1,35 +1,28 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 pragma solidity ^0.8.24;
 
-import {Constants} from '../libraries/Constants.sol';
+import {Locker} from '../libraries/Locker.sol';
 
 contract LockAndMsgSender {
     error ContractLocked();
 
-    address internal constant NOT_LOCKED_FLAG = address(1);
-    address internal lockedBy = NOT_LOCKED_FLAG;
-
+    /// @notice Modifier enforcing a reentrancy lock that allows self-reentrancy
     modifier isNotLocked() {
+        // Apply a reentrancy lock for all external callers
         if (msg.sender != address(this)) {
-            if (lockedBy != NOT_LOCKED_FLAG) revert ContractLocked();
-            lockedBy = msg.sender;
+            if (Locker.isLocked()) revert ContractLocked();
+            Locker.set(msg.sender);
             _;
-            lockedBy = NOT_LOCKED_FLAG;
+            Locker.set(address(0));
         } else {
+            // The contract is allowed to reenter itself, so the lock is not checked
             _;
         }
     }
 
-    /// @notice Calculates the recipient address for a command
-    /// @param recipient The recipient or recipient-flag for the command
-    /// @return output The resultant recipient for the command
-    function map(address recipient) internal view returns (address) {
-        if (recipient == Constants.MSG_SENDER) {
-            return lockedBy;
-        } else if (recipient == Constants.ADDRESS_THIS) {
-            return address(this);
-        } else {
-            return recipient;
-        }
+    /// @notice Function to be used instead of msg.sender, as the contract performs self-reentrancy and at
+    /// times msg.sender == address(this). Instead _msgSender() returns the initiator of the lock
+    function _msgSender() internal view returns (address) {
+        return Locker.get();
     }
 }

contracts/libraries/Locker.sol

@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity ^0.8.24;
+
+/// @notice A library to implement a reentrancy lock in transient storage.
+/// @dev Instead of storing a boolean, the locker's address is stored to allow the contract to know who locked the contract
+/// TODO: This library can be deleted when we have the transient keyword support in solidity.
+library Locker {
+    // The slot holding the locker state, transiently. bytes32(uint256(keccak256("Locker")) - 1)
+    bytes32 constant LOCKER_SLOT = 0x0e87e1788ebd9ed6a7e63c70a374cd3283e41cad601d21fbe27863899ed4a708;
+
+    function set(address locker) internal {
+        // The locker is always msg.sender or address(0) so does not need to be cleaned
+        assembly ("memory-safe") {
+            tstore(LOCKER_SLOT, locker)
+        }
+    }
+
+    function get() internal view returns (address locker) {
+        assembly ("memory-safe") {
+            locker := tload(LOCKER_SLOT)
+        }
+    }
+
+    function isLocked() internal view returns (bool) {
+        return Locker.get() != address(0);
+    }
+}

package.json

@@ -61,12 +61,12 @@
     "typescript": "^3.7.3"
   },
   "scripts": {
-    "compile": "hardhat compile",
-    "test": "hardhat test",
+    "compile": "hardhat compile && forge build",
+    "test:hardhat": "yarn compile && hardhat test",
     "test:gas": "UPDATE_SNAPSHOT=1 yarn test --grep gas",
-    "test:all": "UPDATE_SNAPSHOT=1 yarn test",
+    "test:all": "UPDATE_SNAPSHOT=1 yarn test:hardhat && forge test --isolate",
     "prettier:fix": "prettier --write '**/*.ts' && prettier --write '**/*.json'",
-    "lint:fix": "yarn prettier:fix && forge fmt",
-    "prettier": "prettier --check '**/*.ts' && forge fmt --check"
+    "lint": "yarn prettier:fix && forge fmt",
+    "lint:check": "prettier --check '**/*.ts' && forge fmt --check"
   }
 }

test/foundry-tests/Locker.t.sol

@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.24;
+
+import 'forge-std/Test.sol';
+import {Constants} from '../../contracts/libraries/Constants.sol';
+import {Locker} from '../../contracts/libraries/Locker.sol';
+
+contract LockerTest is Test {
+    function test_fuzz_set_get(address locker1, address locker2, address locker3) public {
+        assertEq(Locker.get(), address(0));
+
+        Locker.set(locker1);
+        assertEq(Locker.get(), locker1);
+
+        Locker.set(locker2);
+        assertEq(Locker.get(), locker2);
+
+        Locker.set(locker3);
+        assertEq(Locker.get(), locker3);
+
+        Locker.set(address(0));
+        assertEq(Locker.get(), address(0));
+    }
+
+    function test_fuzz_isLocked(address locker) public {
+        assertEq(Locker.get(), address(0));
+        assertEq(Locker.isLocked(), false);
+
+        Locker.set(locker);
+        // the contract is locked when the locker is not address(0)
+        assertEq(Locker.isLocked(), locker != address(0));
+
+        Locker.set(address(0));
+        assertEq(Locker.isLocked(), false);
+    }
+
+    function test_lockerSlot() public {
+        bytes32 expectedSlot = bytes32(uint256(keccak256('Locker')) - 1);
+        assertEq(expectedSlot, Locker.LOCKER_SLOT);
+    }
+}

test/integration-tests/gas-tests/__snapshots__/CheckOwnership.gas.test.ts.snap

@@ -3,6 +3,6 @@
 exports[`Check Ownership Gas gas: balance check ERC20 1`] = `
 Object {
   "calldataByteLength": 356,
-  "gasUsed": 39784,
+  "gasUsed": 37644,
 }
 `;

test/integration-tests/gas-tests/__snapshots__/Payments.gas.test.ts.snap

@@ -3,48 +3,48 @@
 exports[`Payments Gas Tests Individual Command Tests gas: SWEEP with ERC20 1`] = `
 Object {
   "calldataByteLength": 356,
-  "gasUsed": 39153,
+  "gasUsed": 37013,
 }
 `;
 
 exports[`Payments Gas Tests Individual Command Tests gas: SWEEP_WITH_FEE 1`] = `
 Object {
   "calldataByteLength": 516,
-  "gasUsed": 67830,
+  "gasUsed": 65690,
 }
 `;
 
 exports[`Payments Gas Tests Individual Command Tests gas: TRANSFER with ERC20 1`] = `
 Object {
   "calldataByteLength": 356,
-  "gasUsed": 38178,
+  "gasUsed": 36038,
 }
 `;
 
 exports[`Payments Gas Tests Individual Command Tests gas: TRANSFER with ETH 1`] = `
 Object {
   "calldataByteLength": 356,
-  "gasUsed": 33749,
+  "gasUsed": 31609,
 }
 `;
 
 exports[`Payments Gas Tests Individual Command Tests gas: UNWRAP_WETH 1`] = `
 Object {
   "calldataByteLength": 324,
-  "gasUsed": 46717,
+  "gasUsed": 44577,
 }
 `;
 
 exports[`Payments Gas Tests Individual Command Tests gas: UNWRAP_WETH_WITH_FEE 1`] = `
 Object {
   "calldataByteLength": 644,
-  "gasUsed": 53045,
+  "gasUsed": 50905,
 }
 `;
 
 exports[`Payments Gas Tests Individual Command Tests gas: WRAP_ETH 1`] = `
 Object {
   "calldataByteLength": 324,
-  "gasUsed": 57354,
+  "gasUsed": 53402,
 }
 `;