Skip to content

Commit

Permalink
Multiple improvements to hybrid echidna (#87)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
Boyan-MILANOV authored Sep 21, 2022
1 parent 5eff7d3 commit 05e9fd8
Show file tree
Hide file tree
Showing 5 changed files with 165 additions and 24 deletions.
85 changes: 68 additions & 17 deletions optik/common/world.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@
ARCH,
contract,
Cst,
evm_get_static_flag,
evm_set_gas_price,
evm_set_static_flag,
EVMContract,
EVMTransaction,
increment_block_number,
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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)

Expand All @@ -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()
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand All @@ -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:
Expand Down
11 changes: 10 additions & 1 deletion optik/echidna/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
17 changes: 17 additions & 0 deletions optik/echidna/interface.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)
73 changes: 68 additions & 5 deletions optik/echidna/runner.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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()
Expand All @@ -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(
Expand Down Expand Up @@ -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]:
Expand Down
3 changes: 2 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,8 @@ dependencies = [
"rlp",
"crytic-compile",
"slither-analyzer",
"solc-select"
"solc-select",
"pyyaml"
]
requires-python = ">=3.7"

Expand Down

0 comments on commit 05e9fd8

Please sign in to comment.