ethereum
diff --git a/‎tests/zkevm/test_worst_bytecode.py
Lines changed: 39 additions & 41 deletions b/‎tests/zkevm/test_worst_bytecode.py
Lines changed: 39 additions & 41 deletions
@@ -15,6 +15,7 @@
     Alloc,
     Block,
     BlockchainTestFiller,
+    Bytecode,
     Environment,
     Hash,
     Transaction,
@@ -36,12 +37,12 @@
     "opcode",
     [
         Op.EXTCODESIZE,
-    ],
-)
-@pytest.mark.parametrize(
-    "attack_gas_limit",
-    [
-        Environment().gas_limit,
+        Op.EXTCODEHASH,
+        Op.CALL,
+        Op.CALLCODE,
+        Op.DELEGATECALL,
+        Op.STATICCALL,
+        Op.EXTCODECOPY,
     ],
 )
 @pytest.mark.slow()
@@ -51,7 +52,6 @@ def test_worst_bytecode_single_opcode(
     pre: Alloc,
     fork: Fork,
     opcode: Op,
-    attack_gas_limit: int,
 ):
     """
     Test a block execution where a single opcode execution maxes out the gas limit,
@@ -68,6 +68,7 @@ def test_worst_bytecode_single_opcode(
     # We use 100G gas limit to be able to deploy a large number of contracts in a single block,
     # avoiding bloating the number of preparing blocks in the test.
     env = Environment(gas_limit=100_000_000_000)
+    attack_gas_limit = Environment().gas_limit
 
     # The initcode will take its address as a starting point to the input to the keccak
     # hash function.
@@ -87,6 +88,10 @@ def test_worst_bytecode_single_opcode(
             ),
             condition=Op.LT(Op.MSIZE, MAX_CONTRACT_SIZE),
         )
+        # Despite the whole contract has random bytecode, we make the first opcode be a STOP
+        # so CALL-like attacks return as soon as possible, while EXTCODE(HASH|SIZE) work as
+        # intended.
+        + Op.MSTORE8(0, 0x00)
         + Op.RETURN(0, MAX_CONTRACT_SIZE)
     )
     initcode_address = pre.deploy_contract(code=initcode)
@@ -105,7 +110,7 @@ def test_worst_bytecode_single_opcode(
             Op.CREATE2(
                 value=0,
                 offset=0,
-                size=Op.MSIZE,
+                size=Op.EXTCODESIZE(initcode_address),
                 salt=Op.SLOAD(0),
             ),
         )
@@ -128,40 +133,25 @@ def test_worst_bytecode_single_opcode(
         gas_costs.G_KECCAK_256  # KECCAK static cost
         + math.ceil(85 / 32) * gas_costs.G_KECCAK_256_WORD  # KECCAK dynamic cost for CREATE2
         + gas_costs.G_VERY_LOW * 3  # ~MSTOREs+ADDs
-        + gas_costs.G_COLD_ACCOUNT_ACCESS  # EXTCODESIZE
+        + gas_costs.G_COLD_ACCOUNT_ACCESS  # Opcode cost
         + 30  # ~Gluing opcodes
     )
-    max_number_of_contract_calls = (
+    num_contracts = (
         # Base available gas = GAS_LIMIT - intrinsic - (out of loop MSTOREs)
         attack_gas_limit - intrinsic_gas_cost_calc() - gas_costs.G_VERY_LOW * 4
     ) // loop_cost
 
-    total_contracts_to_deploy = max_number_of_contract_calls
-    approximate_gas_per_deployment = 4_970_000  # Obtained from evm tracing
-    contracts_deployed_per_tx = env.gas_limit // approximate_gas_per_deployment
-
-    deploy_txs = []
-
-    def generate_deploy_tx(contracts_to_deploy: int):
-        return Transaction(
-            to=factory_caller_address,
-            gas_limit=env.gas_limit,
-            gas_price=10**9,  # Bump required due to the amount of full blocks
-            data=Hash(contracts_deployed_per_tx),
-            sender=pre.fund_eoa(),
-        )
-
-    for _ in range(total_contracts_to_deploy // contracts_deployed_per_tx):
-        deploy_txs.append(generate_deploy_tx(contracts_deployed_per_tx))
-
-    if total_contracts_to_deploy % contracts_deployed_per_tx != 0:
-        deploy_txs.append(
-            generate_deploy_tx(total_contracts_to_deploy % contracts_deployed_per_tx)
-        )
+    contracts_deployment_tx = Transaction(
+        to=factory_caller_address,
+        gas_limit=env.gas_limit,
+        gas_price=10**9,
+        data=Hash(num_contracts),
+        sender=pre.fund_eoa(),
+    )
 
     post = {}
     deployed_contract_addresses = []
-    for i in range(total_contracts_to_deploy):
+    for i in range(num_contracts):
         deployed_contract_address = compute_create2_address(
             address=factory_address,
             salt=i,
@@ -170,29 +160,37 @@ def generate_deploy_tx(contracts_to_deploy: int):
         post[deployed_contract_address] = Account(nonce=1)
         deployed_contract_addresses.append(deployed_contract_address)
 
-    opcode_code = (
+    attack_call = Bytecode()
+    if opcode == Op.EXTCODECOPY:
+        attack_call = Op.EXTCODECOPY(address=Op.SHA3(32 - 20 - 1, 85), dest_offset=85, size=1000)
+    else:
+        # For the rest of the opcodes, we can use the same generic attack call
+        # since all only minimally need the `address` of the target.
+        attack_call = Op.POP(opcode(address=Op.SHA3(32 - 20 - 1, 85)))
+    attack_code = (
         # Setup memory for later CREATE2 address generation loop.
+        # 0xFF+[Address(20bytes)]+[seed(32bytes)]+[initcode keccak(32bytes)]
         Op.MSTORE(0, factory_address)
-        + Op.MSTORE8(32 - 20 - 1, 0xFF)  # 0xFF prefix byte
+        + Op.MSTORE8(32 - 20 - 1, 0xFF)
         + Op.MSTORE(32, 0)
         + Op.MSTORE(64, initcode.keccak256())
         # Main loop
         + While(
-            body=Op.POP(opcode(Op.SHA3(32 - 20 - 1, 85))) + Op.MSTORE(32, Op.ADD(Op.MLOAD(32), 1)),
+            body=attack_call + Op.MSTORE(32, Op.ADD(Op.MLOAD(32), 1)),
         )
     )
 
-    if len(opcode_code) > MAX_CONTRACT_SIZE:
+    if len(attack_code) > MAX_CONTRACT_SIZE:
         # TODO: A workaround could be to split the opcode code into multiple contracts
         # and call them in sequence.
         raise ValueError(
-            f"Code size {len(opcode_code)} exceeds maximum code size {MAX_CONTRACT_SIZE}"
+            f"Code size {len(attack_code)} exceeds maximum code size {MAX_CONTRACT_SIZE}"
         )
-    opcode_address = pre.deploy_contract(code=opcode_code)
+    opcode_address = pre.deploy_contract(code=attack_code)
     opcode_tx = Transaction(
         to=opcode_address,
         gas_limit=attack_gas_limit,
-        gas_price=10**9,  # Bump required due to the amount of full blocks
+        gas_price=10**9,
         sender=pre.fund_eoa(),
     )
 
@@ -201,7 +199,7 @@ def generate_deploy_tx(contracts_to_deploy: int):
         pre=pre,
         post=post,
         blocks=[
-            *[Block(txs=[deploy_tx]) for deploy_tx in deploy_txs],
+            Block(txs=[contracts_deployment_tx]),
             Block(txs=[opcode_tx]),
         ],
         exclude_full_post_state_in_output=True,