Resolve open issues: reentrancy tests, TimeLock Ownable + grace period, Staking invariant

.gas-snapshot

@@ -20,21 +20,29 @@ OwnableTest:test_Renounce_ZerosOwner() (gas: 14479)
 OwnableTest:test_Transfer_OverwritesPendingOwner() (gas: 44002)
 OwnableTest:test_Transfer_RevertWhenNotOwner() (gas: 15260)
 OwnableTest:test_TwoStepTransfer() (gas: 34792)
-StakingInvariantTest:invariant_totalStakedEqualsBalance() (runs: 256, calls: 3840, reverts: 0)
+ReentrancyTest:test_Staking_ReentrancyIsBlocked() (gas: 240602)
+ReentrancyTest:test_Vault_ReentrancyIsBlocked() (gas: 233451)
+StakingInvariantTest:invariant_contractBalanceEqualsTotalStaked() (runs: 256, calls: 3840, reverts: 0)
+StakingInvariantTest:invariant_netFlowEqualsTotalStaked() (runs: 256, calls: 3840, reverts: 0)
+StakingInvariantTest:invariant_totalStakedEqualsActorSum() (runs: 256, calls: 3840, reverts: 0)
 StakingTest:testFuzz_StakeUnstake_TotalAlwaysConsistent(uint96,uint96) (runs: 1000, μ: 113808, ~: 114389)
 StakingTest:test_Stake() (gas: 65043)
 StakingTest:test_Unstake() (gas: 75504)
 StakingTest:test_Unstake_RevertInsufficientStake() (gas: 64892)
-TimeLockTest:testFuzz_ExecuteAtAnyTimeAfterEta(uint256) (runs: 1000, μ: 56458, ~: 56505)
-TimeLockTest:test_Cancel_ClearsQueue() (gas: 36989)
-TimeLockTest:test_Execute_RevertWhenNotQueued() (gas: 20812)
-TimeLockTest:test_Execute_RevertWhenTooEarly() (gas: 47721)
-TimeLockTest:test_OnlyAdmin_CanExecute() (gas: 48914)
-TimeLockTest:test_OnlyAdmin_CanQueue() (gas: 16143)
-TimeLockTest:test_QueueThenExecute() (gas: 55711)
-TimeLockTest:test_Queue_RevertOnDuplicate() (gas: 46068)
+TimeLockTest:testFuzz_ExecuteAnytimeWithinGracePeriod(uint256) (runs: 1000, μ: 56694, ~: 56737)
+TimeLockTest:test_Cancel_ClearsQueue() (gas: 37152)
+TimeLockTest:test_Execute_RevertWhenNotQueued() (gas: 20870)
+TimeLockTest:test_Execute_RevertWhenStale() (gas: 48160)
+TimeLockTest:test_Execute_RevertWhenTooEarly() (gas: 47967)
+TimeLockTest:test_OnlyAdmin_CanExecute() (gas: 49094)
+TimeLockTest:test_OnlyAdmin_CanQueue() (gas: 16223)
+TimeLockTest:test_QueueThenExecute() (gas: 55925)
+TimeLockTest:test_Queue_RevertOnDuplicate() (gas: 46182)
+TimeLockTest:test_TransferAdmin_OnlyOwnerCanStart() (gas: 13305)
+TimeLockTest:test_TransferAdmin_OnlyPendingOwnerCanAccept() (gas: 42016)
+TimeLockTest:test_TransferAdmin_TwoStep() (gas: 61300)
 TokenTest:testFuzz_Approve(address,uint256) (runs: 1000, μ: 36452, ~: 36831)
-TokenTest:testFuzz_Transfer_PreservesTotalSupply(uint256) (runs: 1000, μ: 47249, ~: 47882)
+TokenTest:testFuzz_Transfer_PreservesTotalSupply(uint256) (runs: 1000, μ: 47254, ~: 47882)
 TokenTest:test_Approve_AndTransferFrom() (gas: 73017)
 TokenTest:test_InfiniteAllowance_NotDecremented() (gas: 70002)
 TokenTest:test_Metadata() (gas: 21459)

README.md

