Merge pull request #36 from pentagonxyz/rayquaza7/reentrancy-guard

reentrancy guard

README.md

@@ -57,14 +57,14 @@ utils
 ├─ Data — TODO
 ├─ BitPackLib — "Efficient bit packing library"
 ├─ CustomErrors — "Wrappers for reverting with common error messages"
-├─ ERC1155Receiver — TODO
+├─ ERC1155Receiver — "A minimal interface for receiving ERC1155 tokens"
 ├─ Errors — "Custom error utilities"
 ├─ JumpTableUtil — "Utility macros for retrieving jumpdest pcs from jump tables"
 ├─ LibBit — "A library ported from solady for bit twiddling operations"
 ├─ MerkleProofLib — "Gas optimized merkle proof verification library"
 ├─ Multicallable — "Enables a single call to call multiple methods within a contract"
 ├─ TSOwnable — "An Ownable Implementation using Two-Step Transfer Pattern"
-├─ ReentrancyGuard — TODO
+├─ ReentrancyGuard — "Gas optimized reentrancy protection for smart contracts"
 ├─ SSTORE2 — TODO
 ```
 

foundry.toml

@@ -18,6 +18,7 @@ fs_permissions = [
     { access = "read", path = "./test/utils/mocks/JumpTableUtilWrappers.huff" },
     { access = "read", path = "./test/utils/mocks/MerkleProofLibWrappers.huff" },
     { access = "read", path = "./test/utils/mocks/MulticallableWrappers.huff" },
+    { access = "read", path = "./test/utils/mocks/ReentrancyGuardMock.huff" },
     { access = "read", path = "./test/utils/mocks/LibBitWrapper.huff" },
 ]
 remappings = [

src/utils/CommonErrors.huff

@@ -15,6 +15,12 @@
 #define constant UNSAFE_RECIPIENT_ERROR = 0x554e534146455f524543495049454e5400000000000000000000000000000000
 #define constant UNSAFE_RECIPIENT_LENGTH = 0x10
 
+#define constant REENTRANCY_ERROR = 0xa5245454e5452414e435900000000000000000000000000000000000000000000
+#define constant REENTRANCY_LENGTH = 0x0a
+
+#define constant DISPATCH_ERROR_MESSAGE = 0xe44495350415443485f4552524f52000000000000000000000000000000000000
+#define constant DISPATCH_LENGTH = 0x0e
+
 /// @notice Reverts with an "NOT_AUTHORIZED" message if the condition is false
 #define macro UNAUTHORIZED(condition) = takes (0) returns (0) {
     [NOT_AUTHORIZED_ERROR]      // ["NOT_AUTHORIZED"]
@@ -37,4 +43,20 @@
     [UNSAFE_RECIPIENT_LENGTH]   // [16 (length), "UNSAFE_RECIPIENT"]
     <condition>                 // [condition, 16 (length), "UNSAFE_RECIPIENT"]
     REQUIRE()                   // []
+}
+
+/// @notice Reverts with a "REENTRANCY" message if condition is false
+#define macro REENTRANCY(condition) = takes (0) returns (0) {
+    [REENTRANCY_ERROR]          // ["REENTRANCY"]
+    [REENTRANCY_LENGTH]         // [10 (length), "REENTRANCY"]
+    <condition>                 // [condition, 10 (length), "REENTRANCY"]
+    REQUIRE()                   // []
+}
+
+/// @notice Reverts with a "DISPATCH_ERROR" message if condition is false
+#define macro DISPATCH_ERROR(condition) = takes (0) returns (0) {
+    [DISPATCH_ERROR_MESSAGE]    // ["DISPATCH_ERROR"]
+    [DISPATCH_LENGTH]           // [14 (length), "DISPATCH_ERROR"]
+    <condition>                 // [condition, 14 (length), "DISPATCH_ERROR"]
+    REQUIRE()                   // []
 }

src/utils/ERC1155Receiver.huff

@@ -1,6 +1,6 @@
 /// @title ERC1155 Receiver
 /// @author asnared <https://github.com/abigger87>
-/// @notice A basic interface for receiving ERC1155 tokens.
+/// @notice A minimal interface for receiving ERC1155 tokens.
 
 /// @notice On ERC1155 Received
 /// @notice Returns the function selector for `onERC1155Received`

src/utils/LibBit.huff

@@ -19,7 +19,7 @@
 
     dup1 iszero    // [x == 0, x]
     0x08 shl       // [(x == 0) << 0x08, x]
-    
+
     dup2 [A] lt    // [A < x, r, x]
     0x07 shl       // [(A < x) << 0x07, r, x]
     or             // [r, x]
@@ -38,7 +38,7 @@
     swap1 dup2 shr // [x >> r, r]
     dup1 0x01 shr  // [(x >> r) >> 0x01, x >> r, r]
     or             // [x, r]
-    
+
     dup1 0x02 shr  // [x >> 0x02, x, r]
     or             // [x, r]
 
@@ -112,7 +112,7 @@
 
     0x00 not       // [max, x]
     dup1 dup3 lt   // [is_not_max, max, x]
-    
+
     swap2          // [x, max, is_not_max]
     0x03 dup3 div  // [max / 0x03, x, max, is_not_max]
     dup2 0x01 shr  // [x >> 0x01, max / 0x03, x, max, is_not_max]
@@ -138,7 +138,7 @@
     0x100 xor      // [((max / 0xFF) * x) >> 0xF8) ^ 0x100, is_not_max]
     mul            // [((((max / 0xFF) * x) >> 0xF8) ^ 0x100) * is_not_max]
     0x100 xor      // [(((((max / 0xFF) * x) >> 0xF8) ^ 0x100) * is_not_max) ^ 0x100]
- 
+
     // Return stack:  [c]
 }
 
@@ -158,7 +158,7 @@
 
     0x03 eq succeed jumpi
     0x00 dup1 revert
-    
+
     succeed:
 }
 

src/utils/ReentrancyGuard.huff

@@ -0,0 +1,67 @@
+/// @title Reentrancy Guard
+/// @author rayquaza7
+/// @notice A Reentrancy Guard based off <https://github.com/Rari-Capital/solmate/blob/main/src/utils/ReentrancyGuard.sol>
+/// @license Apache License, Version 2.0
+
+#include "./CommonErrors.huff"
+
+// Constants
+#define constant _UNLOCKED = 0x01
+#define constant _LOCKED = 0x02
+#define constant LOCKED_SLOT = FREE_STORAGE_POINTER()
+
+/// @title Lock
+/// @notice Locks the contract to prevent reentrancy
+#define fn LOCK() = takes (0) returns (0) {
+    [_LOCKED]                // [0x02]
+    dup1                     // [0x02, 0x02]
+    [LOCKED_SLOT]            // [locked_slot, 0x02, 0x02]
+    sload                    // [locked_slot_value, 0x02, 0x02]
+    lt                       // [locked_slot_value < 0x02, 0x02]
+    lock jumpi
+
+    // Otherwise revert with re-entrancy
+    REENTRANCY(0x00)
+
+    lock:
+        [LOCKED_SLOT] sstore
+}
+
+/// @title Unlock
+/// @notice Unlocks the contract
+#define fn UNLOCK() = takes (0) returns (0) {
+    [_UNLOCKED] [LOCKED_SLOT] sstore
+}
+
+// Tests
+
+#define test TEST_LOCK() = takes (0) returns (0) {
+    // Make sure our slot is set to the UNLOCKED state
+    UNLOCK()
+
+    // Lock
+    LOCK()
+    [LOCKED_SLOT] sload
+
+    // We expect the locked slot to be set to 2 - the LOCKED state
+    0x02 eq succeed jumpi
+    0x00 dup1 revert
+
+    succeed:
+}
+
+#define test TEST_UNLOCK() = takes (0) returns (0) {
+    // Make sure our slot is set to the LOCKED state
+    LOCK()
+
+    // Unlock
+    UNLOCK()
+    [LOCKED_SLOT] sload
+
+    // We expect the locked slot to be set to 1 - the UNLOCKED state
+    0x01 eq succeed jumpi
+    0x00 dup1 revert
+
+    succeed:
+}
+

test/utils/ReentrancyGuard.t.sol

@@ -1,2 +1,52 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.15;
+
+import "foundry-huff/HuffDeployer.sol";
+import "forge-std/Test.sol";
+
+// Mock Interface
+interface IGuard {
+    function state() external view returns (uint256);
+    function lock() external;
+    function unlock() external;
+}
+
+contract ReentranctGuardTest is Test {
+    IGuard guard;
+
+    function setUp() public {
+        string memory wrapper_code = vm.readFile("test/utils/mocks/ReentrancyGuardMock.huff");
+        guard = IGuard(HuffDeployer.deploy_with_code("utils/ReentrancyGuard", wrapper_code));
+    }
+
+    /// @notice Test locking
+    function testLocking() public {
+        guard.unlock();
+        uint256 state = guard.state();
+        assertEq(state, 1);
+
+        // We should remain unlocked
+        guard.unlock();
+        state = guard.state();
+        assertEq(state, 1);
+
+        // Let's lock
+        guard.lock();
+        state = guard.state();
+        assertEq(state, 2);
+
+        // We should be able to unlock
+        guard.unlock();
+        state = guard.state();
+        assertEq(state, 1);
+
+        // We cannot lock twice
+        guard.lock();
+        state = guard.state();
+        assertEq(state, 2);
+
+        // vm.expectRevert(bytes4(""));
+        // guard.lock();
+    }
+
+}

test/utils/mocks/ReentrancyGuardMock.huff

@@ -0,0 +1,39 @@
+
+#define function state() view returns (uint256)
+#define function lock() nonpayable returns ()
+#define function unlock() nonpayable returns ()
+
+#define macro LOCK_WRAPPER() = takes (0) returns (0) {
+    LOCK()
+    stop
+}
+
+#define macro UNLOCK_WRAPPER() = takes (0) returns (0) {
+    UNLOCK()
+    stop
+}
+
+#define macro GET_STATE() = takes (0) returns (0) {
+    [LOCKED_SLOT] sload // [LOCKED]
+    0x00 mstore         // []
+    0x20 0x00 return    // []
+}
+
+#define macro MAIN() = takes (0) returns (0) {
+    pc calldataload 0xE0 shr  // [selector]
+
+    dup1 __FUNC_SIG(state) eq state_jump jumpi
+    dup1 __FUNC_SIG(lock) eq lock_jump jumpi
+    dup1 __FUNC_SIG(unlock) eq unlock_jump jumpi
+
+    DISPATCH_ERROR(0x00)
+
+    state_jump:
+        GET_STATE()
+
+    lock_jump:
+        LOCK_WRAPPER()
+
+    unlock_jump:
+        UNLOCK_WRAPPER()
+}