1use std::collections::HashMap;
40use std::fmt;
41
42use alloy_primitives::{Address, B256, U256, keccak256};
43use revm::Inspector;
44use revm::interpreter::interpreter_types::{Jumps, MemoryTr, StackTr};
45use revm::interpreter::{Interpreter, InterpreterTypes};
46
47const OP_KECCAK256: u8 = 0x20;
49const OP_SLOAD: u8 = 0x54;
51const MAX_PREIMAGE_LEN: usize = 4096;
54
55#[derive(Clone, Copy, Debug, PartialEq, Eq)]
61pub enum SlotLayout {
62 SolidityMapping,
64 VyperMapping,
66 Nested,
68 PackedSeed {
71 seed: U256,
73 },
74 ArrayPointer,
76 Opaque,
78}
79
80impl fmt::Display for SlotLayout {
81 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
82 match self {
83 SlotLayout::SolidityMapping => f.write_str("mapping(key‖slot) [Solidity]"),
84 SlotLayout::VyperMapping => f.write_str("mapping(slot‖key) [Vyper]"),
85 SlotLayout::Nested => f.write_str("nested mapping"),
86 SlotLayout::PackedSeed { seed } => write!(f, "packed(addr‖seed={seed:#x}) [Solady]"),
87 SlotLayout::ArrayPointer => f.write_str("array/base pointer"),
88 SlotLayout::Opaque => f.write_str("opaque"),
89 }
90 }
91}
92
93#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
96pub enum Confidence {
97 Low,
99 Heuristic,
101 Medium,
103 High,
105}
106
107impl fmt::Display for Confidence {
108 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109 write!(f, "{self:?}")
110 }
111}
112
113#[derive(Clone, Debug, PartialEq, Eq)]
116pub struct HashSlotAccess {
117 pub slot: B256,
119 pub value: U256,
121 pub keys: Vec<B256>,
124 pub base_slot: U256,
126 pub layout: SlotLayout,
128 pub depth: usize,
130 pub confidence: Confidence,
132}
133
134impl HashSlotAccess {
135 pub fn keyed_by(&self, k: B256) -> bool {
138 self.keys.iter().any(|key| word_matches(*key, k))
139 }
140
141 pub fn as_tracked(&self, contract: Address) -> Option<TrackedMapping> {
145 if self.keys.len() != 1 {
146 return None;
147 }
148 match self.layout {
149 SlotLayout::SolidityMapping
150 | SlotLayout::VyperMapping
151 | SlotLayout::PackedSeed { .. } => Some(TrackedMapping {
152 contract,
153 base_slot: self.base_slot,
154 layout: self.layout,
155 }),
156 _ => None,
157 }
158 }
159}
160
161#[derive(Clone, Copy, Debug, PartialEq, Eq)]
169pub struct TrackedMapping {
170 pub contract: Address,
172 pub base_slot: U256,
174 pub layout: SlotLayout,
176}
177
178impl TrackedMapping {
179 pub fn new(contract: Address, base_slot: U256, layout: SlotLayout) -> Self {
181 Self {
182 contract,
183 base_slot,
184 layout,
185 }
186 }
187
188 pub fn slot_for(&self, key: B256) -> Option<B256> {
194 match self.layout {
195 SlotLayout::SolidityMapping => {
196 let mut pre = [0u8; 64];
197 pre[0..32].copy_from_slice(key.as_slice());
198 pre[32..64].copy_from_slice(self.base_slot_word().as_slice());
199 Some(keccak256(pre))
200 }
201 SlotLayout::VyperMapping => {
202 let mut pre = [0u8; 64];
203 pre[0..32].copy_from_slice(self.base_slot_word().as_slice());
204 pre[32..64].copy_from_slice(key.as_slice());
205 Some(keccak256(pre))
206 }
207 SlotLayout::PackedSeed { seed } => {
208 let mut pre = [0u8; 32];
211 pre[0..20].copy_from_slice(&Address::from_word(key).into_array());
212 let seed_be = seed.to_be_bytes::<32>();
213 pre[20..32].copy_from_slice(&seed_be[20..32]);
214 Some(keccak256(pre))
215 }
216 _ => None,
217 }
218 }
219
220 pub fn slots_for(&self, keys: impl IntoIterator<Item = B256>) -> Vec<(B256, B256)> {
222 keys.into_iter()
223 .filter_map(|k| self.slot_for(k).map(|s| (k, s)))
224 .collect()
225 }
226
227 fn base_slot_word(&self) -> B256 {
228 B256::from(self.base_slot.to_be_bytes::<32>())
229 }
230}
231
232pub type TrackedBalances = (TrackedMapping, Vec<(Address, B256)>);
236
237#[derive(Clone, Debug)]
242struct SloadRecord {
243 slot: B256,
244 value: U256,
245}
246
247#[derive(Clone, Debug, Default)]
265pub struct HashStorageProbe {
266 preimages: HashMap<B256, Vec<u8>>,
268 reads: Vec<SloadRecord>,
270 pending: Option<B256>,
272}
273
274impl HashStorageProbe {
275 pub fn new() -> Self {
277 Self::default()
278 }
279
280 pub fn preimage_count(&self) -> usize {
282 self.preimages.len()
283 }
284
285 pub fn hashed_read_count(&self) -> usize {
287 self.reads
288 .iter()
289 .filter(|r| self.preimages.contains_key(&r.slot))
290 .count()
291 }
292
293 pub fn slots_returning(&self, value: U256) -> Vec<B256> {
301 let mut seen = std::collections::HashSet::new();
302 self.reads
303 .iter()
304 .filter(|r| r.value == value)
305 .map(|r| r.slot)
306 .filter(|slot| seen.insert(*slot))
307 .collect()
308 }
309
310 pub fn accesses(&self, known: &[B256]) -> Vec<HashSlotAccess> {
327 self.reads
328 .iter()
329 .filter(|r| self.preimages.contains_key(&r.slot))
330 .map(|r| resolve(r.slot, r.value, &self.preimages, known))
331 .collect()
332 }
333}
334
335impl<CTX, INTR: InterpreterTypes> Inspector<CTX, INTR> for HashStorageProbe {
336 fn step(&mut self, interp: &mut Interpreter<INTR>, _ctx: &mut CTX) {
337 match interp.bytecode.opcode() {
338 OP_KECCAK256 => {
339 let (offset, size) = {
341 let s = interp.stack.data();
342 let n = s.len();
343 if n < 2 {
344 return;
345 }
346 (s[n - 1], s[n - 2])
347 };
348 let (Some(offset), Some(size)) = (to_usize(offset), to_usize(size)) else {
349 return;
350 };
351 if size == 0 || size > MAX_PREIMAGE_LEN {
352 return;
353 }
354 let preimage = read_mem(interp, offset, size);
355 self.preimages.insert(keccak256(&preimage), preimage);
356 }
357 OP_SLOAD => {
358 if let Some(k) = interp.stack.data().last() {
361 self.pending = Some(word_from(*k));
362 }
363 }
364 _ => {}
365 }
366 }
367
368 fn step_end(&mut self, interp: &mut Interpreter<INTR>, _ctx: &mut CTX) {
369 if let Some(slot) = self.pending.take() {
370 if let Some(value) = interp.stack.data().last().copied() {
372 self.reads.push(SloadRecord { slot, value });
373 }
374 }
375 }
376}
377
378fn resolve(
384 slot: B256,
385 value: U256,
386 pre: &HashMap<B256, Vec<u8>>,
387 known: &[B256],
388) -> HashSlotAccess {
389 let mut keys: Vec<B256> = Vec::new();
390 let mut confidence = Confidence::High;
391 let mut cur = slot;
392
393 let (base_slot, layout) = loop {
394 let Some(p) = pre.get(&cur) else {
395 let l = if keys.is_empty() {
398 SlotLayout::Opaque
399 } else {
400 SlotLayout::Nested
401 };
402 break (U256::from_be_slice(cur.as_slice()), l);
403 };
404
405 match p.len() {
406 64 => {
407 let aw = B256::from_slice(&p[0..32]);
408 let bw = B256::from_slice(&p[32..64]);
409 let a_parent = pre.contains_key(&aw);
410 let b_parent = pre.contains_key(&bw);
411 if a_parent ^ b_parent {
412 let (parent, key) = if a_parent { (aw, bw) } else { (bw, aw) };
414 keys.push(key);
415 cur = parent;
416 continue;
417 }
418 let (slot_word, key_word, key_first, conf) =
419 split_key_slot(&p[0..32], &p[32..64], known);
420 confidence = confidence.min(conf);
421 keys.push(key_word);
422 let layout = if keys.len() > 1 {
423 SlotLayout::Nested
424 } else if key_first {
425 SlotLayout::SolidityMapping
426 } else {
427 SlotLayout::VyperMapping
428 };
429 break (U256::from_be_slice(slot_word.as_slice()), layout);
430 }
431 32 => {
432 let hi = Address::from_slice(&p[0..20]);
434 if hi != Address::ZERO && known.iter().any(|k| Address::from_word(*k) == hi) {
435 keys.push(hi.into_word());
436 confidence = confidence.min(Confidence::Medium);
437 let seed = U256::from_be_slice(&p[20..32]);
438 break (seed, SlotLayout::PackedSeed { seed });
439 }
440 confidence = confidence.min(Confidence::Heuristic);
442 let l = if keys.is_empty() {
443 SlotLayout::ArrayPointer
444 } else {
445 SlotLayout::Nested
446 };
447 break (U256::from_be_slice(&p[0..32]), l);
448 }
449 _ => {
450 confidence = Confidence::Low;
451 break (U256::from_be_slice(cur.as_slice()), SlotLayout::Opaque);
452 }
453 }
454 };
455
456 let depth = keys.len().max(1);
457 HashSlotAccess {
458 slot,
459 value,
460 keys,
461 base_slot,
462 layout,
463 depth,
464 confidence,
465 }
466}
467
468fn split_key_slot(a: &[u8], b: &[u8], known: &[B256]) -> (B256, B256, bool, Confidence) {
471 let aw = B256::from_slice(a);
472 let bw = B256::from_slice(b);
473 let a_known = known.iter().any(|k| word_matches(*k, aw));
474 let b_known = known.iter().any(|k| word_matches(*k, bw));
475 match (a_known, b_known) {
476 (true, false) => (bw, aw, true, Confidence::High), (false, true) => (aw, bw, false, Confidence::High), _ => {
479 let (sa, sb) = (sig(a), sig(b));
482 if sa < sb {
483 (aw, bw, false, Confidence::Medium)
484 } else if sb < sa {
485 (bw, aw, true, Confidence::Medium)
486 } else {
487 (aw, bw, false, Confidence::Low)
488 }
489 }
490 }
491}
492
493fn read_mem<INTR: InterpreterTypes>(
496 interp: &Interpreter<INTR>,
497 offset: usize,
498 len: usize,
499) -> Vec<u8> {
500 let mut out = vec![0u8; len];
501 let size = interp.memory.size();
502 if offset >= size || len == 0 {
503 return out;
504 }
505 let avail = (size - offset).min(len);
506 let chunk = interp.memory.slice_len(offset, avail);
507 out[..avail].copy_from_slice(&chunk[..]);
508 out
509}
510
511fn to_usize(x: U256) -> Option<usize> {
512 let limbs = x.as_limbs();
513 if limbs[1] | limbs[2] | limbs[3] != 0 {
514 return None;
515 }
516 usize::try_from(limbs[0]).ok()
517}
518
519fn word_from(x: U256) -> B256 {
520 B256::from(x.to_be_bytes::<32>())
521}
522
523fn word_matches(a: B256, b: B256) -> bool {
526 a == b || Address::from_word(a) == Address::from_word(b)
527}
528
529fn sig(bytes: &[u8]) -> usize {
531 match bytes.iter().position(|&x| x != 0) {
532 Some(p) => bytes.len() - p,
533 None => 0,
534 }
535}
536
537#[cfg(test)]
542mod tests {
543 use super::*;
544 use alloy_primitives::address;
545
546 fn sol_slot(key: Address, base: u64) -> B256 {
547 let mut pre = [0u8; 64];
548 pre[0..32].copy_from_slice(key.into_word().as_slice());
549 pre[63] = base as u8;
550 keccak256(pre)
551 }
552
553 fn vyper_slot(base: u64, key: Address) -> B256 {
554 let mut pre = [0u8; 64];
555 pre[31] = base as u8;
556 pre[32..64].copy_from_slice(key.into_word().as_slice());
557 keccak256(pre)
558 }
559
560 fn solady_slot(key: Address, seed: u32) -> B256 {
561 let mut pre = [0u8; 32];
562 pre[0..20].copy_from_slice(&key.into_array());
563 pre[28..32].copy_from_slice(&seed.to_be_bytes());
564 keccak256(pre)
565 }
566
567 fn probe_with(preimages: Vec<Vec<u8>>, reads: Vec<(B256, U256)>) -> HashStorageProbe {
569 let mut p = HashStorageProbe::new();
570 for pre in preimages {
571 p.preimages.insert(keccak256(&pre), pre);
572 }
573 for (slot, value) in reads {
574 p.reads.push(SloadRecord { slot, value });
575 }
576 p
577 }
578
579 #[test]
580 fn slots_returning_matches_value_and_dedups() {
581 let key = address!("00000000000000000000000000000000000000A1");
584 let mut pre = vec![0u8; 64];
585 pre[0..32].copy_from_slice(key.into_word().as_slice());
586 pre[63] = 3;
587 let hashed_slot = keccak256(&pre);
588 let plain_slot = B256::from(U256::from(2u64).to_be_bytes::<32>()); let mut probe = probe_with(vec![pre], vec![]);
591 probe.reads.push(SloadRecord {
593 slot: hashed_slot,
594 value: U256::from(100u64),
595 });
596 probe.reads.push(SloadRecord {
597 slot: hashed_slot,
598 value: U256::from(100u64),
599 });
600 probe.reads.push(SloadRecord {
601 slot: plain_slot,
602 value: U256::from(100u64),
603 });
604 probe.reads.push(SloadRecord {
605 slot: plain_slot,
606 value: U256::from(7u64),
607 }); assert_eq!(
611 probe.slots_returning(U256::from(100u64)),
612 vec![hashed_slot, plain_slot]
613 );
614 assert_eq!(probe.hashed_read_count(), 2); assert_eq!(probe.accesses(&[key.into_word()]).len(), 2); }
618
619 #[test]
620 fn resolves_solidity_mapping() {
621 let key = address!("00000000000000000000000000000000000000A1");
622 let mut pre = vec![0u8; 64];
623 pre[0..32].copy_from_slice(key.into_word().as_slice());
624 pre[63] = 3;
625 let slot = keccak256(&pre);
626 let probe = probe_with(vec![pre], vec![(slot, U256::from(42u64))]);
627
628 let a = &probe.accesses(&[key.into_word()])[0];
629 assert_eq!(a.layout, SlotLayout::SolidityMapping);
630 assert_eq!(a.base_slot, U256::from(3u64));
631 assert_eq!(a.keys, vec![key.into_word()]);
632 assert_eq!(a.confidence, Confidence::High);
633 }
634
635 #[test]
636 fn resolves_vyper_order_without_known_key() {
637 let key = address!("28C6c06298d514Db089934071355E5743bf21d60");
640 let mut pre = vec![0u8; 64];
641 pre[31] = 2;
642 pre[32..64].copy_from_slice(key.into_word().as_slice());
643 let slot = keccak256(&pre);
644 let probe = probe_with(vec![pre], vec![(slot, U256::from(1u64))]);
645
646 let a = &probe.accesses(&[])[0];
648 assert_eq!(a.layout, SlotLayout::VyperMapping);
649 assert_eq!(a.base_slot, U256::from(2u64));
650 assert_eq!(a.confidence, Confidence::Medium);
651 }
652
653 #[test]
654 fn resolves_solady_packed() {
655 let key = address!("00000000000000000000000000000000000000A1");
656 let seed = 0x87a2_11a2u32;
657 let mut pre = vec![0u8; 32];
658 pre[0..20].copy_from_slice(&key.into_array());
659 pre[28..32].copy_from_slice(&seed.to_be_bytes());
660 let slot = keccak256(&pre);
661 let probe = probe_with(vec![pre], vec![(slot, U256::from(7u64))]);
662
663 let a = &probe.accesses(&[key.into_word()])[0];
664 assert_eq!(
665 a.layout,
666 SlotLayout::PackedSeed {
667 seed: U256::from(seed)
668 }
669 );
670 assert!(a.keyed_by(key.into_word()));
671 }
672
673 #[test]
674 fn resolves_nested_mapping() {
675 let owner = address!("00000000000000000000000000000000000000A1");
676 let spender = address!("00000000000000000000000000000000000000B2");
677 let mut inner_pre = vec![0u8; 64];
679 inner_pre[0..32].copy_from_slice(owner.into_word().as_slice());
680 inner_pre[63] = 4;
681 let inner = keccak256(&inner_pre);
682 let mut outer_pre = vec![0u8; 64];
683 outer_pre[0..32].copy_from_slice(spender.into_word().as_slice());
684 outer_pre[32..64].copy_from_slice(inner.as_slice());
685 let outer = keccak256(&outer_pre);
686
687 let probe = probe_with(vec![inner_pre, outer_pre], vec![(outer, U256::from(9u64))]);
688 let a = &probe.accesses(&[owner.into_word(), spender.into_word()])[0];
689 assert_eq!(a.layout, SlotLayout::Nested);
690 assert_eq!(a.base_slot, U256::from(4u64));
691 assert_eq!(a.depth, 2);
692 assert_eq!(a.keys, vec![spender.into_word(), owner.into_word()]);
693 }
694
695 #[test]
696 fn tracked_mapping_round_trips_each_layout() {
697 let key = address!("00000000000000000000000000000000000000A1");
698
699 let t = TrackedMapping::new(Address::ZERO, U256::from(3u64), SlotLayout::SolidityMapping);
700 assert_eq!(t.slot_for(key.into_word()).unwrap(), sol_slot(key, 3));
701
702 let t = TrackedMapping::new(Address::ZERO, U256::from(2u64), SlotLayout::VyperMapping);
703 assert_eq!(t.slot_for(key.into_word()).unwrap(), vyper_slot(2, key));
704
705 let seed = 0x87a2_11a2u32;
706 let t = TrackedMapping::new(
707 Address::ZERO,
708 U256::from(seed),
709 SlotLayout::PackedSeed {
710 seed: U256::from(seed),
711 },
712 );
713 assert_eq!(t.slot_for(key.into_word()).unwrap(), solady_slot(key, seed));
714 }
715
716 #[test]
717 fn as_tracked_rejects_nested_and_arrays() {
718 let access = HashSlotAccess {
719 slot: B256::ZERO,
720 value: U256::ZERO,
721 keys: vec![B256::ZERO, B256::ZERO],
722 base_slot: U256::from(4u64),
723 layout: SlotLayout::Nested,
724 depth: 2,
725 confidence: Confidence::High,
726 };
727 assert!(access.as_tracked(Address::ZERO).is_none());
728 }
729}