@@ -43,12 +43,12 @@ forge snapshot --check
 | `src/Counter.sol` | Trivial counter | unit, fuzz |
 | `src/Token.sol` | Minimal ERC20-style token | unit, fuzz |
 | `src/NFT.sol` | Minimal ERC721-style NFT | unit, fuzz |
-| `src/Vault.sol` | ETH vault | unit, fuzz, invariant |
-| `src/Staking.sol` | ETH staking | unit, fuzz, invariant |
+| `src/Vault.sol` | ETH vault | unit, fuzz, invariant, reentrancy |
+| `src/Staking.sol` | ETH staking | unit, fuzz, invariant, reentrancy |
 | `src/TimeLock.sol` | Queue/execute/cancel timelock | unit, fuzz |
 | `src/Ownable.sol` | Two-step ownership mixin | unit |
 
-Tests live in `test/`. The invariant handler in `test/VaultInvariant.t.sol` constrains the random call surface (bounded amounts, a fixed actor set) so invariants converge instead of bouncing off reverts, and it tracks ghost variables to cross-check the contract's own accounting.
+Tests live in `test/`. The invariant handlers in `test/VaultInvariant.t.sol` and `test/StakingInvariant.t.sol` constrain the random call surface (bounded amounts, a fixed actor set) so invariants converge instead of bouncing off reverts, and they track ghost variables to cross-check each contract's own accounting.
 
 ## Why Foundry over Hardhat
 

src/TimeLock.sol

@@ -1,32 +1,33 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.24;
 
-/// Minimal timelock: queue a call, wait `delay`, then execute. Single admin.
+import {Ownable} from "./Ownable.sol";
+
+/// Minimal timelock: queue a call, wait `delay`, then execute. Admin rights come
+/// from the two-step Ownable mixin, so the admin can be handed over safely.
 /// Used to practice testing block.timestamp manipulation with vm.warp.
