iptr_edge_analyzer/control_flow_handler/
fuzz_bitmap.rs

1//! This module contains fuzz bitmap control flow handler logics.
2
3#[cfg(feature = "cache")]
4use std::{num::NonZero, ops::Range};
5
6use crate::{ControlFlowTransitionKind, HandleControlFlow};
7
8/// [`HandleControlFlow`] implementor for maintaining fuzzing bitmap
9pub struct FuzzBitmapControlFlowHandler<M: AsRef<[u8]> + AsMut<[u8]>> {
10    /// Already recorded bitmap indices in current cache.
11    ///
12    /// This is used for quickly locate bitmap count in [`per_cache_bitmap`][Self::per_cache_bitmap].
13    #[cfg(feature = "cache")]
14    per_cache_recorded_bitmap_indices: Vec<u32>,
15    /// Bitmap like fuzzing bitmap, but used for current cache only.
16    ///
17    /// This is created with the same size as fuzzing bitmap, while each element indicates
18    /// the count of transition among this cache in the corresponding position.
19    #[cfg(feature = "cache")]
20    per_cache_bitmap: Box<[u8]>,
21    /// This is the actual structure holding the cache data. The cached key
22    /// is a range into this list, and each element is a (pos, count) pair.
23    ///
24    /// This list will always have one dummy element at decode begin. By this approach,
25    /// we can make sure the real indices into this list are always non-zero, which can
26    /// make the cached key even smaller using Rust's niche optimization.
27    #[cfg(feature = "cache")]
28    bitmap_entries_arena: Vec<CompactBitmapEntry>,
29    /// The fuzzing bitmap needed to be maintained.
30    fuzzing_bitmap: M,
31    /// Range of valid instruction addresses, if given.
32    ///
33    /// For instruction out of the range, it will not update
34    /// hit count of fuzz bitmap.
35    filter_range: Option<Box<[(u64, u64)]>>,
36    /// Previous location used to calculating fuzzing bitmap index.
37    prev_loc: u64,
38}
39
40/// Initial size of [`per_cache_recorded_bitmap_indices`][FuzzBitmapControlFlowHandler::per_cache_recorded_bitmap_indices].
41#[cfg(feature = "cache")]
42const INITIAL_RESULTS_PER_CACHE: usize = 64;
43/// Initial size of [`bitmap_entries_arena`][FuzzBitmapControlFlowHandler::bitmap_entries_arena].
44#[cfg(feature = "cache")]
45const INITIAL_BITMAP_ENTRIES_ARENA_SIZE: usize = 0x100;
46
47impl<M: AsRef<[u8]> + AsMut<[u8]>> FuzzBitmapControlFlowHandler<M> {
48    /// Create a new fuzz bitmap control flow handler.
49    ///
50    /// You can pass things like `&mut [u8]`, `Vec<u8>`, `Box<[u8]>`, or even a mmapped structure
51    /// as `fuzzing_bitmap`. If you want to give range restrictions, pass `filter_range`,
52    /// or you could just pass a [`None`] here to indicate that there is no
53    /// range restrictions.
54    pub fn new(fuzzing_bitmap: M, filter_range: Option<&[(u64, u64)]>) -> Self {
55        #[cfg(feature = "cache")]
56        let bitmap_size = fuzzing_bitmap.as_ref().len();
57        #[cfg(feature = "cache")]
58        let mut bitmap_entries_arena = Vec::with_capacity(INITIAL_BITMAP_ENTRIES_ARENA_SIZE);
59        #[cfg(feature = "cache")]
60        bitmap_entries_arena.push(DUMMY_BITMAP_ENTRY);
61        Self {
62            #[cfg(feature = "cache")]
63            per_cache_recorded_bitmap_indices: Vec::with_capacity(INITIAL_RESULTS_PER_CACHE),
64            #[cfg(feature = "cache")]
65            per_cache_bitmap: vec![0u8; bitmap_size].into_boxed_slice(),
66            #[cfg(feature = "cache")]
67            bitmap_entries_arena,
68            filter_range: filter_range.map(Box::from),
69            fuzzing_bitmap,
70            prev_loc: 0,
71        }
72    }
73
74    #[inline]
75    fn is_addr_in_filter_range(&self, address: u64) -> bool {
76        let Some(filter_range) = &self.filter_range else {
77            return true;
78        };
79        for (start, end) in filter_range {
80            if (*start..=*end).contains(&address) {
81                return true;
82            }
83        }
84        false
85    }
86
87    /// Get fuzz bitmap size as a modulus for calculating bitmap index
88    fn bitmap_size_modulus(&self) -> u64 {
89        self.fuzzing_bitmap.as_ref().len() as u64
90    }
91
92    /// Update [`prev_loc`][FuzzBitmapControlFlowHandler::prev_loc] and calculate bitmap index
93    #[expect(clippy::cast_possible_truncation)]
94    fn on_new_loc(&mut self, new_loc: u64) -> usize {
95        let bitmap_index = self.prev_loc ^ new_loc;
96        self.set_new_loc(new_loc);
97        (bitmap_index % self.bitmap_size_modulus()) as usize
98    }
99
100    /// Set [`prev_loc`][FuzzBitmapControlFlowHandler::prev_loc] without calculating bitmap index
101    fn set_new_loc(&mut self, new_loc: u64) {
102        self.prev_loc = new_loc >> 1;
103    }
104
105    /// Get diagnose information
106    pub fn diagnose(&self) -> FuzzBitmapDiagnosticInformation {
107        FuzzBitmapDiagnosticInformation {
108            #[cfg(feature = "cache")]
109            bitmap_entries_count: self.bitmap_entries_arena.len(),
110        }
111    }
112}
113
114/// Diagnostic information for [`FuzzBitmapControlFlowHandler`].
115///
116/// This struct can be retrieved from [`FuzzBitmapControlFlowHandler::diagnose`]
117pub struct FuzzBitmapDiagnosticInformation {
118    /// Number of raw bitmap entries stored in cache structure
119    #[cfg(feature = "cache")]
120    pub bitmap_entries_count: usize,
121}
122
123impl<M: AsRef<[u8]> + AsMut<[u8]>> HandleControlFlow for FuzzBitmapControlFlowHandler<M> {
124    type Error = std::convert::Infallible;
125    #[cfg(feature = "cache")]
126    type CachedKey = PerCacheBitmapEntries;
127
128    fn at_decode_begin(&mut self) -> Result<(), Self::Error> {
129        self.prev_loc = 0;
130        #[cfg(feature = "cache")]
131        self.clear_current_cache();
132        Ok(())
133    }
134
135    #[inline]
136    #[cfg_attr(feature = "cache", expect(clippy::cast_possible_truncation))]
137    #[expect(clippy::enum_glob_use)]
138    fn on_new_block(
139        &mut self,
140        block_addr: u64,
141        transition_kind: ControlFlowTransitionKind,
142        cache: bool,
143    ) -> Result<(), Self::Error> {
144        use ControlFlowTransitionKind::*;
145        if !self.is_addr_in_filter_range(block_addr) {
146            self.set_new_loc(0);
147            return Ok(());
148        }
149        match transition_kind {
150            ConditionalBranch | Indirect | DirectJump | DirectCall => {
151                let bitmap_index = self.on_new_loc(block_addr);
152                debug_assert!(
153                    bitmap_index < self.fuzzing_bitmap.as_ref().len(),
154                    "Unexpected OOB"
155                );
156                let count = unsafe { self.fuzzing_bitmap.as_mut().get_unchecked_mut(bitmap_index) };
157                *count = count.wrapping_add(1);
158                #[cfg(feature = "cache")]
159                if cache {
160                    // SAFETY: bitmap index is caculated by modulo
161                    debug_assert!(bitmap_index < self.per_cache_bitmap.len(), "Unexpected OOB");
162                    let count = unsafe { self.per_cache_bitmap.get_unchecked_mut(bitmap_index) };
163                    if *count == 0 {
164                        debug_assert!(u32::try_from(bitmap_index).is_ok(), "Bitmap size too large");
165                        self.per_cache_recorded_bitmap_indices
166                            .push(bitmap_index as u32);
167                    }
168                    *count = count.wrapping_add(1);
169                }
170                #[cfg(not(feature = "cache"))]
171                let _ = cache;
172            }
173            NewBlock => {
174                self.set_new_loc(block_addr);
175            }
176        }
177        Ok(())
178    }
179
180    #[cfg(feature = "cache")]
181    #[expect(clippy::cast_possible_truncation)]
182    fn cache_prev_cached_key(&mut self, cached_key: Self::CachedKey) -> Result<(), Self::Error> {
183        let entries_range = cached_key.to_range();
184        // SAFETY: bitmap entries arena will never shrink
185        debug_assert!(
186            entries_range.end <= self.bitmap_entries_arena.len(),
187            "Unexpected OOB"
188        );
189        let bitmap_entries = unsafe { self.bitmap_entries_arena.get_unchecked(entries_range) };
190        for bitmap_entry in bitmap_entries {
191            let bitmap_index = bitmap_entry.bitmap_index();
192            // SAFETY: bitmap index is caculated by modulo
193            debug_assert!(bitmap_index < self.per_cache_bitmap.len(), "Unexpected OOB");
194            let count = unsafe { self.per_cache_bitmap.get_unchecked_mut(bitmap_index) };
195            if *count == 0 {
196                debug_assert!(u32::try_from(bitmap_index).is_ok(), "Bitmap size too large");
197                self.per_cache_recorded_bitmap_indices
198                    .push(bitmap_index as u32);
199            }
200            *count = count.wrapping_add(bitmap_entry.bitmap_count());
201        }
202
203        Ok(())
204    }
205
206    #[cfg(feature = "cache")]
207    fn clear_current_cache(&mut self) -> Result<(), Self::Error> {
208        for bitmap_index in self.per_cache_recorded_bitmap_indices.drain(..) {
209            let bitmap_index = bitmap_index as usize;
210            // SAFETY: bitmap index is caculated by modulo
211            debug_assert!(bitmap_index < self.per_cache_bitmap.len(), "Unexpected OOB");
212            let count = unsafe { self.per_cache_bitmap.get_unchecked_mut(bitmap_index) };
213            *count = 0;
214        }
215        Ok(())
216    }
217
218    #[cfg(feature = "cache")]
219    #[expect(clippy::cast_possible_truncation)]
220    fn take_cache(&mut self) -> Result<Option<Self::CachedKey>, Self::Error> {
221        let start_index = self.bitmap_entries_arena.len();
222        for bitmap_index in self.per_cache_recorded_bitmap_indices.drain(..) {
223            let bitmap_index = bitmap_index as usize;
224            // SAFETY: bitmap index is caculated by modulo
225            debug_assert!(bitmap_index < self.per_cache_bitmap.len(), "Unexpected OOB");
226            let count = unsafe { self.per_cache_bitmap.get_unchecked_mut(bitmap_index) };
227
228            let bitmap_entry = CompactBitmapEntry::new(bitmap_index, *count);
229            self.bitmap_entries_arena.push(bitmap_entry);
230
231            *count = 0;
232        }
233        let end_index = self.bitmap_entries_arena.len();
234        if start_index == end_index {
235            // No new bitmap entries
236            return Ok(None);
237        }
238        // SAFETY: bitmap entries arena always have a dummy first element, so index will never be zero
239        debug_assert!(start_index > 0 && end_index > 0, "Unexpected!");
240        debug_assert!(
241            u32::try_from(start_index).is_ok() && u32::try_from(end_index).is_ok(),
242            "Too many bitmap entries!"
243        );
244        let start_index = unsafe { NonZero::new_unchecked(start_index as u32) };
245        let end_index = unsafe { NonZero::new_unchecked(end_index as u32) };
246
247        Ok(Some(PerCacheBitmapEntries {
248            start: start_index,
249            end: end_index,
250        }))
251    }
252
253    #[cfg(feature = "cache")]
254    fn on_reused_cache(
255        &mut self,
256        cached_key: &Self::CachedKey,
257        new_bb: u64,
258    ) -> Result<(), Self::Error> {
259        let entries_range = cached_key.to_range();
260        // SAFETY: bitmap entries arena will never shrink
261        debug_assert!(
262            entries_range.end <= self.bitmap_entries_arena.len(),
263            "Unexpected OOB"
264        );
265        let bitmap_entries = unsafe { self.bitmap_entries_arena.get_unchecked(entries_range) };
266        // FIXME: This loop should be unrolled, but there is a bug in LLVM: https://github.com/rust-lang/rust/issues/150647
267        for bitmap_entry in bitmap_entries {
268            let bitmap_index = bitmap_entry.bitmap_index();
269            debug_assert!(
270                bitmap_index < self.fuzzing_bitmap.as_ref().len(),
271                "Unexpected OOB"
272            );
273            let count = unsafe { self.fuzzing_bitmap.as_mut().get_unchecked_mut(bitmap_index) };
274            *count = count.wrapping_add(bitmap_entry.bitmap_count());
275        }
276        self.set_new_loc(new_bb);
277
278        Ok(())
279    }
280}
281
282/// Dummy bitmap entry used to make sure the index of [`bitmap_entries_arena`][FuzzBitmapControlFlowHandler::bitmap_entries_arena]
283/// will never be zero
284#[cfg(feature = "cache")]
285const DUMMY_BITMAP_ENTRY: CompactBitmapEntry = CompactBitmapEntry { value: 0 };
286
287/// Compact representation of a (pos, count) pair used for fuzzing bitmap
288#[cfg(feature = "cache")]
289#[derive(Clone, Copy)]
290struct CompactBitmapEntry {
291    /// The actual value.
292    ///
293    /// The upper 24 bits is the pos, and the lower 8 bits is the count
294    value: u32,
295}
296
297#[cfg(feature = "cache")]
298impl CompactBitmapEntry {
299    /// Create a new compact bitmap entry. The bitmap index should never greater
300    /// than `0x00FF_FFFF`.
301    #[expect(clippy::cast_possible_truncation)]
302    fn new(bitmap_index: usize, bitmap_count: u8) -> Self {
303        debug_assert!(bitmap_index <= 0x00FF_FFFF, "Bitmap size too large");
304        let bitmap_index = (bitmap_index as u32) << 8;
305        Self {
306            value: bitmap_index | (bitmap_count as u32),
307        }
308    }
309
310    /// Get the bitmap index
311    fn bitmap_index(self) -> usize {
312        (self.value >> 8) as usize
313    }
314
315    /// Get the bitmap count
316    #[expect(clippy::cast_possible_truncation)]
317    fn bitmap_count(self) -> u8 {
318        self.value as u8
319    }
320}
321
322/// Cached key for [`FuzzBitmapControlFlowHandler`]
323///
324/// The cached key is a range into the [`bitmap_entries_arena`][FuzzBitmapControlFlowHandler::bitmap_entries_arena].
325#[cfg(feature = "cache")]
326#[doc(hidden)]
327#[derive(Clone, Copy)]
328pub struct PerCacheBitmapEntries {
329    /// Start of range, inclusive
330    start: NonZero<u32>,
331    /// End of range, exclusive
332    end: NonZero<u32>,
333}
334
335#[cfg(feature = "cache")]
336impl PerCacheBitmapEntries {
337    /// Get the range of bitmap entries
338    fn to_range(self) -> Range<usize> {
339        (self.start.get() as usize)..(self.end.get() as usize)
340    }
341}
iptr_edge_analyzer/control_flow_handler/fuzz_bitmap.rs

iptr_edge_analyzer/control_flow_handler/
fuzz_bitmap.rs
