Multiple improvements to hybrid echidna (#87)

* Support loading contracts from echidna state initialization file

* Support ContractCreated init event

* Support STATICCALL, DELEGATECALL, support echidna initialisation function call event

* Support setting gas price

* Fix call return data handling

* Lint

optik/common/world.py

@@ -6,6 +6,9 @@
     ARCH,
     contract,
     Cst,
+    evm_get_static_flag,
+    evm_set_gas_price,
+    evm_set_static_flag,
     EVMContract,
     EVMTransaction,
     increment_block_number,
@@ -116,18 +119,22 @@ def __init__(
     def current_runtime(self) -> EVMRuntime:
         return self.runtime_stack[-1]
 
-    def push_runtime(self, tx: Optional[AbstractTx]) -> EVMRuntime:
+    def push_runtime(
+        self, tx: Optional[AbstractTx], share_storage_uid: Optional[int] = None
+    ) -> EVMRuntime:
         """Send a new transaction to the contract
 
         :param tx: The incoming transaction for which to create a new runtime
         :param is_init_runtime: True if the runtime is created
+        :param share_storage_uid: Optional uid of the runtime with whom the new runtime should share
+        its storage. Used solely for DELEGATECALL
         :return: The new runtime created to execute 'tx'
         """
         # Create a new engine that shares runtime code, symbolic
         # variables, and path constraints
         new_engine = self.root_engine._duplicate(share={"mem", "vars", "path"})
         # Create new maat contract runtime for new engine
-        new_evm_runtime(new_engine, self.root_engine)
+        new_evm_runtime(new_engine, self.root_engine, share_storage_uid)
         self.runtime_stack.append(EVMRuntime(new_engine, tx))
         return self.current_runtime
 
@@ -183,6 +190,7 @@ def __init__(self) -> None:
         self.tx_queue: List[AbstractTx] = []
         self.current_tx: Optional[AbstractTx] = None
         self.monitors: List[WorldMonitor] = []
+        self.static_flag_stack: List[bool] = []
         # Counter for transactions being run
         self._current_tx_num: int = 0
         # Root engine
@@ -296,11 +304,14 @@ def current_engine(self) -> MaatEngine:
         return self.current_contract.current_runtime.engine
 
     def _push_runtime(
-        self, runner: ContractRunner, tx: Optional[AbstractTx]
+        self,
+        runner: ContractRunner,
+        tx: Optional[AbstractTx],
+        share_storage_uid: Optional[int] = None,
     ) -> EVMRuntime:
         """Wrapper function to push a new runtime for a contract runner, and
         trigger the corresponding event"""
-        rt = runner.push_runtime(tx)
+        rt = runner.push_runtime(tx, share_storage_uid)
         self._on_event("new_runtime", rt)
         return rt
 
@@ -322,14 +333,18 @@ def run(self) -> STOP:
                 # If no actual transaction, get next one
                 if self.current_tx.tx is None:
                     continue
+                # Update gas price for this transaction
+                evm_set_gas_price(
+                    self.root_engine, self.current_tx.tx.gas_price
+                )
                 # Find contract runner for the target contract
                 contract_addr = self.current_tx.tx.recipient
                 try:
                     runner = self.contracts[contract_addr]
-                except KeyError:
+                except KeyError as e:
                     raise WorldException(
                         f"Transaction recipient is {contract_addr}, but no contract is deployed there"
-                    )
+                    ) from e
                 # Create new runtime to run this transaction
                 self._push_runtime(runner, self.current_tx)
                 # Add to call stack
@@ -371,14 +386,21 @@ def run(self) -> STOP:
                     rt.revert()
 
                 # If contract was still initializing, handle success/failure
+                create_failed = (
+                    not self.current_contract.initialized
+                ) and not succeeded
                 if not self.current_contract.initialized:
                     # This pushes the success status in the caller contract
                     # stack, and deletes the contract runner for current_contract
                     # if the contract creation failed
                     self._handle_CREATE_after(succeeded)
 
                 # Delete the runtime
-                self.current_contract.pop_runtime()
+                # We do it only if CREATE didn't fail. If it failed, then
+                # _handle_CREATE_after will have deleted the entire contract runner
+                #  already
+                if not create_failed:
+                    self.current_contract.pop_runtime()
 
                 # Handle call result in potential caller contract
                 if is_msg_call_return:
@@ -387,11 +409,12 @@ def run(self) -> STOP:
                             self.call_stack[-2]
                         ].current_runtime.engine
                     )
-                    # Handle return of CALL/DELEGATECALL/CALLCODE
+                    # Handle return of call
                     if caller_contract.outgoing_transaction.type in [
                         TX.CALL,
                         TX.CALLCODE,
                         TX.DELEGATECALL,
+                        TX.STATICCALL,
                     ]:
                         self._handle_CALL_after(caller_contract, succeeded)
 
