foundry_evm/inspectors/
edge_cov.rs

1use alloy_primitives::{
2    Address, U256,
3    map::{DefaultHashBuilder, HashMap},
4};
5use core::{
6    fmt,
7    hash::{BuildHasher, Hash, Hasher},
8};
9use revm::{
10    Inspector,
11    bytecode::opcode,
12    interpreter::{
13        Interpreter,
14        interpreter_types::{InputsTr, Jumps},
15    },
16};
17
18// This is the maximum number of edges that can be tracked. There is a tradeoff between performance
19// and precision (less collisions).
20const MAX_EDGE_COUNT: usize = 65536;
21
22// Maximum number of unique comparison sites to track for CmpLog-style feedback.
23const MAX_CMP_LOG_SITES: usize = 1024;
24
25// Maximum number of comparison operand pairs to track per site. This matches the downstream loop
26// detection threshold so a hot loop can be classified without filling the whole log.
27const MAX_CMP_OBSERVATIONS_PER_SITE: u8 = 8;
28
29/// A comparison operand pair captured during execution.
30#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
31pub struct CmpOperands {
32    /// First operand of the comparison.
33    pub op1: U256,
34    /// Second operand of the comparison.
35    pub op2: U256,
36    /// Program counter where the comparison occurred.
37    pub pc: usize,
38    /// Contract address where the comparison occurred.
39    pub address: Address,
40    /// EVM opcode that performed the comparison.
41    pub opcode: u8,
42}
43
44/// An `Inspector` that tracks [edge coverage](https://clang.llvm.org/docs/SanitizerCoverage.html#edge-coverage).
45/// Covered edges will not wrap to zero e.g. a loop edge hit more than 255 will still be retained.
46///
47/// Also tracks comparison operands for CmpLog-style guided fuzzing.
48// see https://github.com/AFLplusplus/AFLplusplus/blob/5777ceaf23f48ae4ceae60e4f3a79263802633c6/instrumentation/afl-llvm-pass.so.cc#L810-L829
49#[derive(Clone)]
50pub struct EdgeCovInspector {
51    /// Map of hitcounts that can be diffed against to determine if new coverage was reached.
52    hitcount: Vec<u8>,
53    hash_builder: DefaultHashBuilder,
54    /// Comparison operand log for CmpLog-style guided fuzzing.
55    cmp_log: Option<Vec<CmpOperands>>,
56    cmp_site_counts: HashMap<CmpSiteKey, u8>,
57}
58
59impl fmt::Debug for EdgeCovInspector {
60    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
61        f.debug_struct("EdgeCovInspector").finish_non_exhaustive()
62    }
63}
64
65impl EdgeCovInspector {
66    /// Create a new `EdgeCovInspector` with `MAX_EDGE_COUNT` size.
67    pub fn new() -> Self {
68        Self {
69            hitcount: vec![0; MAX_EDGE_COUNT],
70            hash_builder: DefaultHashBuilder::default(),
71            cmp_log: None,
72            cmp_site_counts: HashMap::default(),
73        }
74    }
75
76    /// Create a new `EdgeCovInspector` with comparison operand logging enabled.
77    pub fn with_cmp_log() -> Self {
78        let mut inspector = Self::new();
79        inspector.enable_cmp_log(true);
80        inspector
81    }
82
83    /// Set whether to collect comparison operand logs.
84    pub fn enable_cmp_log(&mut self, yes: bool) {
85        if yes {
86            self.cmp_log.get_or_insert_with(|| Vec::with_capacity(MAX_CMP_LOG_SITES));
87        } else {
88            self.cmp_log = None;
89            self.cmp_site_counts.clear();
90        }
91    }
92
93    /// Reset the hitcount to zero and clear the comparison log.
94    pub fn reset(&mut self) {
95        self.hitcount.fill(0);
96        if let Some(cmp_log) = &mut self.cmp_log {
97            cmp_log.clear();
98        }
99        self.cmp_site_counts.clear();
100    }
101
102    /// Get an immutable reference to the hitcount.
103    pub const fn get_hitcount(&self) -> &[u8] {
104        self.hitcount.as_slice()
105    }
106
107    /// Get an immutable reference to the comparison operand log.
108    pub const fn get_cmp_log(&self) -> &[CmpOperands] {
109        match &self.cmp_log {
110            Some(cmp_log) => cmp_log.as_slice(),
111            None => &[],
112        }
113    }
114
115    /// Consume the inspector and take ownership of the hitcount.
116    pub fn into_hitcount(self) -> Vec<u8> {
117        self.hitcount
118    }
119
120    /// Consume the inspector and take ownership of both the hitcount and comparison log.
121    pub fn into_parts(self) -> (Vec<u8>, Vec<CmpOperands>) {
122        (self.hitcount, self.cmp_log.unwrap_or_default())
123    }
124
125    /// Mark the edge, H(address, pc, jump_dest), as hit.
126    fn store_hit(&mut self, address: Address, pc: usize, jump_dest: U256) {
127        let mut hasher = self.hash_builder.build_hasher();
128        address.hash(&mut hasher);
129        pc.hash(&mut hasher);
130        jump_dest.hash(&mut hasher);
131        // The hash is used to index into the hitcount array,
132        // so it must be modulo the maximum edge count.
133        let edge_id = (hasher.finish() % MAX_EDGE_COUNT as u64) as usize;
134        self.hitcount[edge_id] = self.hitcount[edge_id].checked_add(1).unwrap_or(1);
135    }
136
137    /// Store comparison operands for CmpLog-style guided fuzzing.
138    fn store_cmp(&mut self, cmp: CmpOperands) {
139        let Some(cmp_log) = &mut self.cmp_log else {
140            return;
141        };
142
143        let site = CmpSiteKey::new(&cmp);
144        if let Some(count) = self.cmp_site_counts.get_mut(&site) {
145            if *count >= MAX_CMP_OBSERVATIONS_PER_SITE {
146                return;
147            }
148            *count += 1;
149            cmp_log.push(cmp);
150        } else if self.cmp_site_counts.len() < MAX_CMP_LOG_SITES {
151            self.cmp_site_counts.insert(site, 1);
152            cmp_log.push(cmp);
153        }
154    }
155
156    #[cold]
157    fn do_step(&mut self, interp: &mut Interpreter) {
158        let address = interp.input.target_address();
159        let current_pc = interp.bytecode.pc();
160
161        match interp.bytecode.opcode() {
162            opcode::JUMP => {
163                // unconditional jump
164                if let Ok(jump_dest) = interp.stack.peek(0) {
165                    self.store_hit(address, current_pc, jump_dest);
166                }
167            }
168            opcode::JUMPI => {
169                if let Ok(stack_value) = interp.stack.peek(1) {
170                    let jump_dest = if stack_value.is_zero() {
171                        // fall through
172                        Ok(U256::from(current_pc + 1))
173                    } else {
174                        // branch taken
175                        interp.stack.peek(0)
176                    };
177
178                    if let Ok(jump_dest) = jump_dest {
179                        self.store_hit(address, current_pc, jump_dest);
180                    }
181                }
182            }
183            _ => {
184                // no-op
185            }
186        }
187    }
188
189    #[cold]
190    fn do_cmp_step(&mut self, interp: &mut Interpreter) {
191        if self.cmp_log.is_none() {
192            return;
193        }
194
195        let address = interp.input.target_address();
196        let current_pc = interp.bytecode.pc();
197
198        match interp.bytecode.opcode() {
199            op @ (opcode::EQ | opcode::LT | opcode::SLT | opcode::GT | opcode::SGT) => {
200                if let (Ok(op1), Ok(op2)) = (interp.stack.peek(0), interp.stack.peek(1)) {
201                    self.store_cmp(CmpOperands { op1, op2, pc: current_pc, address, opcode: op });
202                }
203            }
204            op @ opcode::ISZERO => {
205                if let Ok(op1) = interp.stack.peek(0) {
206                    self.store_cmp(CmpOperands {
207                        op1,
208                        op2: U256::ZERO,
209                        pc: current_pc,
210                        address,
211                        opcode: op,
212                    });
213                }
214            }
215            _ => {}
216        }
217    }
218}
219
220impl Default for EdgeCovInspector {
221    fn default() -> Self {
222        Self::new()
223    }
224}
225
226impl<CTX> Inspector<CTX> for EdgeCovInspector {
227    #[inline]
228    fn step(&mut self, interp: &mut Interpreter, _context: &mut CTX) {
229        let op = interp.bytecode.opcode();
230        if matches!(op, opcode::JUMP | opcode::JUMPI) {
231            self.do_step(interp);
232        }
233        if self.cmp_log.is_some()
234            && matches!(
235                op,
236                opcode::EQ | opcode::LT | opcode::GT | opcode::SLT | opcode::SGT | opcode::ISZERO
237            )
238        {
239            self.do_cmp_step(interp);
240        }
241    }
242}
243
244#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
245struct CmpSiteKey {
246    address: Address,
247    pc: usize,
248    opcode: u8,
249}
250
251impl CmpSiteKey {
252    const fn new(cmp: &CmpOperands) -> Self {
253        Self { address: cmp.address, pc: cmp.pc, opcode: cmp.opcode }
254    }
255}
256
257#[cfg(test)]
258mod tests {
259    use super::*;
260
261    #[test]
262    fn cmp_operands_defaults_and_clones() {
263        let cmp = CmpOperands {
264            op1: U256::from(123),
265            op2: U256::from(456),
266            pc: 42,
267            address: Address::repeat_byte(0xaa),
268            opcode: opcode::EQ,
269        };
270
271        assert_eq!(cmp.op1, U256::from(123));
272        assert_eq!(cmp.op2, U256::from(456));
273        assert_eq!(cmp.pc, 42);
274
275        assert_eq!(CmpOperands::default(), CmpOperands::default());
276        let cloned = cmp;
277        assert_eq!(cloned, cmp);
278    }
279
280    #[test]
281    fn cmp_log_starts_empty_and_is_returned_by_into_parts() {
282        let inspector = EdgeCovInspector::new();
283
284        assert!(inspector.get_cmp_log().is_empty());
285        assert!(inspector.get_hitcount().iter().all(|&x| x == 0));
286
287        let (hitcount, cmp_log) = inspector.into_parts();
288        assert_eq!(hitcount.len(), MAX_EDGE_COUNT);
289        assert!(cmp_log.is_empty());
290    }
291
292    #[test]
293    fn reset_clears_hitcount_and_cmp_log() {
294        let mut inspector = EdgeCovInspector::with_cmp_log();
295
296        inspector.hitcount[0] = 1;
297        inspector.store_cmp(CmpOperands {
298            op1: U256::from(123),
299            op2: U256::from(456),
300            pc: 42,
301            address: Address::ZERO,
302            opcode: opcode::EQ,
303        });
304
305        inspector.reset();
306
307        assert!(inspector.get_hitcount().iter().all(|&x| x == 0));
308        assert!(inspector.get_cmp_log().is_empty());
309    }
310
311    #[test]
312    fn cmp_log_is_capped_per_site() {
313        let mut inspector = EdgeCovInspector::with_cmp_log();
314
315        for i in 0..usize::from(MAX_CMP_OBSERVATIONS_PER_SITE) + 1 {
316            inspector.store_cmp(CmpOperands {
317                op1: U256::from(i),
318                op2: U256::from(i + 1),
319                pc: 42,
320                address: Address::ZERO,
321                opcode: opcode::EQ,
322            });
323        }
324
325        assert_eq!(inspector.get_cmp_log().len(), usize::from(MAX_CMP_OBSERVATIONS_PER_SITE));
326    }
327
328    #[test]
329    fn cmp_log_caps_sites_without_starving_later_observations() {
330        let mut inspector = EdgeCovInspector::with_cmp_log();
331
332        for i in 0..usize::from(MAX_CMP_OBSERVATIONS_PER_SITE) * 2 {
333            inspector.store_cmp(CmpOperands {
334                op1: U256::from(i),
335                op2: U256::from(i + 1),
336                pc: 1,
337                address: Address::ZERO,
338                opcode: opcode::EQ,
339            });
340        }
341        inspector.store_cmp(CmpOperands {
342            op1: U256::from(123),
343            op2: U256::from(456),
344            pc: 2,
345            address: Address::ZERO,
346            opcode: opcode::EQ,
347        });
348
349        assert_eq!(inspector.get_cmp_log().len(), usize::from(MAX_CMP_OBSERVATIONS_PER_SITE) + 1);
350        assert_eq!(inspector.get_cmp_log().last().unwrap().pc, 2);
351    }
352}
foundry_evm/inspectors/edge_cov.rs

foundry_evm/inspectors/
edge_cov.rs
