iptr_edge_analyzer/control_flow_handler/mod.rs
1//! This module contains the core definition of [`HandleControlFlow`] trait,
2//! and several implementors like [`FuzzBitmapControlFlowHandler`][fuzz_bitmap::FuzzBitmapControlFlowHandler].
3//!
4//! This module also contains a `LogControlFlowHandler` in the path
5//! `iptr_edge_analyzer::control_flow_handler::log::LogControlFlowHandler`. However, due
6//! to some limitations of rustdoc, this struct cannot be displayed in the documentation.
7//! This struct is only accessible if `log_control_flow_handler` feature is on and `cache`
8//! feature is off.
9
10use derive_more::Display;
11
12pub mod combined;
13#[cfg(feature = "fuzz_bitmap")]
14pub mod fuzz_bitmap;
15#[cfg(all(not(feature = "cache"), feature = "log_control_flow_handler"))]
16pub mod log;
17
18/// Kind of control flow transitions
19#[derive(Debug, Display, Clone, Copy)]
20pub enum ControlFlowTransitionKind {
21 /// Conditional Jcc
22 ConditionalBranch,
23 /// Direct JMP
24 DirectJump,
25 /// Direct CALL
26 DirectCall,
27 /// Indirect transition, including `RET`.
28 Indirect,
29 /// New block
30 ///
31 /// Basic blocks that cannot be categorized into
32 /// other reasons. For example, blocks after page fault is resolved.
33 NewBlock,
34}
35
36/// Control flow handler used for [`EdgeAnalyzer`][crate::EdgeAnalyzer]
37///
38/// There are several implementors provided in this crate, such as
39/// [`FuzzBitmapControlFlowHandler`][fuzz_bitmap::FuzzBitmapControlFlowHandler].
40///
41/// # Non-cache mode
42///
43/// For non-cache mode, the usage of this trait is very simple: there are only two methods.
44/// [`at_decode_begin`][HandleControlFlow::at_decode_begin] is invoked at decode begin, and
45/// [`on_new_block`][HandleControlFlow::on_new_block] is invoked every time a basic block is
46/// encountered.
47///
48/// # Cache-mode
49///
50/// The overall workflow when using this trait is like:
51/// 1. Creating a new handler.
52/// 2. Clear cache by [`clear_current_cache`][HandleControlFlow::clear_current_cache]
53/// 3. When a new basic block is met, call [`on_new_block`][HandleControlFlow::on_new_block].
54/// This function should always deal with the impact, and deal with the cache depending on the
55/// `cache` parameter.
56/// 4. When a previous cache is met, call [`on_reused_cache`][HandleControlFlow::on_reused_cache].
57/// This function should only deal with the impact.
58/// 5. Optionally merge caches by [`cache_prev_cached_key`][HandleControlFlow::cache_prev_cached_key].
59/// 6. Collect cache by [`take_cache`][HandleControlFlow::take_cache].
60///
61/// In the documentation of this trait, there are two terms: "impact" and "cache". Let's take
62/// some examples. For fuzzing, "impact" means modification of fuzzing bitmap, and "cache"
63/// means modification of internal cached information. For logging, "impact" means logging the
64/// basic block transition, and "cache" also means the modification of internal cached information.
65///
66/// It should be noted that, the correctness of implementation of cache-mode [`HandleControlFlow`]
67/// will have no impact on the correctness of cache mechanism used by [`EdgeAnalyzer`][crate::EdgeAnalyzer].
68/// Instead, since [`on_new_block`][HandleControlFlow::on_new_block] is only invoked when a non-cache
69/// basic block is encountered, if you incorrectly implemented this trait, you may not accurately
70/// get all basic blocks.
71pub trait HandleControlFlow {
72 /// Error of control flow handler
73 type Error: std::error::Error;
74 /// Cached key returned by [`take_cache`][HandleControlFlow::take_cache].
75 ///
76 /// This can be used by the edge analyzer to tell the control flow handler
77 /// a previous TNT sequence has been met again and the cache is reused instead
78 /// of re-parsing all TNT bits.
79 #[cfg(feature = "cache")]
80 type CachedKey: Clone;
81
82 /// Callback at begin of decoding.
83 ///
84 /// This is useful when using the same handler to process multiple Intel PT
85 /// traces
86 fn at_decode_begin(&mut self) -> Result<(), Self::Error>;
87
88 /// Callback when a new basic block is met.
89 ///
90 /// For non-cache mode, this function is always invoked no matter whether this
91 /// basic block has been encountered before (i.e., this is not a "unique" block).
92 /// For cache mode, this function is only invoked when a non-cached basic block
93 /// is encountered.
94 ///
95 /// The new block's address is `block_addr`, and the reason for getting
96 /// into this block is in `transition_kind`. `cache` is only used in cache mode,
97 /// which indicates whether this block transition should be taken into cache
98 /// by the implementor, which is used as an optimizing hint. If `cache` is false,
99 /// this means this transition will never be folded into cache.
100 /// No matter `cache` is true or false, this function should always deal with
101 /// the impact of new block.
102 ///
103 /// When conducting caching, it should be extremely important, that
104 /// the cached state should always be consistent with `block_addr`.
105 ///
106 /// Suggest marking `#[inline]` on the implementation
107 fn on_new_block(
108 &mut self,
109 block_addr: u64,
110 transition_kind: ControlFlowTransitionKind,
111 cache: bool,
112 ) -> Result<(), Self::Error>;
113
114 /// Merge a previous cached key into cache
115 ///
116 /// When analyzing TNT packets, the cache manager maintains two kinds of cache: 8bits cache
117 /// and 32bits cache. As a result, when creating 32bits caches, we need to merge four 8bits
118 /// caches into one 32bits cache, and this function serves the purpose.
119 ///
120 /// It should be noted that although merged, the previous cache should still be kept.
121 ///
122 /// This function only deals with the caching thing. The previously cached information should
123 /// not have impact in this function. For dealing with impacts, see [`on_reused_cache`][HandleControlFlow::on_reused_cache].
124 #[cfg(feature = "cache")]
125 fn cache_prev_cached_key(&mut self, cached_key: Self::CachedKey) -> Result<(), Self::Error>;
126
127 /// Collect all currently cached information and generate a cached key. This could clear
128 /// the cache depending on your implementing logic.
129 ///
130 /// If there is no cache, this function should return `Ok(None)`.
131 #[cfg(feature = "cache")]
132 fn take_cache(&mut self) -> Result<Option<Self::CachedKey>, Self::Error>;
133
134 /// Clear the cache.
135 ///
136 /// This is NOT clearing all cached information. Instead, this is
137 /// to clear current temporary cache.
138 #[cfg(feature = "cache")]
139 fn clear_current_cache(&mut self) -> Result<(), Self::Error>;
140
141 /// Callback when a given cached key is being reused.
142 ///
143 /// `new_bb` is the next basic block address after the cached key is applied.
144 ///
145 /// This function only deals ith the impact of cached key, and should not add new caches.
146 /// For adding new caches, see [`cache_prev_cached_key`][HandleControlFlow::cache_prev_cached_key].
147 #[cfg(feature = "cache")]
148 fn on_reused_cache(
149 &mut self,
150 cached_key: &Self::CachedKey,
151 new_bb: u64,
152 ) -> Result<(), Self::Error>;
153}