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::interpreter::{opcode, Interpreter};
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.contains_key(&slot) {
35                    return false
36                }
37                self.keys.insert(slot, key);
38                self.parent_slots.insert(slot, parent);
39                self.children.entry(parent).or_default().push(slot);
40                self.insert(parent);
41                true
42            }
43            None => false,
44        }
45    }
46}
47
48impl Cheatcode for startMappingRecordingCall {
49    fn apply(&self, state: &mut Cheatcodes) -> Result {
50        let Self {} = self;
51        if state.mapping_slots.is_none() {
52            state.mapping_slots = Some(Default::default());
53        }
54        Ok(Default::default())
55    }
56}
57
58impl Cheatcode for stopMappingRecordingCall {
59    fn apply(&self, state: &mut Cheatcodes) -> Result {
60        let Self {} = self;
61        state.mapping_slots = None;
62        Ok(Default::default())
63    }
64}
65
66impl Cheatcode for getMappingLengthCall {
67    fn apply(&self, state: &mut Cheatcodes) -> Result {
68        let Self { target, mappingSlot } = self;
69        let result = slot_child(state, target, mappingSlot).map(Vec::len).unwrap_or(0);
70        Ok((result as u64).abi_encode())
71    }
72}
73
74impl Cheatcode for getMappingSlotAtCall {
75    fn apply(&self, state: &mut Cheatcodes) -> Result {
76        let Self { target, mappingSlot, idx } = self;
77        let result = slot_child(state, target, mappingSlot)
78            .and_then(|set| set.get(idx.saturating_to::<usize>()))
79            .copied()
80            .unwrap_or_default();
81        Ok(result.abi_encode())
82    }
83}
84
85impl Cheatcode for getMappingKeyAndParentOfCall {
86    fn apply(&self, state: &mut Cheatcodes) -> Result {
87        let Self { target, elementSlot: slot } = self;
88        let mut found = false;
89        let mut key = &B256::ZERO;
90        let mut parent = &B256::ZERO;
91        if let Some(slots) = mapping_slot(state, target) {
92            if let Some(key2) = slots.keys.get(slot) {
93                found = true;
94                key = key2;
95                parent = &slots.parent_slots[slot];
96            } else if let Some((key2, parent2)) = slots.seen_sha3.get(slot) {
97                found = true;
98                key = key2;
99                parent = parent2;
100            }
101        }
102        Ok((found, key, parent).abi_encode_params())
103    }
104}
105
106fn mapping_slot<'a>(state: &'a Cheatcodes, target: &'a Address) -> Option<&'a MappingSlots> {
107    state.mapping_slots.as_ref()?.get(target)
108}
109
110fn slot_child<'a>(
111    state: &'a Cheatcodes,
112    target: &'a Address,
113    slot: &'a B256,
114) -> Option<&'a Vec<B256>> {
115    mapping_slot(state, target)?.children.get(slot)
116}
117
118#[cold]
119pub(crate) fn step(mapping_slots: &mut AddressHashMap<MappingSlots>, interpreter: &Interpreter) {
120    match interpreter.current_opcode() {
121        opcode::KECCAK256 => {
122            if interpreter.stack.peek(1) == Ok(U256::from(0x40)) {
123                let address = interpreter.contract.target_address;
124                let offset = interpreter.stack.peek(0).expect("stack size > 1").saturating_to();
125                let data = interpreter.shared_memory.slice(offset, 0x40);
126                let low = B256::from_slice(&data[..0x20]);
127                let high = B256::from_slice(&data[0x20..]);
128                let result = keccak256(data);
129
130                mapping_slots.entry(address).or_default().seen_sha3.insert(result, (low, high));
131            }
132        }
133        opcode::SSTORE => {
134            if let Some(mapping_slots) = mapping_slots.get_mut(&interpreter.contract.target_address)
135            {
136                if let Ok(slot) = interpreter.stack.peek(0) {
137                    mapping_slots.insert(slot.into());
138                }
139            }
140        }
141        _ => {}
142    }
143}
foundry_cheatcodes/evm/mapping.rs

foundry_cheatcodes/evm/
mapping.rs
