diff --git a/.gas-snapshot b/.gas-snapshot index e7bfe77..e10e473 100644 --- a/.gas-snapshot +++ b/.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) diff --git a/README.md b/README.md index cf51801..32dd05e 100644 --- a/README.md +++ b/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 diff --git a/src/TimeLock.sol b/src/TimeLock.sol index f895cdc..191b20b 100644 --- a/src/TimeLock.sol +++ b/src/TimeLock.sol @@ -1,11 +1,16 @@ // 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 @@ -13,20 +18,16 @@ contract TimeLock { 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); diff --git a/test/Reentrancy.t.sol b/test/Reentrancy.t.sol new file mode 100644 index 0000000..514196e --- /dev/null +++ b/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); + } +} diff --git a/test/Staking.t.sol b/test/Staking.t.sol index 3ddb706..b799846 100644 --- a/test/Staking.t.sol +++ b/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); - } -} diff --git a/test/StakingInvariant.t.sol b/test/StakingInvariant.t.sol new file mode 100644 index 0000000..e14a049 --- /dev/null +++ b/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()); + } +} diff --git a/test/TimeLock.t.sol b/test/TimeLock.t.sol index 3785512..0ec8004 100644 --- a/test/TimeLock.t.sol +++ b/test/TimeLock.t.sol @@ -4,6 +4,7 @@ pragma solidity ^0.8.24; import {Test} from "forge-std/Test.sol"; import {TimeLock} from "../src/TimeLock.sol"; import {Counter} from "../src/Counter.sol"; +import {Ownable} from "../src/Ownable.sol"; contract TimeLockTest is Test { TimeLock public timelock; @@ -11,9 +12,12 @@ contract TimeLockTest is Test { address admin = makeAddr("admin"); address attacker = makeAddr("attacker"); uint256 constant DELAY = 2 days; + uint256 constant GRACE_PERIOD = 14 days; function setUp() public { - timelock = new TimeLock(admin, DELAY); + // Deploy as the admin so Ownable makes it the owner. + vm.prank(admin); + timelock = new TimeLock(DELAY, GRACE_PERIOD); counter = new Counter(); } @@ -51,6 +55,21 @@ contract TimeLockTest is Test { timelock.execute(address(counter), 0, data, salt); } + function test_Execute_RevertWhenStale() public { + bytes memory data = _setNumberCalldata(1); + bytes32 salt = bytes32(uint256(1)); + + vm.prank(admin); + timelock.queue(address(counter), 0, data, salt); + + // Past eta plus the grace period the operation is stale. + vm.warp(block.timestamp + DELAY + GRACE_PERIOD + 1); + + vm.expectRevert(); // TooLate + vm.prank(admin); + timelock.execute(address(counter), 0, data, salt); + } + function test_Execute_RevertWhenNotQueued() public { vm.expectRevert(TimeLock.NotQueued.selector); vm.prank(admin); @@ -87,7 +106,7 @@ contract TimeLockTest is Test { } function test_OnlyAdmin_CanQueue() public { - vm.expectRevert(TimeLock.NotAdmin.selector); + vm.expectRevert(Ownable.NotOwner.selector); vm.prank(attacker); timelock.queue(address(counter), 0, _setNumberCalldata(1), bytes32(0)); } @@ -101,13 +120,13 @@ contract TimeLockTest is Test { vm.warp(block.timestamp + DELAY); - vm.expectRevert(TimeLock.NotAdmin.selector); + vm.expectRevert(Ownable.NotOwner.selector); vm.prank(attacker); timelock.execute(address(counter), 0, data, salt); } - function testFuzz_ExecuteAtAnyTimeAfterEta(uint256 wait) public { - wait = bound(wait, DELAY, DELAY + 365 days); + function testFuzz_ExecuteAnytimeWithinGracePeriod(uint256 wait) public { + wait = bound(wait, DELAY, DELAY + GRACE_PERIOD); bytes memory data = _setNumberCalldata(99); bytes32 salt = bytes32(uint256(1)); @@ -120,4 +139,42 @@ contract TimeLockTest is Test { timelock.execute(address(counter), 0, data, salt); assertEq(counter.number(), 99); } + + function test_TransferAdmin_TwoStep() public { + address newAdmin = makeAddr("newAdmin"); + + vm.prank(admin); + timelock.transferOwnership(newAdmin); + // Nothing changes until the new admin accepts. + assertEq(timelock.owner(), admin); + assertEq(timelock.pendingOwner(), newAdmin); + + vm.prank(newAdmin); + timelock.acceptOwnership(); + assertEq(timelock.owner(), newAdmin); + + // The new admin can queue now, the old one no longer can. + vm.prank(newAdmin); + timelock.queue(address(counter), 0, _setNumberCalldata(1), bytes32(0)); + + vm.expectRevert(Ownable.NotOwner.selector); + vm.prank(admin); + timelock.queue(address(counter), 0, _setNumberCalldata(2), bytes32(uint256(2))); + } + + function test_TransferAdmin_OnlyOwnerCanStart() public { + vm.expectRevert(Ownable.NotOwner.selector); + vm.prank(attacker); + timelock.transferOwnership(attacker); + } + + function test_TransferAdmin_OnlyPendingOwnerCanAccept() public { + address newAdmin = makeAddr("newAdmin"); + vm.prank(admin); + timelock.transferOwnership(newAdmin); + + vm.expectRevert(Ownable.NotPendingOwner.selector); + vm.prank(attacker); + timelock.acceptOwnership(); + } }