foundry_cheatcodes/evm/
mapping.rs

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

foundry_cheatcodes/evm/
mapping.rs
