foundry_common/
mapping_slots.rs

1use alloy_primitives::{
2    B256, U256, keccak256,
3    map::{AddressHashMap, B256HashMap},
4};
5use revm::{
6    bytecode::opcode,
7    interpreter::{Interpreter, interpreter_types::Jumps},
8};
9
10/// Recorded mapping slots.
11#[derive(Clone, Debug, Default)]
12pub struct MappingSlots {
13    /// Holds mapping parent (slots => slots)
14    pub parent_slots: B256HashMap<B256>,
15
16    /// Holds mapping key (slots => key)
17    pub keys: B256HashMap<B256>,
18
19    /// Holds mapping child (slots => slots[])
20    pub children: B256HashMap<Vec<B256>>,
21
22    /// Holds the last sha3 result `sha3_result => (data_low, data_high)`, this would only record
23    /// when sha3 is called with `size == 0x40`, and the lower 256 bits would be stored in
24    /// `data_low`, higher 256 bits in `data_high`.
25    /// This is needed for mapping_key detect if the slot is for some mapping and record that.
26    pub seen_sha3: B256HashMap<(B256, B256)>,
27}
28
29impl MappingSlots {
30    /// Tries to insert a mapping slot. Returns true if it was inserted.
31    pub fn insert(&mut self, slot: B256) -> bool {
32        match self.seen_sha3.get(&slot).copied() {
33            Some((key, parent)) => {
34                if self.keys.insert(slot, key).is_some() {
35                    return false;
36                }
37                self.parent_slots.insert(slot, parent);
38                self.children.entry(parent).or_default().push(slot);
39                self.insert(parent);
40                true
41            }
42            None => false,
43        }
44    }
45}
46
47/// Function to be used in Inspector::step to record mapping slots and keys
48#[cold]
49pub fn step(mapping_slots: &mut AddressHashMap<MappingSlots>, interpreter: &Interpreter) {
50    match interpreter.bytecode.opcode() {
51        opcode::KECCAK256 => {
52            if interpreter.stack.peek(1) == Ok(U256::from(0x40)) {
53                let address = interpreter.input.target_address;
54                let offset = interpreter.stack.peek(0).expect("stack size > 1").saturating_to();
55                let data = interpreter.memory.slice_len(offset, 0x40);
56                let low = B256::from_slice(&data[..0x20]);
57                let high = B256::from_slice(&data[0x20..]);
58                let result = keccak256(&*data);
59
60                mapping_slots.entry(address).or_default().seen_sha3.insert(result, (low, high));
61            }
62        }
63        opcode::SSTORE => {
64            if let Some(mapping_slots) = mapping_slots.get_mut(&interpreter.input.target_address)
65                && let Ok(slot) = interpreter.stack.peek(0)
66            {
67                mapping_slots.insert(slot.into());
68            }
69        }
70        _ => {}
71    }
72}
foundry_common/mapping_slots.rs

foundry_common/
mapping_slots.rs
