foundry_common/
mapping_slots.rs

1use alloy_primitives::{
2    Address, B256, U256,
3    map::{AddressHashMap, B256HashMap},
4};
5use revm::{
6    bytecode::opcode,
7    interpreter::{Interpreter, interpreter_types::Jumps},
8};
9
10/// Provenance recovered for a Solidity mapping storage slot.
11#[derive(Clone, Debug, PartialEq, Eq)]
12pub struct MappingProvenance {
13    /// The terminal mapping root slot.
14    pub root_slot: B256,
15    /// Mapping keys ordered from the root mapping to the accessed value.
16    pub keys: Vec<B256>,
17}
18
19/// Recorded mapping slots.
20#[derive(Clone, Debug, Default)]
21pub struct MappingSlots {
22    /// Holds mapping parent (slots => slots)
23    pub parent_slots: B256HashMap<B256>,
24
25    /// Holds mapping key (slots => key)
26    pub keys: B256HashMap<B256>,
27
28    /// Holds mapping child (slots => slots[])
29    pub children: B256HashMap<Vec<B256>>,
30
31    /// Holds the last sha3 result `sha3_result => (data_low, data_high)`, this would only record
32    /// when sha3 is called with `size == 0x40`, and the lower 256 bits would be stored in
33    /// `data_low`, higher 256 bits in `data_high`.
34    /// This is needed for mapping_key detect if the slot is for some mapping and record that.
35    pub seen_sha3: B256HashMap<(B256, B256)>,
36}
37
38impl MappingSlots {
39    /// Records the result and two input words of a 64-byte `KECCAK256` operation.
40    pub fn record_hash(&mut self, result: B256, key: B256, parent: B256) {
41        self.seen_sha3.insert(result, (key, parent));
42    }
43
44    /// Resolves a computed slot to its terminal mapping root and root-to-leaf keys.
45    pub fn resolve(&self, slot: B256) -> Option<MappingProvenance> {
46        let mut current = slot;
47        let mut keys = Vec::new();
48        while let Some((key, parent)) = self.seen_sha3.get(&current).copied() {
49            keys.push(key);
50            current = parent;
51            if keys.len() > self.seen_sha3.len() {
52                return None;
53            }
54        }
55        if keys.is_empty() {
56            return None;
57        }
58        keys.reverse();
59        Some(MappingProvenance { root_slot: current, keys })
60    }
61
62    /// Tries to insert a mapping slot. Returns true if it was inserted.
63    pub fn insert(&mut self, slot: B256) -> bool {
64        match self.seen_sha3.get(&slot).copied() {
65            Some((key, parent)) => {
66                if self.keys.insert(slot, key).is_some() {
67                    return false;
68                }
69                self.parent_slots.insert(slot, parent);
70                self.children.entry(parent).or_default().push(slot);
71                self.insert(parent);
72                true
73            }
74            None => false,
75        }
76    }
77}
78
79/// A pending 64-byte Keccak operation captured before execution.
80#[derive(Clone, Copy, Debug)]
81pub struct PendingMappingHash {
82    /// The effective storage address of the executing frame.
83    pub address: Address,
84    /// The memory offset containing the Keccak preimage.
85    pub offset: usize,
86}
87
88/// Captures a 64-byte Keccak operation before execution.
89pub fn capture_hash(interpreter: &Interpreter) -> Option<PendingMappingHash> {
90    if interpreter.bytecode.opcode() != opcode::KECCAK256
91        || interpreter.stack.peek(1).ok()? != U256::from(0x40)
92    {
93        return None;
94    }
95    Some(PendingMappingHash {
96        address: interpreter.input.target_address,
97        offset: interpreter.stack.peek(0).ok()?.try_into().ok()?,
98    })
99}
100
101/// Records a successfully executed 64-byte Keccak operation after memory expansion.
102pub fn record_hash(
103    mapping_slots: &mut AddressHashMap<MappingSlots>,
104    interpreter: &Interpreter,
105    pending: PendingMappingHash,
106) {
107    let Ok(result) = interpreter.stack.peek(0) else { return };
108    let data = interpreter.memory.slice_len(pending.offset, 0x40);
109    let key = B256::from_slice(&data[..0x20]);
110    let parent = B256::from_slice(&data[0x20..]);
111    mapping_slots.entry(pending.address).or_default().record_hash(result.into(), key, parent);
112}
113
114/// Function to be used in `Inspector::step` to record mapping slots.
115#[cold]
116pub fn step(mapping_slots: &mut AddressHashMap<MappingSlots>, interpreter: &Interpreter) {
117    if interpreter.bytecode.opcode() == opcode::SSTORE
118        && let Some(mapping_slots) = mapping_slots.get_mut(&interpreter.input.target_address)
119        && let Ok(slot) = interpreter.stack.peek(0)
120    {
121        mapping_slots.insert(slot.into());
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128    use alloy_primitives::keccak256;
129
130    #[test]
131    fn resolves_mapping_keys_from_root_to_leaf() {
132        let root = B256::with_last_byte(1);
133        let owner = B256::with_last_byte(2);
134        let spender = B256::with_last_byte(3);
135        let inner = keccak256([owner.as_slice(), root.as_slice()].concat());
136        let slot = keccak256([spender.as_slice(), inner.as_slice()].concat());
137        let mut slots = MappingSlots::default();
138        slots.record_hash(inner, owner, root);
139        slots.record_hash(slot, spender, inner);
140
141        assert_eq!(
142            slots.resolve(slot),
143            Some(MappingProvenance { root_slot: root, keys: vec![owner, spender] })
144        );
145    }
146
147    #[test]
148    fn rejects_plain_slots_offsets_and_cycles() {
149        let root = B256::with_last_byte(1);
150        let key = B256::with_last_byte(2);
151        let slot = keccak256([key.as_slice(), root.as_slice()].concat());
152        let mut slots = MappingSlots::default();
153        slots.record_hash(slot, key, root);
154
155        assert!(slots.resolve(root).is_none());
156        assert!(slots.resolve(B256::from(U256::from_be_bytes(slot.0) + U256::ONE)).is_none());
157
158        slots.record_hash(root, key, slot);
159        assert!(slots.resolve(slot).is_none());
160    }
161}
foundry_common/mapping_slots.rs

foundry_common/
mapping_slots.rs
