use alloy_dyn_abi::{DynSolType, DynSolValue};
use alloy_primitives::{Address, Bytes, U256};
use edb_common::{
types::{CallResult, EdbSolValue, ExecutionFrameId, Trace},
EdbContext,
};
use eyre::Result;
use foundry_compilers::{artifacts::Contract, Artifact};
use revm::{
bytecode::OpCode,
context::{ContextTr, CreateScheme, JournalTr},
database::CacheDB,
interpreter::{
interpreter_types::{InputsTr, Jumps},
CallInputs, CallOutcome, CreateInputs, CreateOutcome, Interpreter,
},
Database, DatabaseCommit, DatabaseRef, Inspector,
};
use serde::{Deserialize, Serialize};
use std::{
collections::HashMap,
ops::{Deref, DerefMut},
sync::Arc,
};
use tracing::{debug, error};
use crate::{
analysis::{dyn_sol_type, AnalysisResult, UserDefinedTypeRef, VariableRef, UVID},
USID,
};
pub const MAGIC_SNAPSHOT_NUMBER: U256 = U256::from_be_bytes([
0x20, 0x15, 0x05, 0x02, 0xff, 0xff, 0xff, 0xff, 0x20, 0x24, 0x01, 0x02, 0xff, 0xff, 0xff, 0xff,
0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
]);
pub const MAGIC_VARIABLE_UPDATE_NUMBER: U256 = U256::from_be_bytes([
0x20, 0x25, 0x02, 0x08, 0xff, 0x20, 0x25, 0x09, 0x16, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
]);
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HookSnapshot<DB>
where
DB: Database + DatabaseCommit + DatabaseRef + Clone,
<CacheDB<DB> as Database>::Error: Clone,
<DB as Database>::Error: Clone,
{
pub target_address: Address,
pub bytecode_address: Address,
pub database: Arc<CacheDB<DB>>,
pub locals: HashMap<String, Option<Arc<EdbSolValue>>>,
pub state_variables: HashMap<String, Option<Arc<EdbSolValue>>>,
pub usid: USID,
}
#[derive(Debug, Clone)]
pub struct HookSnapshots<DB>
where
DB: Database + DatabaseCommit + DatabaseRef + Clone,
<CacheDB<DB> as Database>::Error: Clone,
<DB as Database>::Error: Clone,
{
snapshots: Vec<(ExecutionFrameId, Option<HookSnapshot<DB>>)>,
}
impl<DB> Default for HookSnapshots<DB>
where
DB: Database + DatabaseCommit + DatabaseRef + Clone,
<CacheDB<DB> as Database>::Error: Clone,
<DB as Database>::Error: Clone,
{
fn default() -> Self {
Self { snapshots: Vec::new() }
}
}
impl<DB> Deref for HookSnapshots<DB>
where
DB: Database + DatabaseCommit + DatabaseRef + Clone,
<CacheDB<DB> as Database>::Error: Clone,
<DB as Database>::Error: Clone,
{
type Target = Vec<(ExecutionFrameId, Option<HookSnapshot<DB>>)>;
fn deref(&self) -> &Self::Target {
&self.snapshots
}
}
impl<DB> DerefMut for HookSnapshots<DB>
where
DB: Database + DatabaseCommit + DatabaseRef + Clone,
<CacheDB<DB> as Database>::Error: Clone,
<DB as Database>::Error: Clone,
{
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.snapshots
}
}
impl<DB> IntoIterator for HookSnapshots<DB>
where
DB: Database + DatabaseCommit + DatabaseRef + Clone,
<CacheDB<DB> as Database>::Error: Clone,
<DB as Database>::Error: Clone,
{
type Item = (ExecutionFrameId, Option<HookSnapshot<DB>>);
type IntoIter = std::vec::IntoIter<Self::Item>;
fn into_iter(self) -> Self::IntoIter {
self.snapshots.into_iter()
}
}
impl<DB> HookSnapshots<DB>
where
DB: Database + DatabaseCommit + DatabaseRef + Clone,
<CacheDB<DB> as Database>::Error: Clone,
<DB as Database>::Error: Clone,
{
pub fn get_snapshot(&self, frame_id: ExecutionFrameId) -> Option<&HookSnapshot<DB>> {
self.snapshots
.iter()
.find(|(id, _)| *id == frame_id)
.and_then(|(_, snapshot)| snapshot.as_ref())
}
pub fn get_frames_with_hooks(&self) -> Vec<ExecutionFrameId> {
self.snapshots
.iter()
.filter_map(
|(frame_id, snapshot)| {
if snapshot.is_some() {
Some(*frame_id)
} else {
None
}
},
)
.collect()
}
fn add_frame_placeholder(&mut self, frame_id: ExecutionFrameId) {
self.snapshots.push((frame_id, None));
}
fn update_last_frame_with_snapshot(
&mut self,
frame_id: ExecutionFrameId,
snapshot: HookSnapshot<DB>,
) {
if let Some((last_frame_id, slot)) = self.snapshots.last_mut() {
if last_frame_id != &frame_id {
error!("Mismatched frame IDs: expected {}, got {}", last_frame_id, frame_id);
}
if slot.is_none() {
*slot = Some(snapshot);
return;
}
}
self.snapshots.push((frame_id, Some(snapshot)));
}
}
#[derive(Debug)]
pub struct HookSnapshotInspector<'a, DB>
where
DB: Database + DatabaseCommit + DatabaseRef + Clone,
<CacheDB<DB> as Database>::Error: Clone,
<DB as Database>::Error: Clone,
{
trace: &'a Trace,
analysis: &'a HashMap<Address, AnalysisResult>,
pub snapshots: HookSnapshots<DB>,
frame_stack: Vec<ExecutionFrameId>,
current_trace_id: usize,
creation_hooks: Vec<(Bytes, Bytes, Bytes)>,
uvid_values: HashMap<UVID, Arc<EdbSolValue>>,
}
impl<'a, DB> HookSnapshotInspector<'a, DB>
where
DB: Database + DatabaseCommit + DatabaseRef + Clone,
<CacheDB<DB> as Database>::Error: Clone,
<DB as Database>::Error: Clone,
{
pub fn new(trace: &'a Trace, analysis: &'a HashMap<Address, AnalysisResult>) -> Self {
Self {
trace,
analysis,
snapshots: HookSnapshots::default(),
frame_stack: Vec::new(),
current_trace_id: 0,
creation_hooks: Vec::new(),
uvid_values: HashMap::new(),
}
}
pub fn with_creation_hooks(
&mut self,
hooks: Vec<(&Contract, &Contract, &Bytes)>,
) -> Result<()> {
for (original, hooked, args) in hooks {
self.creation_hooks.push((
original
.get_bytecode_bytes()
.ok_or(eyre::eyre!("Failed to get bytecode for contract"))?
.as_ref()
.clone(),
hooked
.get_bytecode_bytes()
.ok_or(eyre::eyre!("Failed to get bytecode for contract"))?
.as_ref()
.clone(),
args.clone(),
));
}
Ok(())
}
pub fn into_snapshots(self) -> HookSnapshots<DB> {
self.snapshots
}
fn current_frame_id(&self) -> Option<ExecutionFrameId> {
self.frame_stack.last().copied()
}
fn push_frame(&mut self, trace_id: usize) {
let frame_id = ExecutionFrameId::new(trace_id, 0);
self.frame_stack.push(frame_id);
self.snapshots.add_frame_placeholder(frame_id);
}
fn pop_frame(&mut self) -> Option<ExecutionFrameId> {
if let Some(frame_id) = self.frame_stack.pop() {
if let Some(parent_frame_id) = self.frame_stack.last_mut() {
parent_frame_id.increment_re_entry();
}
if let Some(current_frame_id) = self.current_frame_id() {
self.snapshots.add_frame_placeholder(current_frame_id);
}
Some(frame_id)
} else {
None
}
}
fn check_and_record_hook(
&mut self,
data: &[u8],
interp: &Interpreter,
ctx: &mut EdbContext<DB>,
) {
let address = self
.current_frame_id()
.and_then(|frame_id| self.trace.get(frame_id.trace_entry_id()))
.map(|entry| entry.code_address)
.unwrap_or(interp.input.target_address());
let usid_opt = if data.len() >= 32 {
U256::from_be_slice(&data[..32]).try_into().ok()
} else {
error!("KECCAK256 input data too short for snapshot, skipping");
return;
};
let Some(usid) = usid_opt else {
error!("Hook call data does not contain valid USID, skipping snapshot");
return;
};
let mut inner = ctx.journal().to_inner();
let changes = inner.finalize();
let mut snap = ctx.db().clone();
snap.commit(changes);
let Some(step) = self.analysis.get(&address).and_then(|a| a.usid_to_step.get(&usid)) else {
error!(
address=?address,
usid=?usid,
"No analysis step found for address and USID, skipping hook snapshot",
);
return;
};
let mut locals = HashMap::new();
for variable in &step.read().accessible_variables {
if variable.declaration().state_variable {
continue;
}
let uvid = variable.id();
let name = variable.declaration().name.clone();
locals.insert(name, self.uvid_values.get(&uvid).cloned());
}
if let Some(current_frame_id) = self.current_frame_id() {
if let Some(entry) = self.trace.get(current_frame_id.trace_entry_id()) {
let hook_snapshot = HookSnapshot {
target_address: entry.target,
bytecode_address: entry.code_address,
database: Arc::new(snap),
locals,
usid,
state_variables: HashMap::new(), };
self.snapshots.update_last_frame_with_snapshot(current_frame_id, hook_snapshot);
} else {
error!("No trace entry found for frame {}", current_frame_id);
}
} else {
error!("No current frame to update with hook snapshot");
}
}
fn check_and_record_variable_update(
&mut self,
data: &[u8],
interp: &Interpreter,
_ctx: &mut EdbContext<DB>,
) {
let address = self
.current_frame_id()
.and_then(|frame_id| self.trace.get(frame_id.trace_entry_id()))
.map(|entry| entry.code_address)
.unwrap_or(interp.input.target_address());
if data.len() < 96 {
error!(
address=?address,
"KECCAK256 input data too short for variable update value, skipping"
);
return;
}
let Some(uvid) = U256::from_be_slice(&data[..32]).try_into().ok() else {
error!("Hook call data does not contain valid UVID, skipping snapshot");
return;
};
let offset = U256::from_be_slice(&data[32..64]);
if offset != U256::from(0x60) {
error!(
address=?address,
uvid=?uvid,
offset=?offset,
"Unexpected offset for variable update value, skipping"
);
return;
}
let length = U256::from_be_slice(&data[64..96]);
let length_usize = match usize::try_from(length) {
Ok(l) => l,
Err(_) => {
error!(
address=?address,
uvid=?uvid,
length=?length,
"Variable update value length too large, skipping"
);
return;
}
};
let decoded_data = &data[96..96 + length_usize];
let Some(analysis) = self.analysis.get(&address) else {
error!(
address=?address,
uvid=?uvid,
"No analysis found for address, skipping variable update recording",
);
return;
};
let Some(variable) = analysis.uvid_to_variable.get(&uvid) else {
error!(
address=?address,
uvid=?uvid,
"No variable found for address and UVID, skipping variable update recording",
);
return;
};
let value =
match decode_variable_value(&analysis.user_defined_types, variable, decoded_data) {
Ok(v) => v,
Err(e) => {
error!(
address=?address,
uvid=?uvid,
variable=?variable.declaration().type_descriptions.type_string,
type_name = ?variable.declaration().type_name,
data=?hex::encode(decoded_data),
error=?e,
);
return;
}
};
debug!(
uvid=?uvid,
address=?address,
variable=?variable.declaration().name,
value=?value,
"Found variable update",
);
self.uvid_values.insert(uvid, Arc::new(value.into()));
}
fn check_and_apply_creation_hooks(
&mut self,
inputs: &mut CreateInputs,
ctx: &mut EdbContext<DB>,
) {
let Ok(account) = ctx.journaled_state.load_account(inputs.caller) else {
error!("Failed to load account for caller {:?}", inputs.caller);
return;
};
let nonce = account.info.nonce;
let predicted_address = inputs.created_address(nonce);
for (original_bytecode, hooked_bytecode, constructor_args) in &self.creation_hooks {
if inputs.init_code.len() >= constructor_args.len() {
let input_args_start = inputs.init_code.len() - constructor_args.len();
let input_args = &inputs.init_code[input_args_start..];
if input_args == constructor_args.as_ref() {
let input_bytecode = &inputs.init_code[..input_args_start];
if input_bytecode == original_bytecode.as_ref() {
let mut new_init_code = Vec::from(hooked_bytecode.as_ref());
new_init_code.extend_from_slice(constructor_args.as_ref());
inputs.init_code = Bytes::from(new_init_code);
inputs.scheme = CreateScheme::Custom { address: predicted_address };
debug!(
"Replaced creation bytecode with hooked version for {:?} -> {:?}",
inputs.caller, predicted_address
);
break; }
}
}
}
}
pub fn clear(&mut self) {
self.snapshots.snapshots.clear();
self.frame_stack.clear();
self.current_trace_id = 0;
}
}
impl<'a, DB> Inspector<EdbContext<DB>> for HookSnapshotInspector<'a, DB>
where
DB: Database + DatabaseCommit + DatabaseRef + Clone,
<CacheDB<DB> as Database>::Error: Clone,
<DB as Database>::Error: Clone,
{
fn step(&mut self, interp: &mut Interpreter, ctx: &mut EdbContext<DB>) {
let opcode = unsafe { OpCode::new_unchecked(interp.bytecode.opcode()) };
if opcode != OpCode::KECCAK256 {
return;
}
let Some(data) = interp.stack.pop().ok().and_then(|offset_u256| {
let data = interp.stack.pop().ok().and_then(|len_u256| {
let offset = usize::try_from(offset_u256).ok()?;
let len = usize::try_from(len_u256).ok()?;
let data = interp.memory.slice_len(offset, len);
let _ = interp.stack.push(len_u256);
Some(data)
});
let _ = interp.stack.push(offset_u256);
data
}) else {
error!("Failed to read KECCAK256 input data from stack");
return;
};
if data.len() < 32 {
return;
}
let magic_number = U256::from_be_slice(&data[..32]);
if magic_number == MAGIC_SNAPSHOT_NUMBER {
self.check_and_record_hook(&data[32..], interp, ctx);
} else if magic_number == MAGIC_VARIABLE_UPDATE_NUMBER {
self.check_and_record_variable_update(&data[32..], interp, ctx);
}
}
fn call(
&mut self,
_context: &mut EdbContext<DB>,
_inputs: &mut CallInputs,
) -> Option<CallOutcome> {
self.push_frame(self.current_trace_id);
self.current_trace_id += 1;
None
}
fn call_end(
&mut self,
_context: &mut EdbContext<DB>,
inputs: &CallInputs,
outcome: &mut CallOutcome,
) {
let Some(frame_id) = self.pop_frame() else { return };
let Some(entry) = self.trace.get(frame_id.trace_entry_id()) else { return };
if entry.result != Some(outcome.into()) {
error!(
target_address = inputs.target_address.to_string(),
bytecode_address = inputs.bytecode_address.to_string(),
"Call outcome mismatch at frame {}: expected {:?}, got {:?} ({:?})",
frame_id,
entry.result,
Into::<CallResult>::into(&outcome),
outcome,
);
}
}
fn create(
&mut self,
context: &mut EdbContext<DB>,
inputs: &mut CreateInputs,
) -> Option<CreateOutcome> {
self.check_and_apply_creation_hooks(inputs, context);
self.push_frame(self.current_trace_id);
self.current_trace_id += 1;
None
}
fn create_end(
&mut self,
_context: &mut EdbContext<DB>,
_inputs: &CreateInputs,
outcome: &mut CreateOutcome,
) {
let Some(frame_id) = self.pop_frame() else { return };
let Some(entry) = self.trace.get(frame_id.trace_entry_id()) else { return };
if entry.result.as_ref().map(|r| r.result()) != Some(outcome.result.result) {
error!(
"Create outcome mismatch at frame {}: expected {:?}, got {:?}",
frame_id, entry.result, outcome
);
}
}
}
impl<DB> HookSnapshots<DB>
where
DB: Database + DatabaseCommit + DatabaseRef + Clone,
<CacheDB<DB> as Database>::Error: Clone,
<DB as Database>::Error: Clone,
{
pub fn print_summary(&self) {
println!(
"\n\x1b[36m╔══════════════════════════════════════════════════════════════════╗\x1b[0m"
);
println!(
"\x1b[36m║ HOOK SNAPSHOT INSPECTOR SUMMARY ║\x1b[0m"
);
println!(
"\x1b[36m╚══════════════════════════════════════════════════════════════════╝\x1b[0m\n"
);
let total_frames = self.len();
let hook_frames = self.get_frames_with_hooks().len();
println!("\x1b[33m📊 Overall Statistics:\x1b[0m");
println!(" Total frames tracked: \x1b[32m{total_frames}\x1b[0m");
println!(" Frames with hooks: \x1b[32m{hook_frames}\x1b[0m");
println!(
" Hook trigger rate: \x1b[32m{:.1}%\x1b[0m",
if total_frames > 0 { hook_frames as f64 / total_frames as f64 * 100.0 } else { 0.0 }
);
if self.is_empty() {
println!("\n\x1b[90m No execution frames were tracked.\x1b[0m");
return;
}
println!("\n\x1b[33m🎯 Hook Trigger Details:\x1b[0m");
println!(
"\x1b[90m─────────────────────────────────────────────────────────────────\x1b[0m"
);
use std::collections::HashMap;
let mut frame_groups: HashMap<ExecutionFrameId, Vec<&HookSnapshot<DB>>> = HashMap::new();
let mut frame_order = Vec::new();
for (frame_id, snapshot) in &self.snapshots {
if !frame_groups.contains_key(frame_id) {
frame_order.push(*frame_id);
}
match snapshot {
Some(hook_snapshot) => {
frame_groups.entry(*frame_id).or_default().push(hook_snapshot);
}
None => {
frame_groups.entry(*frame_id).or_default();
}
}
}
for (display_idx, frame_id) in frame_order.iter().enumerate() {
let hooks = frame_groups.get(frame_id).unwrap();
if hooks.is_empty() {
println!(
" \x1b[90m[{:3}] Frame {}\x1b[0m (trace.{}, re-entry {}) - No hooks",
display_idx,
frame_id,
frame_id.trace_entry_id(),
frame_id.re_entry_count()
);
} else {
let usids: Vec<_> = hooks.iter().map(|h| h.usid).collect();
let hook_count = hooks.len();
#[allow(deprecated)]
let addresses: std::collections::HashSet<_> =
hooks.iter().map(|h| h.bytecode_address).collect();
println!(
"\n \x1b[32m[{:3}] Frame {}\x1b[0m (trace.{}, re-entry {})",
display_idx,
frame_id,
frame_id.trace_entry_id(),
frame_id.re_entry_count()
);
println!(
" └─ \x1b[33m{} Hook{} Triggered\x1b[0m",
hook_count,
if hook_count == 1 { "" } else { "s" }
);
for address in &addresses {
println!(" ├─ Address: \x1b[36m{address:?}\x1b[0m");
}
if usids.len() == 1 {
println!(" └─ USID: \x1b[36m{}\x1b[0m", usids[0]);
} else if usids.len() <= 10 {
let usid_list: Vec<String> = usids.iter().map(|u| u.to_string()).collect();
println!(" └─ USIDs: \x1b[36m[{}]\x1b[0m", usid_list.join(", "));
} else {
let first_few: Vec<String> =
usids.iter().take(3).map(|u| u.to_string()).collect();
let last_few: Vec<String> =
usids.iter().rev().take(3).rev().map(|u| u.to_string()).collect();
if first_few.last() == last_few.first() {
println!(
" └─ USIDs: \x1b[36m[{} ... {} total]\x1b[0m",
first_few.join(", "),
usids.len()
);
} else {
println!(
" └─ USIDs: \x1b[36m[{}, ... {}, {} total]\x1b[0m",
first_few.join(", "),
last_few.join(", "),
usids.len()
);
}
}
}
}
println!(
"\n\x1b[90m─────────────────────────────────────────────────────────────────\x1b[0m"
);
println!("\x1b[33m💡 Magic Snapshot Number:\x1b[0m {MAGIC_SNAPSHOT_NUMBER:?}");
}
}
pub fn decode_variable_value(
user_defined_types: &HashMap<usize, UserDefinedTypeRef>,
variable: &VariableRef,
data: &[u8],
) -> Result<DynSolValue> {
let type_name = variable
.type_name()
.ok_or(eyre::eyre!("Failed to get variable type: no type name in the declaration"))?;
let Some(variable_type): Option<DynSolType> = dyn_sol_type(user_defined_types, type_name)
else {
return Err(eyre::eyre!("Failed to get variable type: no type string in the declaration"));
};
let value = variable_type
.abi_decode(data)
.map_err(|e| eyre::eyre!("Failed to decode variable value: {}", e))?;
Ok(value)
}