WARP: Allow skipping trivial function constraints

3 files changed, 54 insertions, 8 deletions

diff --git a/plugins/warp/src/cache.rs b/plugins/warp/src/cache.rs
index 33d70548..652ddc39 100644
--- a/plugins/warp/src/cache.rs
+++ b/plugins/warp/src/cache.rs
@@ -93,15 +93,21 @@ pub fn cached_call_site_constraints(function: &BNFunction) -> HashSet<FunctionCo
     }
 }
 
-pub fn cached_adjacency_constraints(function: &BNFunction) -> HashSet<FunctionConstraint> {
+pub fn cached_adjacency_constraints<F>(
+    function: &BNFunction,
+    filter: F,
+) -> HashSet<FunctionConstraint>
+where
+    F: Fn(&BNFunction) -> bool,
+{
     let view = function.view();
     let view_id = ViewID::from(view);
     let guid_cache = GUID_CACHE.get_or_init(Default::default);
     match guid_cache.get(&view_id) {
-        Some(cache) => cache.adjacency_constraints(function),
+        Some(cache) => cache.adjacency_constraints(function, filter),
         None => {
             let cache = GUIDCache::default();
-            let constraints = cache.adjacency_constraints(function);
+            let constraints = cache.adjacency_constraints(function, filter);
             guid_cache.insert(view_id, cache);
             constraints
         }
@@ -226,6 +232,7 @@ impl GUIDCache {
                         if cs_ref_func_id != func_id {
                             let call_site_offset: i64 =
                                 call_site.address.wrapping_sub(func_start) as i64;
+                            // TODO: If the function is thunk we should also insert the called function.
                             constraints
                                 .insert(self.function_constraint(&cs_ref_func, call_site_offset));
                         }
@@ -249,7 +256,14 @@ impl GUIDCache {
         constraints
     }
 
-    pub fn adjacency_constraints(&self, function: &BNFunction) -> HashSet<FunctionConstraint> {
+    pub fn adjacency_constraints<F>(
+        &self,
+        function: &BNFunction,
+        filter: F,
+    ) -> HashSet<FunctionConstraint>
+    where
+        F: Fn(&BNFunction) -> bool,
+    {
         let view = function.view();
         let func_id = FunctionID::from(function);
         let func_start = function.start();
@@ -259,7 +273,7 @@ impl GUIDCache {
             // NOTE: We could potentially have dozens of functions all at the same start address.
             for curr_func in &view.functions_at(func_start_addr) {
                 let curr_func_id = FunctionID::from(curr_func.as_ref());
-                if curr_func_id != func_id {
+                if curr_func_id != func_id && filter(curr_func.as_ref()) {
                     // NOTE: For this to work the GUID has to have already been cached. If not it will just be the symbol.
                     // Function adjacent to another function, constrain on the pattern.
                     let curr_addr_offset = (func_start_addr as i64) - func_start as i64;
@@ -350,7 +364,8 @@ impl TypeRefCache {
             Some(cache) => cache.to_owned(),
             None => match type_ref.target(view) {
                 Some(raw_ty) => {
-                    let computed_ty = ComputedType::new(from_bn_type_internal(view, visited_refs, &raw_ty, 255));
+                    let computed_ty =
+                        ComputedType::new(from_bn_type_internal(view, visited_refs, &raw_ty, 255));
                     self.cache
                         .entry(ntr_id)
                         .insert(Some(computed_ty))
diff --git a/plugins/warp/src/lib.rs b/plugins/warp/src/lib.rs
index 2c50d978..e49e7605 100644
--- a/plugins/warp/src/lib.rs
+++ b/plugins/warp/src/lib.rs
@@ -35,7 +35,8 @@ pub fn build_function<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
         ty: from_bn_type(&func.view(), &bn_fn_ty, 255),
         constraints: FunctionConstraints {
             // NOTE: Adding adjacent only works if analysis is complete.
-            adjacent: cached_adjacency_constraints(func),
+            // NOTE: We do not filter out adjacent functions here.
+            adjacent: cached_adjacency_constraints(func, |_| true),
             call_sites: cached_call_site_constraints(func),
             // TODO: Add caller sites (when adjacent and call sites are minimal)
             // NOTE: Adding caller sites only works if analysis is complete.
diff --git a/plugins/warp/src/matcher.rs b/plugins/warp/src/matcher.rs
index 4dfc3e35..6ad3fdc7 100644
--- a/plugins/warp/src/matcher.rs
+++ b/plugins/warp/src/matcher.rs
@@ -251,8 +251,17 @@ impl Matcher {
         function: &BNFunction,
         matched_functions: &'a [Function],
     ) -> Option<&'a Function> {
+        // Filter out adjacent functions which are trivial, this helps avoid false positives.
+        // NOTE: If the user sets `trivial_function_adjacent_allowed` to true we will always match.
+        // TODO: Expand on this more later. We might want to match on adjacent functions smaller than this.
+        let adjacent_function_filter = |adj_func: &BNFunction| {
+            let adj_func_len = adj_func.highest_address() - adj_func.lowest_address();
+            adj_func_len > self.settings.trivial_function_len
+                || self.settings.trivial_function_adjacent_allowed
+        };
+
         let call_sites = cached_call_site_constraints(function);
-        let adjacent = cached_adjacency_constraints(function);
+        let adjacent = cached_adjacency_constraints(function, adjacent_function_filter);
 
         // "common" being the intersection between the observed and matched.
         fn find_highest_common_count<'a, F, T>(
@@ -370,6 +379,10 @@ pub struct MatcherSettings {
     ///
     /// This is set to [MatcherSettings::DEFAULT_TRIVIAL_FUNCTION_LEN] by default.
     pub minimum_matched_constraints: usize,
+    /// For a successful constrained function match the number of matches must be above this.
+    ///
+    /// This is set to [MatcherSettings::DEFAULT_TRIVIAL_FUNCTION_LEN] by default.
+    pub trivial_function_adjacent_allowed: bool,
 }
 
 impl MatcherSettings {
@@ -382,6 +395,9 @@ impl MatcherSettings {
     pub const MINIMUM_MATCHED_CONSTRAINTS_DEFAULT: usize = 1;
     pub const MINIMUM_MATCHED_CONSTRAINTS_SETTING: &'static str =
         "analysis.warp.minimumMatchedConstraints";
+    pub const TRIVIAL_FUNCTION_ADJACENT_ALLOWED_DEFAULT: bool = false;
+    pub const TRIVIAL_FUNCTION_ADJACENT_ALLOWED_SETTING: &'static str =
+        "analysis.warp.trivialFunctionAdjacentAllowed";
 
     /// Populates the [MatcherSettings] to the current Binary Ninja settings instance.
     ///
@@ -439,6 +455,18 @@ impl MatcherSettings {
             Self::MINIMUM_MATCHED_CONSTRAINTS_SETTING,
             minimum_matched_constraints_props.to_string(),
         );
+
+        let trivial_function_adjacent_allowed_props = json!({
+            "title" : "Trivial Function Adjacent Constraints Allowed",
+            "type" : "boolean",
+            "default" : Self::TRIVIAL_FUNCTION_ADJACENT_ALLOWED_DEFAULT,
+            "description" : "When function constraints are checked if this is enabled functions can match based off trivial adjacent functions.",
+            "ignore" : ["SettingsProjectScope", "SettingsResourceScope"]
+        });
+        bn_settings.register_setting_json(
+            Self::TRIVIAL_FUNCTION_ADJACENT_ALLOWED_SETTING,
+            trivial_function_adjacent_allowed_props.to_string(),
+        );
     }
 
     pub fn global() -> Self {
@@ -474,6 +502,8 @@ impl Default for MatcherSettings {
             minimum_function_len: MatcherSettings::MINIMUM_FUNCTION_LEN_DEFAULT,
             maximum_function_len: None,
             minimum_matched_constraints: MatcherSettings::MINIMUM_MATCHED_CONSTRAINTS_DEFAULT,
+            trivial_function_adjacent_allowed:
+                MatcherSettings::TRIVIAL_FUNCTION_ADJACENT_ALLOWED_DEFAULT,
         }
     }
 }