foundry_cheatcodes/evm/
mapping.rs

1use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
2use alloy_primitives::{
3    Address, B256, U256, keccak256,
4    map::{AddressHashMap, B256HashMap},
5};
6use alloy_sol_types::SolValue;
7use revm::{
8    bytecode::opcode,
9    interpreter::{Interpreter, interpreter_types::Jumps},
10};
11
12/// Recorded mapping slots.
13#[derive(Clone, Debug, Default)]
14pub struct MappingSlots {
15    /// Holds mapping parent (slots => slots)
16    pub parent_slots: B256HashMap<B256>,
17
18    /// Holds mapping key (slots => key)
19    pub keys: B256HashMap<B256>,
20
21    /// Holds mapping child (slots => slots[])
22    pub children: B256HashMap<Vec<B256>>,
23
24    /// Holds the last sha3 result `sha3_result => (data_low, data_high)`, this would only record
25    /// when sha3 is called with `size == 0x40`, and the lower 256 bits would be stored in
26    /// `data_low`, higher 256 bits in `data_high`.
27    /// This is needed for mapping_key detect if the slot is for some mapping and record that.
28    pub seen_sha3: B256HashMap<(B256, B256)>,
29}
30
31impl MappingSlots {
32    /// Tries to insert a mapping slot. Returns true if it was inserted.
33    pub fn insert(&mut self, slot: B256) -> bool {
34        match self.seen_sha3.get(&slot).copied() {
35            Some((key, parent)) => {
36                if self.keys.contains_key(&slot) {
37                    return false;
38                }
39                self.keys.insert(slot, key);
40                self.parent_slots.insert(slot, parent);
41                self.children.entry(parent).or_default().push(slot);
42                self.insert(parent);
43                true
44            }
45            None => false,
46        }
47    }
48}
49
50impl Cheatcode for startMappingRecordingCall {
51    fn apply(&self, state: &mut Cheatcodes) -> Result {
52        let Self {} = self;
53        if state.mapping_slots.is_none() {
54            state.mapping_slots = Some(Default::default());
55        }
56        Ok(Default::default())
57    }
58}
59
60impl Cheatcode for stopMappingRecordingCall {
61    fn apply(&self, state: &mut Cheatcodes) -> Result {
62        let Self {} = self;
63        state.mapping_slots = None;
64        Ok(Default::default())
65    }
66}
67
68impl Cheatcode for getMappingLengthCall {
69    fn apply(&self, state: &mut Cheatcodes) -> Result {
70        let Self { target, mappingSlot } = self;
71        let result = slot_child(state, target, mappingSlot).map(Vec::len).unwrap_or(0);
72        Ok((result as u64).abi_encode())
73    }
74}
75
76impl Cheatcode for getMappingSlotAtCall {
77    fn apply(&self, state: &mut Cheatcodes) -> Result {
78        let Self { target, mappingSlot, idx } = self;
79        let result = slot_child(state, target, mappingSlot)
80            .and_then(|set| set.get(idx.saturating_to::<usize>()))
81            .copied()
82            .unwrap_or_default();
83        Ok(result.abi_encode())
84    }
85}
86
87impl Cheatcode for getMappingKeyAndParentOfCall {
88    fn apply(&self, state: &mut Cheatcodes) -> Result {
89        let Self { target, elementSlot: slot } = self;
90        let mut found = false;
91        let mut key = &B256::ZERO;
92        let mut parent = &B256::ZERO;
93        if let Some(slots) = mapping_slot(state, target) {
94            if let Some(key2) = slots.keys.get(slot) {
95                found = true;
96                key = key2;
97                parent = &slots.parent_slots[slot];
98            } else if let Some((key2, parent2)) = slots.seen_sha3.get(slot) {
99                found = true;
100                key = key2;
101                parent = parent2;
102            }
103        }
104        Ok((found, key, parent).abi_encode_params())
105    }
106}
107
108fn mapping_slot<'a>(state: &'a Cheatcodes, target: &'a Address) -> Option<&'a MappingSlots> {
109    state.mapping_slots.as_ref()?.get(target)
110}
111
112fn slot_child<'a>(
113    state: &'a Cheatcodes,
114    target: &'a Address,
115    slot: &'a B256,
116) -> Option<&'a Vec<B256>> {
117    mapping_slot(state, target)?.children.get(slot)
118}
119
120#[cold]
121pub(crate) fn step(mapping_slots: &mut AddressHashMap<MappingSlots>, interpreter: &Interpreter) {
122    match interpreter.bytecode.opcode() {
123        opcode::KECCAK256 => {
124            if interpreter.stack.peek(1) == Ok(U256::from(0x40)) {
125                let address = interpreter.input.target_address;
126                let offset = interpreter.stack.peek(0).expect("stack size > 1").saturating_to();
127                let data = interpreter.memory.slice_len(offset, 0x40);
128                let low = B256::from_slice(&data[..0x20]);
129                let high = B256::from_slice(&data[0x20..]);
130                let result = keccak256(&*data);
131
132                mapping_slots.entry(address).or_default().seen_sha3.insert(result, (low, high));
133            }
134        }
135        opcode::SSTORE => {
136            if let Some(mapping_slots) = mapping_slots.get_mut(&interpreter.input.target_address)
137                && let Ok(slot) = interpreter.stack.peek(0)
138            {
139                mapping_slots.insert(slot.into());
140            }
141        }
142        _ => {}
143    }
144}