[WARP] Improve matcher speed and add fast fail for pathological cases

Improves the speed by moving chunks of functions into worker threads, because of how the functions and possible functions are gathered we have many locations to insert fast fails, which is also partially addressed by this commit (see `maximum_possible_functions`).

2 files changed, 80 insertions, 22 deletions

diff --git a/plugins/warp/src/matcher.rs b/plugins/warp/src/matcher.rs
index c6ab7f9f..52a5b04f 100644
--- a/plugins/warp/src/matcher.rs
+++ b/plugins/warp/src/matcher.rs
@@ -47,6 +47,18 @@ impl Matcher {
             return None;
         }
 
+        // The number of possible functions is too high, skip.
+        // This can happen if the function is extremely common, in cases like that we are already unlikely to match.
+        // It is unfortunate that we have to do this, but it is the best we can do. In the future we
+        // may find a way to chunk up the possible functions and only match on a subset of them.
+        if self
+            .settings
+            .maximum_possible_functions
+            .is_some_and(|max| max < matched_functions.len() as u64)
+        {
+            return None;
+        }
+
         // If we have a single possible match than that must be our function.
         // We must also not be a trivial function, as those will likely be artifacts of an incomplete dataset
         if matched_functions.len() == 1 && !is_function_trivial {
@@ -296,6 +308,10 @@ pub struct MatcherSettings {
     ///
     /// This is set to [MatcherSettings::TRIVIAL_FUNCTION_ADJACENT_ALLOWED_DEFAULT] by default.
     pub trivial_function_adjacent_allowed: bool,
+    /// The maximum number of WARP functions that can be used to match a Binary Ninja function.
+    ///
+    /// This is set to [MatcherSettings::MAXIMUM_POSSIBLE_FUNCTIONS_DEFAULT] by default.
+    pub maximum_possible_functions: Option<u64>,
 }
 
 impl MatcherSettings {
@@ -311,6 +327,9 @@ impl MatcherSettings {
     pub const TRIVIAL_FUNCTION_ADJACENT_ALLOWED_DEFAULT: bool = false;
     pub const TRIVIAL_FUNCTION_ADJACENT_ALLOWED_SETTING: &'static str =
         "analysis.warp.trivialFunctionAdjacentAllowed";
+    pub const MAXIMUM_POSSIBLE_FUNCTIONS_SETTING: &'static str =
+        "analysis.warp.maximumPossibleFunctions";
+    pub const MAXIMUM_POSSIBLE_FUNCTIONS_DEFAULT: u64 = 1000;
 
     /// Populates the [MatcherSettings] to the current Binary Ninja settings instance.
     ///
@@ -378,6 +397,18 @@ impl MatcherSettings {
             Self::TRIVIAL_FUNCTION_ADJACENT_ALLOWED_SETTING,
             &trivial_function_adjacent_allowed_props.to_string(),
         );
+
+        let maximum_possible_functions_props = json!({
+            "title" : "Maximum Possible Functions",
+            "type" : "number",
+            "default" : Self::MAXIMUM_POSSIBLE_FUNCTIONS_DEFAULT,
+            "description" : "When matching any function that has a list of possible functions greater than this number will be skipped. A value of 0 will disable this check.",
+            "ignore" : []
+        });
+        bn_settings.register_setting_json(
+            Self::MAXIMUM_POSSIBLE_FUNCTIONS_SETTING,
+            &maximum_possible_functions_props.to_string(),
+        );
     }
 
     /// Retrieve matcher settings from [`BNSettings`].
@@ -407,6 +438,14 @@ impl MatcherSettings {
             settings.trivial_function_adjacent_allowed = bn_settings
                 .get_bool_with_opts(Self::TRIVIAL_FUNCTION_ADJACENT_ALLOWED_SETTING, query_opts);
         }
+        if bn_settings.contains(Self::MAXIMUM_POSSIBLE_FUNCTIONS_SETTING) {
+            match bn_settings
+                .get_integer_with_opts(Self::MAXIMUM_POSSIBLE_FUNCTIONS_SETTING, query_opts)
+            {
+                0 => settings.maximum_possible_functions = None,
+                len => settings.maximum_possible_functions = Some(len),
+            }
+        }
         settings
     }
 }
