revm/inspector/
handler_register.rs

1use crate::{
2    db::Database,
3    handler::register::EvmHandler,
4    interpreter::{opcode, InstructionResult, Interpreter},
5    primitives::EVMError,
6    Context, FrameOrResult, FrameResult, Inspector, JournalEntry,
7};
8use core::cell::RefCell;
9use revm_interpreter::opcode::DynInstruction;
10use std::{rc::Rc, sync::Arc, vec::Vec};
11
12/// Provides access to an `Inspector` instance.
13pub trait GetInspector<DB: Database> {
14    /// Returns the associated `Inspector`.
15    fn get_inspector(&mut self) -> &mut impl Inspector<DB>;
16}
17
18impl<DB: Database, INSP: Inspector<DB>> GetInspector<DB> for INSP {
19    #[inline]
20    fn get_inspector(&mut self) -> &mut impl Inspector<DB> {
21        self
22    }
23}
24
25/// Register Inspector handles that interact with Inspector instance.
26///
27///
28/// # Note
29///
30/// Inspector handle register does not override any existing handlers, and it
31/// calls them before (or after) calling Inspector. This means that it is safe
32/// to use this register with any other register.
33///
34/// A few instructions handlers are wrapped twice once for `step` and `step_end`
35/// and in case of Logs and Selfdestruct wrapper is wrapped again for the
36/// `log` and `selfdestruct` calls.
37pub fn inspector_handle_register<DB: Database, EXT: GetInspector<DB>>(
38    handler: &mut EvmHandler<'_, EXT, DB>,
39) {
40    let table = &mut handler.instruction_table;
41
42    // Update all instructions to call inspector step and step_end.
43    table.update_all(inspector_instruction);
44
45    // Register inspector LOG* instructions.
46    for opcode in opcode::LOG0..=opcode::LOG4 {
47        table.update_boxed(opcode, move |prev, interpreter, host| {
48            let prev_log_len = host.evm.journaled_state.logs.len();
49            prev(interpreter, host);
50            // check if log was added. It is possible that revert happened
51            // cause of gas or stack underflow.
52            if host.evm.journaled_state.logs.len() == prev_log_len + 1 {
53                // clone log.
54                // TODO decide if we should remove this and leave the comment
55                // that log can be found as journaled_state.
56                let last_log = host.evm.journaled_state.logs.last().unwrap().clone();
57                // call Inspector
58                host.external
59                    .get_inspector()
60                    .log(interpreter, &mut host.evm, &last_log);
61            }
62        });
63    }
64
65    // Register selfdestruct function.
66    table.update_boxed(opcode::SELFDESTRUCT, |prev, interpreter, host| {
67        // execute selfdestruct
68        prev(interpreter, host);
69        // check if selfdestruct was successful and if journal entry is made.
70        match host.evm.journaled_state.journal.last().unwrap().last() {
71            Some(JournalEntry::AccountDestroyed {
72                address,
73                target,
74                had_balance,
75                ..
76            }) => {
77                host.external
78                    .get_inspector()
79                    .selfdestruct(*address, *target, *had_balance);
80            }
81            Some(JournalEntry::BalanceTransfer {
82                from, to, balance, ..
83            }) => {
84                host.external
85                    .get_inspector()
86                    .selfdestruct(*from, *to, *balance);
87            }
88            _ => {}
89        }
90    });
91
92    // call and create input stack shared between handlers. They are used to share
93    // inputs in *_end Inspector calls.
94    let call_input_stack = Rc::<RefCell<Vec<_>>>::default();
95    let create_input_stack = Rc::<RefCell<Vec<_>>>::default();
96    let eofcreate_input_stack = Rc::<RefCell<Vec<_>>>::default();
97
98    // Create handler
99    let create_input_stack_inner = create_input_stack.clone();
100    let prev_handle = handler.execution.create.clone();
101    handler.execution.create = Arc::new(
102        move |ctx, mut inputs| -> Result<FrameOrResult, EVMError<DB::Error>> {
103            let inspector = ctx.external.get_inspector();
104            // call inspector create to change input or return outcome.
105            if let Some(outcome) = inspector.create(&mut ctx.evm, &mut inputs) {
106                create_input_stack_inner.borrow_mut().push(inputs.clone());
107                return Ok(FrameOrResult::Result(FrameResult::Create(outcome)));
108            }
109            create_input_stack_inner.borrow_mut().push(inputs.clone());
110
111            let mut frame_or_result = prev_handle(ctx, inputs);
112            if let Ok(FrameOrResult::Frame(frame)) = &mut frame_or_result {
113                ctx.external
114                    .get_inspector()
115                    .initialize_interp(frame.interpreter_mut(), &mut ctx.evm)
116            }
117            frame_or_result
118        },
119    );
120
121    // Call handler
122    let call_input_stack_inner = call_input_stack.clone();
123    let prev_handle = handler.execution.call.clone();
124    handler.execution.call = Arc::new(move |ctx, mut inputs| {
125        // Call inspector to change input or return outcome.
126        let outcome = ctx.external.get_inspector().call(&mut ctx.evm, &mut inputs);
127        call_input_stack_inner.borrow_mut().push(inputs.clone());
128        if let Some(outcome) = outcome {
129            return Ok(FrameOrResult::Result(FrameResult::Call(outcome)));
130        }
131
132        let mut frame_or_result = prev_handle(ctx, inputs);
133        if let Ok(FrameOrResult::Frame(frame)) = &mut frame_or_result {
134            ctx.external
135                .get_inspector()
136                .initialize_interp(frame.interpreter_mut(), &mut ctx.evm)
137        }
138        frame_or_result
139    });
140
141    // Calls inspector `eofcreate` and `initialize_interp` functions. Queues the inputs for the `eofcreate_end`` function.
142    // Calls the old handler, and in case of inspector returning outcome,
143    // returns the outcome without executing eofcreate.
144    let eofcreate_input_stack_inner = eofcreate_input_stack.clone();
145    let prev_handle = handler.execution.eofcreate.clone();
146    handler.execution.eofcreate = Arc::new(move |ctx, mut inputs| {
147        // Call inspector to change input or return outcome.
148        let outcome = ctx
149            .external
150            .get_inspector()
151            .eofcreate(&mut ctx.evm, &mut inputs);
152        eofcreate_input_stack_inner
153            .borrow_mut()
154            .push(inputs.clone());
155        if let Some(outcome) = outcome {
156            return Ok(FrameOrResult::Result(FrameResult::EOFCreate(outcome)));
157        }
158
159        let mut frame_or_result = prev_handle(ctx, inputs);
160        if let Ok(FrameOrResult::Frame(frame)) = &mut frame_or_result {
161            ctx.external
162                .get_inspector()
163                .initialize_interp(frame.interpreter_mut(), &mut ctx.evm)
164        }
165        frame_or_result
166    });
167
168    // Pops eofcreate input from the stack and calls inspector `eofcreate_end` function.
169    // preserve the old handler and calls it with the outcome.
170    let eofcreate_input_stack_inner = eofcreate_input_stack.clone();
171    let prev_handle = handler.execution.insert_eofcreate_outcome.clone();
172    handler.execution.insert_eofcreate_outcome = Arc::new(move |ctx, frame, mut outcome| {
173        let create_inputs = eofcreate_input_stack_inner.borrow_mut().pop().unwrap();
174        outcome = ctx
175            .external
176            .get_inspector()
177            .eofcreate_end(&mut ctx.evm, &create_inputs, outcome);
178        prev_handle(ctx, frame, outcome)
179    });
180
181    // call outcome
182    let call_input_stack_inner = call_input_stack.clone();
183    let prev_handle = handler.execution.insert_call_outcome.clone();
184    handler.execution.insert_call_outcome =
185        Arc::new(move |ctx, frame, shared_memory, mut outcome| {
186            let call_inputs = call_input_stack_inner.borrow_mut().pop().unwrap();
187            outcome = ctx
188                .external
189                .get_inspector()
190                .call_end(&mut ctx.evm, &call_inputs, outcome);
191            prev_handle(ctx, frame, shared_memory, outcome)
192        });
193
194    // create outcome
195    let create_input_stack_inner = create_input_stack.clone();
196    let prev_handle = handler.execution.insert_create_outcome.clone();
197    handler.execution.insert_create_outcome = Arc::new(move |ctx, frame, mut outcome| {
198        let create_inputs = create_input_stack_inner.borrow_mut().pop().unwrap();
199        outcome = ctx
200            .external
201            .get_inspector()
202            .create_end(&mut ctx.evm, &create_inputs, outcome);
203        prev_handle(ctx, frame, outcome)
204    });
205
206    // last frame outcome
207    let prev_handle = handler.execution.last_frame_return.clone();
208    handler.execution.last_frame_return = Arc::new(move |ctx, frame_result| {
209        let inspector = ctx.external.get_inspector();
210        match frame_result {
211            FrameResult::Call(outcome) => {
212                let call_inputs = call_input_stack.borrow_mut().pop().unwrap();
213                *outcome = inspector.call_end(&mut ctx.evm, &call_inputs, outcome.clone());
214            }
215            FrameResult::Create(outcome) => {
216                let create_inputs = create_input_stack.borrow_mut().pop().unwrap();
217                *outcome = inspector.create_end(&mut ctx.evm, &create_inputs, outcome.clone());
218            }
219            FrameResult::EOFCreate(outcome) => {
220                let eofcreate_inputs = eofcreate_input_stack.borrow_mut().pop().unwrap();
221                *outcome =
222                    inspector.eofcreate_end(&mut ctx.evm, &eofcreate_inputs, outcome.clone());
223            }
224        }
225        prev_handle(ctx, frame_result)
226    });
227}
228
229fn inspector_instruction<INSP, DB>(
230    prev: &DynInstruction<'_, Context<INSP, DB>>,
231    interpreter: &mut Interpreter,
232    host: &mut Context<INSP, DB>,
233) where
234    INSP: GetInspector<DB>,
235    DB: Database,
236{
237    // SAFETY: as the PC was already incremented we need to subtract 1 to preserve the
238    // old Inspector behavior.
239    interpreter.instruction_pointer = unsafe { interpreter.instruction_pointer.sub(1) };
240
241    // Call step.
242    host.external
243        .get_inspector()
244        .step(interpreter, &mut host.evm);
245    if interpreter.instruction_result != InstructionResult::Continue {
246        return;
247    }
248
249    // Reset PC to previous value.
250    interpreter.instruction_pointer = unsafe { interpreter.instruction_pointer.add(1) };
251
252    // Execute instruction.
253    prev(interpreter, host);
254
255    // Call step_end.
256    host.external
257        .get_inspector()
258        .step_end(interpreter, &mut host.evm);
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264    use crate::{
265        inspectors::NoOpInspector,
266        interpreter::{CallInputs, CallOutcome, CreateInputs, CreateOutcome},
267        Evm, EvmContext,
268    };
269
270    #[derive(Default, Debug)]
271    struct StackInspector {
272        initialize_interp_called: bool,
273        step: u32,
274        step_end: u32,
275        call: bool,
276        call_end: bool,
277    }
278
279    impl<DB: Database> Inspector<DB> for StackInspector {
280        fn initialize_interp(&mut self, _interp: &mut Interpreter, _context: &mut EvmContext<DB>) {
281            if self.initialize_interp_called {
282                unreachable!("initialize_interp should not be called twice")
283            }
284            self.initialize_interp_called = true;
285        }
286
287        fn step(&mut self, _interp: &mut Interpreter, _context: &mut EvmContext<DB>) {
288            self.step += 1;
289        }
290
291        fn step_end(&mut self, _interp: &mut Interpreter, _context: &mut EvmContext<DB>) {
292            self.step_end += 1;
293        }
294
295        fn call(
296            &mut self,
297            context: &mut EvmContext<DB>,
298            _call: &mut CallInputs,
299        ) -> Option<CallOutcome> {
300            if self.call {
301                unreachable!("call should not be called twice")
302            }
303            self.call = true;
304            assert_eq!(context.journaled_state.depth(), 0);
305            None
306        }
307
308        fn call_end(
309            &mut self,
310            context: &mut EvmContext<DB>,
311            _inputs: &CallInputs,
312            outcome: CallOutcome,
313        ) -> CallOutcome {
314            if self.call_end {
315                unreachable!("call_end should not be called twice")
316            }
317            assert_eq!(context.journaled_state.depth(), 0);
318            self.call_end = true;
319            outcome
320        }
321
322        fn create(
323            &mut self,
324            context: &mut EvmContext<DB>,
325            _call: &mut CreateInputs,
326        ) -> Option<CreateOutcome> {
327            assert_eq!(context.journaled_state.depth(), 0);
328            None
329        }
330
331        fn create_end(
332            &mut self,
333            context: &mut EvmContext<DB>,
334            _inputs: &CreateInputs,
335            outcome: CreateOutcome,
336        ) -> CreateOutcome {
337            assert_eq!(context.journaled_state.depth(), 0);
338            outcome
339        }
340    }
341
342    #[test]
343    fn test_inspector_handlers() {
344        use crate::{
345            db::BenchmarkDB,
346            inspector::inspector_handle_register,
347            interpreter::opcode,
348            primitives::{address, Bytecode, Bytes, TxKind},
349            Evm,
350        };
351
352        let contract_data: Bytes = Bytes::from(vec![
353            opcode::PUSH1,
354            0x1,
355            opcode::PUSH1,
356            0xb,
357            opcode::PUSH1,
358            0x1,
359            opcode::PUSH1,
360            0x1,
361            opcode::PUSH1,
362            0x1,
363            opcode::CREATE,
364            opcode::STOP,
365        ]);
366        let bytecode = Bytecode::new_raw(contract_data);
367
368        let mut evm: Evm<'_, StackInspector, BenchmarkDB> = Evm::builder()
369            .with_db(BenchmarkDB::new_bytecode(bytecode.clone()))
370            .with_external_context(StackInspector::default())
371            .modify_tx_env(|tx| {
372                tx.clear();
373                tx.caller = address!("1000000000000000000000000000000000000000");
374                tx.transact_to = TxKind::Call(address!("0000000000000000000000000000000000000000"));
375                tx.gas_limit = 21100;
376            })
377            .append_handler_register(inspector_handle_register)
378            .build();
379
380        // run evm.
381        evm.transact().unwrap();
382
383        let inspector = evm.into_context().external;
384
385        assert_eq!(inspector.step, 6);
386        assert_eq!(inspector.step_end, 6);
387        assert!(inspector.initialize_interp_called);
388        assert!(inspector.call);
389        assert!(inspector.call_end);
390    }
391
392    #[test]
393    fn test_inspector_reg() {
394        let mut noop = NoOpInspector;
395        let _evm = Evm::builder()
396            .with_external_context(&mut noop)
397            .append_handler_register(inspector_handle_register)
398            .build();
399    }
400}