Skip to content

Commit

Permalink
✨ Transient storage
Browse files Browse the repository at this point in the history
  • Loading branch information
transmissions11 committed Feb 9, 2024
1 parent 3c32b3c commit e1639bf
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 0 deletions.
52 changes: 52 additions & 0 deletions src/test/TransientReentrancyGuard.t.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
// SPDX-License-Identifier: MIT
pragma solidity 0.8.24;

import {DSTestPlus} from "./utils/DSTestPlus.sol";

import {TransientReentrancyGuard} from "../utils/TransientReentrancyGuard.sol";

contract RiskyContract is TransientReentrancyGuard {
uint256 public enterTimes;

function unprotectedCall() public {
enterTimes++;

if (enterTimes > 1) return;

this.protectedCall();
}

function protectedCall() public nonReentrant {
enterTimes++;

if (enterTimes > 1) return;

this.protectedCall();
}

function overprotectedCall() public nonReentrant {}
}

contract TransientReentrancyGuardTest is DSTestPlus {
RiskyContract riskyContract;

function setUp() public {
riskyContract = new RiskyContract();
}

function testFailUnprotectedCall() public {
riskyContract.unprotectedCall();

assertEq(riskyContract.enterTimes(), 1);
}

function testProtectedCall() public {
try riskyContract.protectedCall() {
fail("Reentrancy Guard Failed To Stop Attacker");
} catch {}
}

function testNoReentrancy() public {
riskyContract.overprotectedCall();
}
}
31 changes: 31 additions & 0 deletions src/utils/TransientReentrancyGuard.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
// SPDX-License-Identifier: AGPL-3.0-only
pragma solidity >=0.8.0;

/// @notice Gas optimized reentrancy protection for smart contracts. Leverages Cancuntransient storage.
/// @author Solmate (https://github.com/transmissions11/solmate/blob/main/src/utils/TransientReentrancyGuard.sol)
/// @author Modified from Soledge (https://github.com/Vectorized/soledge/blob/main/src/utils/ReentrancyGuard.sol)
abstract contract TransientReentrancyGuard {
/// Warning: Be careful to avoid collisions with this hand picked slot!
uint256 private constant REENTRANCY_GUARD_SLOT = 0x1FACE81BADDEADBEEF;


modifier nonReentrant() virtual {
bool noReentrancy;

/// @solidity memory-safe-assembly
assembly {
noReentrancy := iszero(tload(REENTRANCY_GUARD_SLOT))

tstore(REENTRANCY_GUARD_SLOT, address())
}

require(noReentrancy, "REENTRANCY");

_;

/// @solidity memory-safe-assembly
assembly {
tstore(REENTRANCY_GUARD_SLOT, 0)
}
}
}

0 comments on commit e1639bf

Please sign in to comment.