Puzzle Wallet

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;
pragma experimental ABIEncoderV2;

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

contract PuzzleProxy is UpgradeableProxy {
    address public pendingAdmin;
    address public admin;

    constructor(address _admin, address _implementation, bytes memory _initData) UpgradeableProxy(_implementation, _initData) {
        admin = _admin;
    }

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

    function proposeNewAdmin(address _newAdmin) external {
        pendingAdmin = _newAdmin;
    }

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

    function upgradeTo(address _newImplementation) external onlyAdmin {
        _upgradeTo(_newImplementation);
    }
}

contract PuzzleWallet {
    address public owner;
    uint256 public maxBalance;
    mapping(address => bool) public whitelisted;
    mapping(address => uint256) public balances;

    function init(uint256 _maxBalance) public {
        require(maxBalance == 0, "Already initialized");
        maxBalance = _maxBalance;
        owner = msg.sender;
    }

    modifier onlyWhitelisted {
        require(whitelisted[msg.sender], "Not whitelisted");
        _;
    }

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

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

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

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

    function multicall(bytes[] calldata data) external payable onlyWhitelisted {
        bool depositCalled = false;
        for (uint256 i = 0; i < data.length; i++) {
            bytes memory _data = data[i];
            bytes4 selector;
            assembly {
                selector := mload(add(_data, 32))
            }
            if (selector == this.deposit.selector) {
                require(!depositCalled, "Deposit can only be called once");
                // Protect against reusing msg.value
                depositCalled = true;
            }
            (bool success, ) = address(this).delegatecall(data[i]);
            require(success, "Error while delegating call");
        }
    }
}

Goal

The goal of this ctf is to hijack the PuzzleWallet by becoming admin of the PuzzleProxy.

Exploit

In order to finish this CTF we will have to use another smart contract. We can name this contract Attacker.

This contract utilizes a library to store two different times for two different timezones. The constructor creates two instances of the library for each time to be stored.

PuzzleWallet is an upgradeable contract which is following the Proxy Upgrade Pattern pattern and can be upgraded by calling the upgradeTo function from PuzzleProxy.

When dealing with upgradeable proxies one must be aware that the storage layout of the proxy contract and the implementation (logic) contract must be the same, and this is where we can start looking for possible exploits.

If we compare the storage variables in PuzzleProxy (proxy contract) and PuzzleWallet (implementation contract) we can notice that pendingAdmin variable inside PuzzleProxy corresponds to owner inside PuzzleWallet and admin variable corresponds to maxBalance. This means if we can find a way to update maxBalance variable we can update the admin of PuzzleProxy.

After digging into the PuzzleWallet contract we can notice that the only way to update maxBalance is through the setMaxBalance function, however there are two requirements that we need to satisfy, first one is we need to be whitelisted, second one is the balance of ether in the PuzzleWalet should be equal to 0.

Let's start with passing the first one. In order to get whitelisted we need to call addToWhitelist function inside PuzzleWallet, but there is additional requirement here, the caller of addToWhitelist must be the owner. Remember when we said that owner state variable corresponds to pendingAdmin variable, this means if we find a way to update the pendingAdmin variable, beacause of the nature of Proxy Upgrade Pattern and delegatecall, owner variable will also get updated and we can set the Attacker contract as the owner. To update pendingAdmin, we will call proposeNewAdmin function inside PuzzleProxy, this will update owner inside PuzzleWallet, and therefore we can call addToWhitelist and add the Attacker contract address in the whitelisted mapping.

It means we are now whitelisted and our next task is to find a way how to drain the ETH balance from PuzzleWallet.

The only function that is sending ETH away from PuzleWallet is execute, but in order to call this function successfully we must have first deposited ETH in PuzzleWallet, the contract is tracking this using the balances mapping. Given that this contract already has a balance of 0.001 ETH locked in it (we can check this by loading the PuzzleWallet contract at address inside Remix, we can find this address in the Ethernaut console), we must fund a way to depoist 0.001 ETH but update the balances mapping to show as if we have deposited 0.002 (after depositing 0.001 ETH, PuzzleWallet contract balance will become 0.002 ETH, and we must take all of this ETH from the contract).

Something that we must have in mind is that when using delegatecall, msg.value and msg.sender are context preserving, meaning they will have the same value as in the caller contract.

In order to do so, we will make use of the deposit & multicall functions. First we can see that inside multicall we can pass a bytes array data which is later used when delegatecall low level call is performed inside multicall. When performing a delegatecall we can notice that the PuzzleWallet contract is actually calling "iteslf":

(bool success, ) = address(this).delegatecall(data[i]);

What we can do is prepare the data argument, in a way that when we call multicall, the first parameter from the data array will cause a delegatecall to the deposit function of the same (PuzzleWallet) contract , and the second paramter from the data array will cause a delegatecall to the multicall function which will then call the deposit function inside PuzzleWallet (we must delegatecall to multicall frist and to deposit directly because of the following check: if (selector == this.deposit.selector)).

After passing the two requirements at the end we are only left to call setMaxBalance function from PuzzleWallet and send the uint256 representation of our EOA address uint256(uint160(someAddress)) as a function parameter.

To exploit the contract we need to:

1. create IPuzzleWallet interface and define all the functions that we are going to use from PuzzleProxy & PuzzleWallet
2. create Attacker contract, passing the IPuzzleWallet interface object as a constructor parameter
3. inside the constructor of the Attacker contract we will:
   
   • call proposeNewAdmin() function from PuzzleProxy passing the address of the Attacker contract as an argument for the _newAdmin parameter
   • call addToWhitelist() function from PuzzleWallet passing the address of the Attacker contract as an argument for the addr parameter
   • prepare the data array that we will need to pass to the multicall function
   • call multicall() function from PuzzleWallet passing the previously prepared data array as an argument for the data parameter and send 0.001 ETH as a value
   • call execute() function from PuzzleWallet passing the address of the Attacker contract as an argument for the to parameter, 0.002 ETH as argument for the value parameter and empty string as argument for the data parameter
   • call setMaxBalance() function from PuzzleWallet passing the uint256 representation of our EOA address
   • check if the admin of PuzzleProxy contract is set to our account address

Attacker contract code

interface IPuzzleWallet {
    function proposeNewAdmin(address _newAdmin) external;
    function addToWhitelist(address addr) external;
    function deposit() external payable;
    function multicall(bytes[] calldata data) external payable;
    function execute(address to, uint256 value, bytes calldata data) external payable;
    function setMaxBalance(uint256 _maxBalance) external;
    function admin() external view returns (address);
}

contract Attacker {
    constructor(IPuzzleWallet puzzleWallet) payable {
        puzzleWallet.proposeNewAdmin(address(this));
        puzzleWallet.addToWhitelist(address(this));

        bytes[] memory depositData = new bytes[](1);
        depositData[0] = abi.encodeWithSelector(puzzleWallet.deposit.selector);

        bytes[] memory data = new bytes[](2);
        data[0] = depositData[0];
        data[1] = abi.encodeWithSelector(puzzleWallet.multicall.selector, depositData);

        puzzleWallet.multicall{value: 0.001 ether}(data);
        puzzleWallet.execute(address(this), 0.002 ether, "");
        puzzleWallet.setMaxBalance(uint256(uint160(msg.sender)));

        require(puzzleWallet.admin() == msg.sender, "attack failed");
    }
}