-contract TimeLock {
-    address public admin;
+contract TimeLock is Ownable {
     uint256 public immutable delay;
+    // Once eta passes there is a window of gracePeriod to execute in. After that
+    // the queued operation is stale and has to be queued again.
+    uint256 public immutable gracePeriod;
 
     mapping(bytes32 => uint256) public queuedAt; // 0 means not queued
 
     event Queued(bytes32 indexed id, address target, uint256 value, bytes data, uint256 eta);
     event Executed(bytes32 indexed id, address target, uint256 value, bytes data);
     event Cancelled(bytes32 indexed id);
 
-    error NotAdmin();
     error AlreadyQueued();
     error NotQueued();
     error TooEarly(uint256 eta, uint256 now_);
+    error TooLate(uint256 deadline, uint256 now_);
     error CallFailed(bytes returndata);
 
-    modifier onlyAdmin() {
-        if (msg.sender != admin) revert NotAdmin();
-        _;
-    }
-
-    constructor(address admin_, uint256 delay_) {
-        admin = admin_;
+    // Ownable's constructor sets the deployer as the owner (the admin here).
+    constructor(uint256 delay_, uint256 gracePeriod_) {
         delay = delay_;
+        gracePeriod = gracePeriod_;
     }
 
     function hashOp(address target, uint256 value, bytes calldata data, bytes32 salt) public pure returns (bytes32) {
@@ -35,7 +36,7 @@ contract TimeLock {
 
     function queue(address target, uint256 value, bytes calldata data, bytes32 salt)
         external
-        onlyAdmin
+        onlyOwner
         returns (bytes32 id)
     {
         id = hashOp(target, value, data, salt);
@@ -45,7 +46,7 @@ contract TimeLock {
         emit Queued(id, target, value, data, eta);
     }
 
-    function cancel(bytes32 id) external onlyAdmin {
+    function cancel(bytes32 id) external onlyOwner {
         if (queuedAt[id] == 0) revert NotQueued();
         delete queuedAt[id];
         emit Cancelled(id);
@@ -54,13 +55,15 @@ contract TimeLock {
     function execute(address target, uint256 value, bytes calldata data, bytes32 salt)
         external
         payable
-        onlyAdmin
+        onlyOwner
         returns (bytes memory)
     {
         bytes32 id = hashOp(target, value, data, salt);
         uint256 eta = queuedAt[id];
         if (eta == 0) revert NotQueued();
         if (block.timestamp < eta) revert TooEarly(eta, block.timestamp);
+        uint256 deadline = eta + gracePeriod;
+        if (block.timestamp > deadline) revert TooLate(deadline, block.timestamp);
 
         delete queuedAt[id];
         (bool ok, bytes memory ret) = target.call{value: value}(data);

test/Reentrancy.t.sol

@@ -0,0 +1,97 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import {Test} from "forge-std/Test.sol";
+import {Vault} from "../src/Vault.sol";
+import {Staking} from "../src/Staking.sol";
+
+// A receiver that tries to re-enter Vault.withdraw from inside its receive().
+contract VaultReentrant {
+    Vault public vault;
+    bool public reentered;
+
+    constructor(Vault vault_) {
+        vault = vault_;
+    }
+
+    function attack(uint256 amount) external {
+        vault.deposit{value: amount}();
+        vault.withdraw(amount);
+    }
+
+    receive() external payable {
+        // Re-enter once. Vault zeroes our balance before it sends the ETH, so
+        // this second withdraw reverts and takes the whole attack down with it.
+        if (!reentered) {
+            reentered = true;
+            vault.withdraw(msg.value);
+        }
+    }
+}
+
+// The same idea aimed at Staking.unstake.
+contract StakingReentrant {
+    Staking public staking;
+    bool public reentered;
+
+    constructor(Staking staking_) {
+        staking = staking_;
+    }
+
+    function attack(uint256 amount) external {
+        staking.stake{value: amount}();
+        staking.unstake(amount);
+    }
+
+    receive() external payable {
+        if (!reentered) {
+            reentered = true;
+            staking.unstake(msg.value);
+        }
+    }
+}
+
+contract ReentrancyTest is Test {
+    Vault vault;
+    Staking staking;
+    address honest = makeAddr("honest");
+
+    function setUp() public {
+        vault = new Vault();
+        staking = new Staking();
+
+        // A legitimate user whose funds an attacker might try to reach.
+        vm.deal(honest, 20 ether);
+        vm.startPrank(honest);
+        vault.deposit{value: 10 ether}();
+        staking.stake{value: 10 ether}();
+        vm.stopPrank();
+    }
+
+    function test_Vault_ReentrancyIsBlocked() public {
+        VaultReentrant attacker = new VaultReentrant(vault);
+        vm.deal(address(attacker), 1 ether);
+
+        // The re-entrant withdraw reverts, which unwinds the whole attack.
+        vm.expectRevert();
+        attacker.attack(1 ether);
+
+        // Nothing was stolen: the honest deposit and the vault total are intact.
+        assertEq(vault.balances(honest), 10 ether);
+        assertEq(vault.balances(address(attacker)), 0);
+        assertEq(address(vault).balance, 10 ether);
+    }
+
+    function test_Staking_ReentrancyIsBlocked() public {
+        StakingReentrant attacker = new StakingReentrant(staking);
+        vm.deal(address(attacker), 1 ether);
+
+        vm.expectRevert();
+        attacker.attack(1 ether);
+
+        assertEq(staking.stakes(honest), 10 ether);
+        assertEq(staking.stakes(address(attacker)), 0);
+        assertEq(staking.totalStaked(), 10 ether);
+        assertEq(address(staking).balance, 10 ether);
+    }
+}

test/Staking.t.sol

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.24;
 
-import {Test, console} from "forge-std/Test.sol";
+import {Test} from "forge-std/Test.sol";
 import {Staking} from "../src/Staking.sol";
 
 contract StakingTest is Test {
@@ -61,51 +61,3 @@ contract StakingTest is Test {
         assertEq(staking.totalStaked(), staking.stakes(alice) + staking.stakes(bob));
     }
 }
-
-/// Invariant test handler - foundry calls random sequences of these
-contract StakingInvariantHandler is Test {
-    Staking public staking;
-    address[] public actors;
-
-    constructor(Staking _staking) {
-        staking = _staking;
-        actors.push(makeAddr("actor1"));
-        actors.push(makeAddr("actor2"));
-        actors.push(makeAddr("actor3"));
-        for (uint256 i; i < actors.length; i++) {
-            vm.deal(actors[i], 100 ether);
-        }
-    }
-
-    function stake(uint256 actorSeed, uint256 amount) external {
-        address actor = actors[actorSeed % actors.length];
-        amount = bound(amount, 1, 5 ether);
-        vm.prank(actor);
-        staking.stake{value: amount}();
-    }
-
-    function unstake(uint256 actorSeed, uint256 amount) external {
-        address actor = actors[actorSeed % actors.length];
-        uint256 current = staking.stakes(actor);
-        if (current == 0) return;
-        amount = bound(amount, 1, current);
-        vm.prank(actor);
-        staking.unstake(amount);
-    }
-}
-
-contract StakingInvariantTest is Test {
-    Staking public staking;
-    StakingInvariantHandler public handler;
-
-    function setUp() public {
-        staking = new Staking();
-        handler = new StakingInvariantHandler(staking);
-        targetContract(address(handler));
-    }
-
-    /// totalStaked must always equal the contract's ETH balance
-    function invariant_totalStakedEqualsBalance() public view {
-        assertEq(staking.totalStaked(), address(staking).balance);
-    }
-}

test/StakingInvariant.t.sol

@@ -0,0 +1,77 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import {Test} from "forge-std/Test.sol";
+import {Staking} from "../src/Staking.sol";
+
+/// Stateful invariant handler - foundry picks random sequences of these calls
+/// and runs the invariants after each. Mirrors test/VaultInvariant.t.sol.
+contract StakingInvariantHandler is Test {
+    Staking public staking;
+    address[] public actors;
+
+    // Tracked alongside the contract so the invariants can compare against an
+    // independently-maintained sum.
+    uint256 public ghostTotalStaked;
+    uint256 public ghostTotalUnstaked;
+
+    constructor(Staking _staking) {
+        staking = _staking;
+        actors.push(makeAddr("s_actor1"));
+        actors.push(makeAddr("s_actor2"));
+        actors.push(makeAddr("s_actor3"));
+        for (uint256 i; i < actors.length; i++) {
+            vm.deal(actors[i], 100 ether);
+        }
+    }
+
+    function stake(uint256 actorSeed, uint256 amount) external {
+        address actor = actors[actorSeed % actors.length];
+        amount = bound(amount, 1, 5 ether);
+        vm.prank(actor);
+        staking.stake{value: amount}();
+        ghostTotalStaked += amount;
+    }
+
+    function unstake(uint256 actorSeed, uint256 amount) external {
+        address actor = actors[actorSeed % actors.length];
+        uint256 current = staking.stakes(actor);
+        if (current == 0) return;
+        amount = bound(amount, 1, current);
+        vm.prank(actor);
+        staking.unstake(amount);
+        ghostTotalUnstaked += amount;
+    }
+
+    function sumActorStakes() external view returns (uint256 sum) {
+        for (uint256 i; i < actors.length; i++) {
+            sum += staking.stakes(actors[i]);
+        }
+    }
+}
+
+contract StakingInvariantTest is Test {
+    Staking public staking;
+    StakingInvariantHandler public handler;
+
+    function setUp() public {
+        staking = new Staking();
+        handler = new StakingInvariantHandler(staking);
+        targetContract(address(handler));
+    }
+
+    /// totalStaked must always equal the sum of the recorded per-actor stakes.
+    function invariant_totalStakedEqualsActorSum() public view {
+        assertEq(staking.totalStaked(), handler.sumActorStakes());
+    }
+
+    /// Ghost accounting: net flow (staked minus unstaked) must equal totalStaked.
+    function invariant_netFlowEqualsTotalStaked() public view {
+        assertEq(staking.totalStaked(), handler.ghostTotalStaked() - handler.ghostTotalUnstaked());
+    }
+
+    /// The contract's ETH balance must match its own accounting at all times.
+    function invariant_contractBalanceEqualsTotalStaked() public view {
+        assertEq(address(staking).balance, staking.totalStaked());
+    }
+}