Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
39 changes: 17 additions & 22 deletions src/detect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 {
Expand Down
112 changes: 81 additions & 31 deletions src/proxy_inspector.rs
Original file line number Diff line number Diff line change
@@ -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::{
Expand Down Expand Up @@ -83,14 +92,17 @@ static ADDR_XOR: Lazy<U256> = 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<Address, Bytes>,
values_to_storage: HashMap<Address, U256>,
delegatecalls: Vec<Address>
pub values_to_storage: HashMap<Address, U256>,
pub delegatecalls: Vec<Address>
}


Expand All @@ -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);
}
}
Expand Down Expand Up @@ -148,7 +160,8 @@ impl Database for ProxyDetectDB {
todo!()
}

fn storage(&mut self, address: Address,index: U256) -> Result<U256,Self::Error> {
fn storage(&mut self, address: Address, index: U256) -> Result<U256,Self::Error> {
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);
Expand All @@ -163,25 +176,32 @@ impl Database for ProxyDetectDB {
}
}


impl Inspector<ProxyDetectDB> for ProxyInspector {

impl<CTX, INTR> Inspector<CTX, INTR> 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<ProxyDetectDB>,
interp: &mut Interpreter<INTR>,
_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);
}
Expand All @@ -193,36 +213,66 @@ impl Inspector<ProxyDetectDB> for ProxyInspector {
#[inline(always)]
fn call(
&mut self,
context: &mut EvmContext<ProxyDetectDB>,
context: &mut CTX,
call: &mut CallInputs,
) -> Option<CallOutcome> {
// 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<BLOCK, TX, CFG, JOURNAL> ProxyDetectDBAccess for Context<BLOCK, TX, CFG, ProxyDetectDB, JOURNAL>
where
JOURNAL: JournalTr<Database = ProxyDetectDB>,
{
fn get_proxy_detect_db(&mut self) -> &mut ProxyDetectDB {
self.journaled_state.db_mut()
}
}