evm_lens_core/storage/resolvers/
heuristic.rs

1use async_trait::async_trait;
2use revm::bytecode::OpCode;
3
4use crate::disassemble;
5use crate::storage::layout::{Provenance, StorageEntry, StorageLayout, StorageType};
6use crate::storage::resolver::StorageLayoutResolver;
7
8/// Heuristic that scans for PUSH <const> followed by SLOAD/SSTORE.
9/// If found, treat the const as a direct slot index. Types are Unknown.
10pub struct HeuristicResolver;
11
12#[async_trait]
13impl StorageLayoutResolver for HeuristicResolver {
14    async fn resolve(&self, input: &[u8]) -> color_eyre::Result<Option<StorageLayout>> {
15        if input.is_empty() {
16            return Ok(Some(StorageLayout::new()));
17        }
18
19        let ops = match disassemble(input) {
20            Ok(v) => v,
21            Err(_) => return Ok(Some(StorageLayout::new())),
22        };
23
24        let mut layout = StorageLayout::new();
25
26        // Windowed scan: PUSH[1..32] followed within small window by SLOAD/SSTORE
27        let window = 6usize;
28
29        for (idx, (_pos, op)) in ops.iter().enumerate() {
30            let push_n = match *op {
31                OpCode::PUSH0 => Some(0u8),
32                OpCode::PUSH1 => Some(1),
33                OpCode::PUSH2 => Some(2),
34                OpCode::PUSH3 => Some(3),
35                OpCode::PUSH4 => Some(4),
36                OpCode::PUSH5 => Some(5),
37                OpCode::PUSH6 => Some(6),
38                OpCode::PUSH7 => Some(7),
39                OpCode::PUSH8 => Some(8),
40                OpCode::PUSH9 => Some(9),
41                OpCode::PUSH10 => Some(10),
42                OpCode::PUSH11 => Some(11),
43                OpCode::PUSH12 => Some(12),
44                OpCode::PUSH13 => Some(13),
45                OpCode::PUSH14 => Some(14),
46                OpCode::PUSH15 => Some(15),
47                OpCode::PUSH16 => Some(16),
48                OpCode::PUSH17 => Some(17),
49                OpCode::PUSH18 => Some(18),
50                OpCode::PUSH19 => Some(19),
51                OpCode::PUSH20 => Some(20),
52                OpCode::PUSH21 => Some(21),
53                OpCode::PUSH22 => Some(22),
54                OpCode::PUSH23 => Some(23),
55                OpCode::PUSH24 => Some(24),
56                OpCode::PUSH25 => Some(25),
57                OpCode::PUSH26 => Some(26),
58                OpCode::PUSH27 => Some(27),
59                OpCode::PUSH28 => Some(28),
60                OpCode::PUSH29 => Some(29),
61                OpCode::PUSH30 => Some(30),
62                OpCode::PUSH31 => Some(31),
63                OpCode::PUSH32 => Some(32),
64                _ => None,
65            };
66
67            if push_n.is_none() {
68                continue;
69            }
70
71            // Look ahead up to `window` ops for SLOAD/SSTORE as a conservative signal
72            let mut found = false;
73            for j in 1..=window {
74                if let Some((_, next_op)) = ops.get(idx + j) {
75                    match *next_op {
76                        OpCode::SLOAD | OpCode::SSTORE => {
77                            found = true;
78                            break;
79                        }
80                        _ => {}
81                    }
82                }
83            }
84            if !found {
85                continue;
86            }
87
88            // Extract pushed immediate as big-endian integer from underlying byte slice.
89            // We conservatively re-scan input bytes at the recorded opcode position to get immediates.
90            // For v0.3 simplicity, we approximate by reading from the raw hex window preceding the op.
91            let (pos, _op) = ops[idx];
92            // The immediate begins at pos+1 and spans push_n bytes; guard bounds.
93            let n = push_n.unwrap() as usize;
94            if pos + 1 + n > input.len() {
95                continue;
96            }
97            let imm = &input[pos + 1..pos + 1 + n];
98            // Convert to u128 if fits; skip otherwise
99            if n > 16 {
100                continue;
101            }
102            let mut slot: u128 = 0;
103            // Shift left by 8 bits for each byte
104            for &b in imm {
105                slot = (slot << 8) | b as u128;
106            }
107
108            // Deduplicate per slot
109            if layout.entries.iter().any(|e| e.slot == slot) {
110                continue;
111            }
112            layout.add_entry(StorageEntry {
113                slot,
114                offset: None,
115                size: None,
116                r#type: StorageType::Unknown,
117                label: None,
118                provenance: Provenance::HeuristicTrace,
119            });
120        }
121
122        Ok(Some(layout))
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[tokio::test]
131    async fn detects_push1_followed_by_sload() {
132        // 60 00 54 00 => PUSH1 0x00; SLOAD; STOP
133        let resolver = HeuristicResolver;
134        let bytes = hex::decode("60005400").unwrap();
135        let layout = resolver.resolve(&bytes).await.unwrap().unwrap();
136        assert_eq!(layout.entries.len(), 1);
137        assert_eq!(layout.entries[0].slot, 0);
138    }
139
140    #[tokio::test]
141    async fn detects_push2_followed_by_sstore() {
142        // 61 ab cd 55 00 => PUSH2 0xABCD; SSTORE; STOP
143        let resolver = HeuristicResolver;
144        let bytes = hex::decode("61abcd5500").unwrap();
145        let layout = resolver.resolve(&bytes).await.unwrap().unwrap();
146        assert_eq!(layout.entries.len(), 1);
147        assert_eq!(layout.entries[0].slot, 0xABCD);
148    }
149
150    #[tokio::test]
151    async fn no_detection_beyond_window() {
152        // PUSH1 0x00, then 7 non-storage ops (DUP1), then SLOAD => beyond window=6
153        // 60 00 80 80 80 80 80 80 80 54 00
154        let resolver = HeuristicResolver;
155        let bytes = hex::decode("6000808080808080805400").unwrap();
156        let layout = resolver.resolve(&bytes).await.unwrap().unwrap();
157        assert_eq!(layout.entries.len(), 0);
158    }
159
160    #[tokio::test]
161    async fn dedup_same_slot_multiple_occurrences() {
162        // PUSH1 0x01; SLOAD; PUSH1 0x01; SSTORE; STOP
163        // 60 01 54 60 01 55 00
164        let resolver = HeuristicResolver;
165        let bytes = hex::decode("60015460015500").unwrap();
166        let layout = resolver.resolve(&bytes).await.unwrap().unwrap();
167        assert_eq!(layout.entries.len(), 1);
168        assert_eq!(layout.entries[0].slot, 1);
169    }
170
171    #[tokio::test]
172    async fn skip_push32_immediate() {
173        // 7f followed by 32 bytes of 0x00, then SLOAD => n>16 should skip
174        let resolver = HeuristicResolver;
175        let bytes = hex::decode(&format!("7f{}54", "00".repeat(32))).unwrap();
176        let layout = resolver.resolve(&bytes).await.unwrap().unwrap();
177        assert_eq!(layout.entries.len(), 0);
178    }
179}
evm_lens_core/storage/resolvers/heuristic.rs

evm_lens_core/storage/resolvers/
heuristic.rs