@@ -420,6 +459,7 @@ impl Default for MatcherSettings {
             minimum_matched_constraints: MatcherSettings::MINIMUM_MATCHED_CONSTRAINTS_DEFAULT,
             trivial_function_adjacent_allowed:
                 MatcherSettings::TRIVIAL_FUNCTION_ADJACENT_ALLOWED_DEFAULT,
+            maximum_possible_functions: Some(MatcherSettings::MAXIMUM_POSSIBLE_FUNCTIONS_DEFAULT),
         }
     }
 }
diff --git a/plugins/warp/src/plugin/workflow.rs b/plugins/warp/src/plugin/workflow.rs
index 8da69342..47f4b2a7 100644
--- a/plugins/warp/src/plugin/workflow.rs
+++ b/plugins/warp/src/plugin/workflow.rs
@@ -13,6 +13,8 @@ use binaryninja::command::Command;
 use binaryninja::settings::{QueryOptions, Settings};
 use binaryninja::workflow::{activity, Activity, AnalysisContext, Workflow, WorkflowBuilder};
 use itertools::Itertools;
+use rayon::iter::IntoParallelIterator;
+use rayon::iter::ParallelIterator;
 use std::collections::HashMap;
 use std::time::Instant;
 use warp::r#type::class::function::{Location, RegisterLocation, StackLocation};
@@ -112,31 +114,47 @@ pub fn run_matcher(view: &BinaryView) {
                 .sources_with_function_guids(target, guids)
                 .unwrap_or_default();
 
-            for (guid, sources) in &function_guid_with_sources {
-                let matched_functions: Vec<Function> = sources
-                    .iter()
-                    .flat_map(|source| {
-                        container
-                            .functions_with_guid(target, source, guid)
-                            .unwrap_or_default()
-                    })
-                    .collect();
+            function_guid_with_sources
+                .into_par_iter()
+                .for_each(|(guid, sources)| {
+                    let matched_functions: Vec<Function> = sources
+                        .iter()
+                        .flat_map(|source| {
+                            container
+                                .functions_with_guid(target, source, &guid)
+                                .unwrap_or_default()
+                        })
+                        .collect();
 
-                let functions = functions_by_target_and_guid
-                    .get(&(*guid, target.clone()))
-                    .expect("Function guid not found");
-
-                for function in functions {
-                    // Match on all the possible functions
-                    if let Some(matched_function) =
-                        matcher.match_function_from_constraints(function, &matched_functions)
+                    // NOTE: See the comment in `match_function_from_constraints` about this fast fail.
+                    if matcher
+                        .settings
+                        .maximum_possible_functions
+                        .is_some_and(|max| max < matched_functions.len() as u64)
                     {
-                        // We were able to find a match, add it to the match cache and then mark the function
-                        // as requiring updates; this is so that we know about it in the applier activity.
-                        insert_cached_function_match(function, Some(matched_function.clone()));
+                        log::warn!(
+                            "Skipping {}, too many possible functions: {}",
+                            guid,
+                            matched_functions.len()
+                        );
+                        return;
                     }
-                }
-            }
+
+                    let functions = functions_by_target_and_guid
+                        .get(&(guid, target.clone()))
+                        .expect("Function guid not found");
+
+                    for function in functions {
+                        // Match on all the possible functions
+                        if let Some(matched_function) =
+                            matcher.match_function_from_constraints(function, &matched_functions)
+                        {
+                            // We were able to find a match, add it to the match cache and then mark the function
+                            // as requiring updates; this is so that we know about it in the applier activity.
+                            insert_cached_function_match(function, Some(matched_function.clone()));
+                        }
+                    }
+                });
         }
     });