diff --git a/Cargo.toml b/Cargo.toml index 073c146..d439e6d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,15 +25,15 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"]} anyhow = "1.0" clap = { version = "4", features = ["derive", "env", "unicode", "wrap_help"] } num-traits = "0.2" -ruint = { version = "1.11", features = ["num-traits"] } +ruint = { version = "1.11.1", features = ["num-traits"] } thiserror = "1.0" twoway = "0.2" ## alloy -alloy-primitives = "0.8" +alloy-primitives = { version = "1.5.2", features = ["serde"] } ## revm -revm = { version = "14"} +revm = { version = "33.1.0", default-features = false, features = ["optional_block_gas_limit", "std", "secp256k1"] } # revm-interpreter = { version = "5.0", features = ["serde"] } bytes = "1.4" diff --git a/src/detect.rs b/src/detect.rs index 41ff61f..0870fd9 100644 --- a/src/detect.rs +++ b/src/detect.rs @@ -3,7 +3,7 @@ use crate::consts::{EIP_1967_DEFAULT_STORAGE, DIAMOND_STANDARD_STORAGE_SLOT_LESS // use hardfork::Hardfork; use crate::proxy_inspector::{ProxyInspector, ProxyDetectDB, InspectorData}; use once_cell::sync::Lazy; -use revm::{inspector_handle_register, primitives::{TransactTo, TxEnv}, EvmBuilder}; +use revm::{Context, context::TxEnv, primitives::TxKind, MainContext, MainBuilder, InspectEvm}; use alloy_primitives::{Address, Bytes, U256}; use tracing::debug; use twoway::find_bytes; @@ -142,28 +142,23 @@ impl StorageCallTaint { let inspector = ProxyInspector::new(); - let mut evm = EvmBuilder::default() + // Build EVM with the new revm 33 API + let tx = TxEnv::builder() + .kind(TxKind::Call(self.address)) + .data(calldata) + .value(U256::ZERO) + .gas_limit(30_000_000) + .build() + .expect("Failed to build TxEnv"); + + let mut evm = Context::mainnet() .with_db(db) - .with_external_context(inspector) - .append_handler_register(inspector_handle_register) - .modify_tx_env(|tx: &mut TxEnv| { - tx.transact_to = TransactTo::Call(self.address); - tx.data = calldata; - tx.value = U256::ZERO; - // Block gas limit is 30M - tx.gas_limit = 30_000_000; - }) - .build(); - - let _res = evm.transact(); - // if let Ok(ok_res) = res { - // println!("success"); - // } else { - // println!("fail"); - // } - // println!("res: {:?}", res); - // let db = evm.db.unwrap(); - evm.context.external.collect() + .build_mainnet_with_inspector(inspector); + + let _res = evm.inspect_one_tx(tx); + + // Get the inspector from the EVM and collect results + evm.inspector.collect() } fn identify_proxy_by_storage(storage: &U256) -> ProxyType { diff --git a/src/proxy_inspector.rs b/src/proxy_inspector.rs index 37155db..a54f160 100644 --- a/src/proxy_inspector.rs +++ b/src/proxy_inspector.rs @@ -1,8 +1,17 @@ -use std::{collections::HashMap, ops::{BitAnd, BitXor}}; +use std::collections::HashMap; use once_cell::sync::Lazy; use revm::{ - interpreter::{opcode, CallInputs, CallOutcome, CallScheme, Gas, InstructionResult, Interpreter, InterpreterResult, OpCode}, primitives::{AccountInfo, Bytecode}, Database, EvmContext, Inspector + interpreter::{ + CallInputs, CallOutcome, CallScheme, Gas, InstructionResult, + Interpreter, InterpreterResult, InterpreterTypes, + interpreter_types::{Jumps, StackTr}, + }, + state::{AccountInfo, Bytecode}, + database_interface::DBErrorMarker, + Database, Inspector, Context, + bytecode::opcode, + context::JournalTr, }; use alloy_primitives::{ @@ -83,14 +92,17 @@ static ADDR_XOR: Lazy = Lazy::new(|| U256::from_be_bytes(hex_literal::hex! #[derive(Clone, Debug, Error)] pub enum ProxyDetectError { - + #[error("Custom error: {0}")] + Custom(String), } +impl DBErrorMarker for ProxyDetectError {} + pub struct ProxyDetectDB { - contract_address: Address, + pub contract_address: Address, code: HashMap, - values_to_storage: HashMap, - delegatecalls: Vec
+ pub values_to_storage: HashMap, + pub delegatecalls: Vec
} @@ -108,7 +120,7 @@ impl ProxyDetectDB { self.code.insert(address, code.clone()); } - fn insert_delegatecall(&mut self, contract: Address) { + pub fn insert_delegatecall(&mut self, contract: Address) { self.delegatecalls.push(contract); } } @@ -148,7 +160,8 @@ impl Database for ProxyDetectDB { todo!() } - fn storage(&mut self, address: Address,index: U256) -> Result { + fn storage(&mut self, address: Address, index: U256) -> Result { + use std::ops::{BitAnd, BitXor}; let magic_value = index.bitand(*ADDR_MASK).bitxor(*ADDR_XOR); let magic_address = Address::from_word(FixedBytes::from_slice(&magic_value.to_be_bytes::<32>())); debug!("storage(): {:x} -> {:x} = {:x}", address, index, magic_value); @@ -163,25 +176,32 @@ impl Database for ProxyDetectDB { } } - -impl Inspector for ProxyInspector { - +impl Inspector for ProxyInspector +where + CTX: ProxyDetectDBAccess, + INTR: InterpreterTypes, + INTR::Bytecode: Jumps, + INTR::Stack: StackTr, +{ #[inline(always)] fn step( &mut self, - interpreter: &mut Interpreter, - _context: &mut EvmContext, + interp: &mut Interpreter, + _context: &mut CTX, ) { // debug!("addr: {}", interpreter.contract.address); - // debug!("opcode: {}", interpreter.current_opcode()); - trace!("opcode: {}", OpCode::new(interpreter.current_opcode()).unwrap()); - for mem in interpreter.stack().data() { + let op = interp.bytecode.opcode(); + trace!("opcode: {}", revm::bytecode::opcode::OpCode::new(op).unwrap()); + for mem in interp.stack.data() { trace!("STACK: {:x}", mem); } trace!("--"); - match interpreter.current_opcode() { + match op { opcode::SLOAD => { - if let Ok(memory) = interpreter.stack.peek(0) { + // Try to get stack value at position 0 + let stack_data = interp.stack.data(); + if !stack_data.is_empty() { + let memory = stack_data[stack_data.len() - 1]; self.storage_access.push(memory); trace!("SLOAD detected {}", memory); } @@ -193,36 +213,66 @@ impl Inspector for ProxyInspector { #[inline(always)] fn call( &mut self, - context: &mut EvmContext, + context: &mut CTX, call: &mut CallInputs, ) -> Option { // println!("call!!! {:?} {}", call.scheme, call.target_address); - // return (InstructionResult::Continue, Gas::new(call.gas_limit), Bytes::new()); - if call.scheme == CallScheme::Call && call.target_address == context.db.contract_address { + let db = context.get_proxy_detect_db(); + if call.scheme == CallScheme::Call && call.target_address == db.contract_address { return None; } + + // Get the input bytes for function selector extraction + let input_bytes: Bytes = match &call.input { + revm::interpreter::CallInput::Bytes(bytes) => bytes.clone(), + revm::interpreter::CallInput::SharedBuffer(_) => { + // For shared buffer, we can't easily access the bytes without context + // Just use empty bytes as fallback + Bytes::new() + } + }; + match call.scheme { CallScheme::DelegateCall => { - context.db.delegatecalls.push(call.bytecode_address); - if let Some(storage) = context.db.values_to_storage.get(&call.bytecode_address) { + db.delegatecalls.push(call.bytecode_address); + if let Some(storage) = db.values_to_storage.get(&call.bytecode_address) { self.delegatecall_storage.push(*storage); } else { self.delegatecall_unknown.push(call.bytecode_address); } - context.db.insert_delegatecall(call.bytecode_address); + db.insert_delegatecall(call.bytecode_address); }, CallScheme::Call | CallScheme::CallCode | CallScheme::StaticCall => { - if call.input.len() >= 4 { - let fun = slice_as_u32_be(&call.input); + if input_bytes.len() >= 4 { + let fun = slice_as_u32_be(&input_bytes); self.external_calls.push((call.target_address, fun)); debug!("external call detected {:x}: {:x}", call.target_address, fun); } - } - CallScheme::ExtCall | CallScheme::ExtDelegateCall | CallScheme::ExtStaticCall => { - panic!("EIP-7069 not supported"); - } }; - Some(CallOutcome { result: InterpreterResult { result: InstructionResult::Return, output: Bytes::new(), gas: Gas::new(call.gas_limit) }, memory_offset: 0..0 }) + Some(CallOutcome::new( + InterpreterResult { + result: InstructionResult::Return, + output: Bytes::new(), + gas: Gas::new(call.gas_limit) + }, + 0..0 + )) + } +} + +/// Trait to access the ProxyDetectDB from the context. +/// This is needed because we need access to the DB in the Inspector implementation. +pub trait ProxyDetectDBAccess { + fn get_proxy_detect_db(&mut self) -> &mut ProxyDetectDB; +} + +// Implement for the Context type that revm uses +impl ProxyDetectDBAccess for Context +where + JOURNAL: JournalTr, +{ + fn get_proxy_detect_db(&mut self) -> &mut ProxyDetectDB { + self.journaled_state.db_mut() } }