Skip to content

Commit

Permalink
✨ added reentrancyguard
Browse files Browse the repository at this point in the history
  • Loading branch information
atarpara committed Jan 5, 2024
1 parent b3d4ad9 commit 268af43
Show file tree
Hide file tree
Showing 3 changed files with 78 additions and 0 deletions.
4 changes: 4 additions & 0 deletions .gas-snapshot
Original file line number Diff line number Diff line change
Expand Up @@ -930,6 +930,10 @@ RedBlackTreeLibTest:testRedBlackTreeRemoveViaPointer() (gas: 58199)
RedBlackTreeLibTest:testRedBlackTreeTreeFullReverts() (gas: 50293)
RedBlackTreeLibTest:testRedBlackTreeTryInsertAndRemove() (gas: 56127)
RedBlackTreeLibTest:test__codesize() (gas: 13237)
ReentrancyGuardTest:testProtectedCall() (gas: 45314)
ReentrancyGuardTest:testProtectedCallRevertsOnReentrancy() (gas: 48600)
ReentrancyGuardTest:testUnprotectedCall() (gas: 25824)
ReentrancyGuardTest:test__codesize() (gas: 2648)
SSTORE2Test:testReadInvalidPointerCustomBoundsReverts() (gas: 3242)
SSTORE2Test:testReadInvalidPointerCustomBoundsReverts(address,uint256,uint256) (runs: 256, μ: 832570, ~: 629512)
SSTORE2Test:testReadInvalidPointerCustomStartBoundReverts() (gas: 3241)
Expand Down
32 changes: 32 additions & 0 deletions src/utils/ReentrancyGuard.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.4;

/// @notice Reentrancy guard mixin.
/// @author Solady (https://github.com/vectorized/solady/blob/main/src/utils/ReentrancyGuard.sol)
abstract contract ReentrancyGuard {
/*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
/* STORAGE */
/*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/

/// @dev Equivalent to: `uint72(bytes9(keccak256("_REENTRANCY_GUARD_SLOT")))`.
/// Large enough to avoid collisions with lower slots,
/// but not too large to prevent bytecode bloat.
uint256 private constant _REENTRANCY_GUARD_SLOT = 0x929eee149b4bd21268;

/*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
/* REENTRANCY GUARD */
/*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/

modifier nonReentrant() virtual {
/// @solidity memory-safe-assembly
assembly {
if eq(sload(_REENTRANCY_GUARD_SLOT), 2) { revert(codesize(), 0x00) }
sstore(_REENTRANCY_GUARD_SLOT, 2)
}
_;
/// @solidity memory-safe-assembly
assembly {
sstore(_REENTRANCY_GUARD_SLOT, 1)
}
}
}
42 changes: 42 additions & 0 deletions test/ReentrancyGuard.t.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
// SPDX-License-Identifier: MIT
pragma solidity ^0.8.21;

import "./utils/SoladyTest.sol";
import {ReentrancyGuard} from "../src/utils/ReentrancyGuard.sol";

contract ReentrancyGuardTest is SoladyTest, ReentrancyGuard {
uint256 enterTimes;

function testUnprotectedCall() public {
this.unprotectedCall(0);
assertEq(enterTimes, (1 + 0));
this.unprotectedCall(2);
assertEq(enterTimes, (1 + 0) + (1 + 2));
}

function testProtectedCall() public {
this.protectedCall(0);
assertEq(enterTimes, 1);
}

function testProtectedCallRevertsOnReentrancy() public {
vm.expectRevert();
this.protectedCall(1);
}

function unprotectedCall(uint256 recurse) public {
unchecked {
enterTimes++;
if (recurse == 0) return;
this.unprotectedCall(recurse - 1);
}
}

function protectedCall(uint256 recurse) public nonReentrant {
unchecked {
enterTimes++;
if (recurse == 0) return;
this.protectedCall(recurse - 1);
}
}
}

0 comments on commit 268af43

Please sign in to comment.