@@ -408,7 +431,7 @@ def run(self) -> STOP:
                 # Handle message call
                 if out_tx.type in [TX.CREATE, TX.CREATE2]:
                     self._handle_CREATE()
-                elif out_tx.type == TX.CALL:
+                elif out_tx.type in [TX.CALL, TX.STATICCALL, TX.DELEGATECALL]:
                     # Handle ETH transfers to EOAs
                     if not self.is_contract(out_tx.recipient):
                         self._handle_ETH_transfer()
@@ -531,6 +554,9 @@ def _handle_CALL(self) -> None:
 
         rt: EVMRuntime = self.current_contract.current_runtime
         out_tx = contract(rt.engine).outgoing_transaction
+        share_storage_uid = (
+            self.current_engine.uid if out_tx.type == TX.DELEGATECALL else None
+        )
 
         # Get runner for target contract
         try:
@@ -548,8 +574,9 @@ def _handle_CALL(self) -> None:
             self.current_tx.block_timestamp_inc,
             VarContext(),
         )
-        self._push_runtime(contract_runner, tx)
+        self._push_runtime(contract_runner, tx, share_storage_uid)
         self.call_stack.append(contract_runner.address)
+        self.static_flag_stack.append(evm_get_static_flag(self.root_engine))
 
     def _handle_CALL_after(
         self, caller_contract: EVMContract, succeeded: bool
@@ -559,18 +586,42 @@ def _handle_CALL_after(
         # Push the success status of transaction in caller stack
         # if the terminating runtime was called with
         caller_contract.stack.push(Cst(256, 1 if succeeded else 0))
-        # Write the return data in memory
-        if succeeded:
+        # Reset the EVM static flag to previous value
+        evm_set_static_flag(self.root_engine, self.static_flag_stack.pop())
+        # Write the return data in memory (if call ended by RETURN)
+        if succeeded and caller_contract.result_from_last_call:
+            out_size = min(
+                caller_contract.outgoing_transaction.ret_len.as_uint(),
+                caller_contract.result_from_last_call.return_data_size,
+            )
+            # Adjust returned data size to space reserved by caller contract
             if (
-                caller_contract.outgoing_transaction.ret_len.as_uint()
+                out_size
                 < caller_contract.result_from_last_call.return_data_size
             ):
-                raise WorldException(
-                    "Message call returned more bytes than the buffer-size allocated by the caller contract"
-                )
+                return_data = []
+                size = 0  # in bytes
+                for val in caller_contract.result_from_last_call.return_data:
+                    if size == out_size:
+                        break
+                    if size + val.size // 8 > out_size:
+                        return_data.append(
+                            Extract(
+                                val,
+                                val.size - 1,
+                                val.size - 8 * (out_size - size),
+                            )
+                        )
+                        break
+                    else:
+                        return_data.append(val)
+                        size += val.size // 8
+            else:
+                return_data = caller_contract.result_from_last_call.return_data
+            # Write call result data into caller contract's memory
             caller_contract.memory.write_buffer(
                 caller_contract.outgoing_transaction.ret_offset,
-                caller_contract.result_from_last_call.return_data,
+                return_data,
             )
 
     def _update_block_info(self, m: MaatEngine, tx: AbstractTx) -> None:

optik/echidna/__main__.py

@@ -23,6 +23,7 @@
     get_latest_coverage_file,
     extract_contract_bytecode,
     extract_cases_from_json_output,
+    get_echidna_init_file,
 )
 from .runner import replay_inputs, generate_new_inputs, run_echidna_campaign
 from ..common.exceptions import ArgumentParsingError, InitializationError
@@ -83,6 +84,8 @@ def run_hybrid_echidna(arguments: List[str]) -> None:
         logger.error(f"Invalid deployer address: {args.deployer}")
         return
 
+    echidna_init_file = get_echidna_init_file(args)
+
     display.sym_solver_timeout = args.solver_timeout
 
     # Logging stream
@@ -278,7 +281,13 @@ def run_hybrid_echidna(arguments: List[str]) -> None:
         display.fuzz_last_cases_cnt = new_echidna_inputs_cnt
         # Replay new corpus inputs symbolically
         cov.bifurcations = []
-        replay_inputs(new_inputs, contract_file, deployer, cov)
+        replay_inputs(
+            new_inputs,
+            contract_file,
+            deployer,
+            cov,
+            get_echidna_init_file(args),
+        )
 
         # Find inputs to reach new code
         new_inputs_cnt, timeout_cnt = generate_new_inputs(cov, args)

optik/echidna/interface.py

@@ -1,6 +1,8 @@
 import json
 import os
 import tempfile
+import yaml
+import argparse
 from typing import Dict, Final, List, Optional, Tuple, Union
 
 from maat import Cst, EVMTransaction, Value, Var, VarContext
@@ -524,3 +526,18 @@ def count_unique_pc(output: str) -> int:
         unique_pcs = set([x[0] for x in cov])
         res += len(unique_pcs)
     return res
+
+
+def get_echidna_init_file(args: argparse.Namespace) -> Optional[str]:
+    """Return the echidna init file name, or None if no file or config
+    was specified"""
+    if args.config is None:
+        return None
+    with open(args.config, "r") as f:
+        try:
+            config = yaml.safe_load(f)
+        except Exception as e:
+            raise EchidnaException(
+                f"Failed to parse config file {args.config}"
+            ) from e
+        return config.get("initialize", None)

optik/echidna/runner.py

@@ -1,8 +1,9 @@
 import argparse
+import json
 import os
 import subprocess
 from datetime import datetime
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 
 from maat import (
     Cst,
@@ -17,16 +18,16 @@
 from ..common.logger import logger
 from ..common.world import AbstractTx, EVMWorld
 from ..coverage import Coverage
+from ..common.exceptions import EchidnaException, WorldException
 
 
 # TODO(boyan): pass contract bytecode instead of extracting to file
-
-
 def replay_inputs(
     corpus_files: List[str],
     contract_file: str,
     contract_deployer: int,
     cov: Coverage,
+    echidna_init_file: Optional[str],
 ) -> None:
 
     display.reset_current_task()
@@ -47,6 +48,9 @@ def replay_inputs(
         # some of the Coverage classes that rely on on_attach(). We'll probably
         # have to detach() and re-attach() them
         world = EVMWorld()
+        if echidna_init_file:
+            init_world(world, echidna_init_file)
+
         contract_addr = tx_seq[0].tx.recipient
         # Push initial transaction that initialises the target contract
         world.push_transaction(
@@ -77,12 +81,71 @@ def replay_inputs(
         world.push_transactions(tx_seq)
         cov.set_input_uid(corpus_file)
 
-        # Run
-        assert world.run() == STOP.EXIT
+        # Run and ensure the execution terminated properly
+        status = world.run()
+        if status in [STOP.FATAL, STOP.ERROR]:
+            raise WorldException("Engine stopped because of an error")
+        elif status == STOP.HOOK:
+            raise WorldException(
+                "Engine stopped by an event hook before the end of transaction"
+            )
+        elif status == STOP.NONE:
+            raise WorldException(
+                "Engine stopped before the end of transaction for an unknown reason"
+            )
+        elif status != STOP.EXIT:
+            raise WorldException(f"Unexpected engine status: {status}")
 
     return cov
 
 
+def init_world(world: EVMWorld, init_file: str) -> None:
+    """Setup contracts and EOAs in an EVMWorld according to
+    an Echidna state initialisation file"""
+    with open(init_file, "r") as f:
+        data = json.loads(f.read())
+        for event in data:
+            if event["event"] == "ContractCreated":
+                bytecode = bytes.fromhex(event["data"][2:])
+                world.deploy(
+                    "",  # No file, bytecode is in the tx data
+                    int(event["contract_address"], 16),
+                    int(event["from"], 16),
+                    args=[bytecode],
+                    run_init_bytecode=True,
+                )
+            elif event["event"] == "AccountCreated":
+                pass
+            elif event["event"] == "FunctionCall":
+                sender = Cst(160, int(event["from"], 16))
+                data = bytes.fromhex(event["data"][2:])
+                tx = EVMTransaction(
+                    sender,  # origin
+                    sender,  # sender
+                    int(event["to"], 16),  # recipient
+                    Cst(256, event["value"][2:], 16),  # value
+                    [Cst(8, x) for x in data],  # data
+                    Cst(256, int(event["gas_price"], 16)),  # gas price
+                    Cst(256, int(event["gas_used"], 16) * 2),  # gas_limit
+                )
+                world.push_transaction(
+                    AbstractTx(
+                        tx,
+                        Cst(256, 0),
+                        Cst(256, 0),
+                        VarContext(),
+                    )
+                )
+                status = world.run()
+                if status != STOP.EXIT:
+                    raise WorldException(
+                        f"Failed to properly execute initialisation transaction: {event}"
+                    )
+                assert not world.has_pending_transactions
+            else:
+                raise EchidnaException(f"Unsupported event: {event['event']}")
+
+
 def generate_new_inputs(
     cov: Coverage, args: argparse.Namespace, solve_duplicates: bool = False
 ) -> Tuple[int, int]:

pyproject.toml

@@ -14,7 +14,8 @@ dependencies = [
     "rlp",
     "crytic-compile",
     "slither-analyzer",
-    "solc-select"
+    "solc-select",
+    "pyyaml"
 ]
 requires-python = ">=3.7"