[WARP] Allow multi-round matching to resolve matching function dependencies

Fixes some possible nondeterminism due to parallelized matching

2 files changed, 132 insertions, 69 deletions

diff --git a/plugins/warp/src/matcher.rs b/plugins/warp/src/matcher.rs
index 48adee3a..f96128e5 100644
--- a/plugins/warp/src/matcher.rs
+++ b/plugins/warp/src/matcher.rs
@@ -312,6 +312,10 @@ pub struct MatcherSettings {
     ///
     /// This is set to [MatcherSettings::MAXIMUM_POSSIBLE_FUNCTIONS_DEFAULT] by default.
     pub maximum_possible_functions: Option<u64>,
+    /// The maximum number of matching rounds to run, consecutive rounds are ran until no new matched functions are found.
+    ///
+    /// This is set to [MatcherSettings::MAXIMUM_MATCHING_ROUNDS_DEFAULT] by default.
+    pub maximum_matching_rounds: Option<u64>,
 }
 
 impl MatcherSettings {
@@ -342,6 +346,9 @@ impl MatcherSettings {
     pub const MAXIMUM_POSSIBLE_FUNCTIONS_SETTING: &'static str =
         "warp.matcher.maximumPossibleFunctions";
     pub const MAXIMUM_POSSIBLE_FUNCTIONS_DEFAULT: u64 = 1000;
+    pub const MAXIMUM_MATCHING_ROUNDS_SETTING: &'static str = "warp.matcher.maximumMatchingRounds";
+    pub const MAXIMUM_MATCHING_ROUNDS_DEFAULT: u64 = 0;
+
     /// Populates the [MatcherSettings] to the current Binary Ninja settings instance.
     ///
     /// Call this once when you initialize so that the settings exist.
@@ -426,6 +433,18 @@ impl MatcherSettings {
             Self::MAXIMUM_POSSIBLE_FUNCTIONS_SETTING,
             &maximum_possible_functions_props.to_string(),
         );
+
+        let maximum_matching_rounds_props = json!({
+            "title" : "Maximum Matching Rounds",
+            "type" : "number",
+            "default" : Self::MAXIMUM_MATCHING_ROUNDS_DEFAULT,
+            "description" : "The maximum number of matching rounds to run, consecutive rounds are ran until no new matched functions are found. A value of 0 will disable this check.",
+            "ignore" : [],
+        });
+        bn_settings.register_setting_json(
+            Self::MAXIMUM_MATCHING_ROUNDS_SETTING,
+            &maximum_matching_rounds_props.to_string(),
+        );
     }
 
     /// Retrieve matcher settings from [`BNSettings`].
@@ -463,6 +482,14 @@ impl MatcherSettings {
                 len => settings.maximum_possible_functions = Some(len),
             }
         }
+        if bn_settings.contains(Self::MAXIMUM_MATCHING_ROUNDS_SETTING) {
+            match bn_settings
+                .get_integer_with_opts(Self::MAXIMUM_MATCHING_ROUNDS_SETTING, query_opts)
+            {
+                0 => settings.maximum_matching_rounds = None,
+                len => settings.maximum_matching_rounds = Some(len),
+            }
+        }
         settings
     }
 }
@@ -477,6 +504,7 @@ impl Default for MatcherSettings {
             trivial_function_adjacent_allowed:
                 MatcherSettings::TRIVIAL_FUNCTION_ADJACENT_ALLOWED_DEFAULT,
             maximum_possible_functions: Some(MatcherSettings::MAXIMUM_POSSIBLE_FUNCTIONS_DEFAULT),
+            maximum_matching_rounds: Some(MatcherSettings::MAXIMUM_MATCHING_ROUNDS_DEFAULT),
         }
     }
 }
diff --git a/plugins/warp/src/plugin/workflow.rs b/plugins/warp/src/plugin/workflow.rs
index da988581..e5ce9b76 100644
--- a/plugins/warp/src/plugin/workflow.rs
+++ b/plugins/warp/src/plugin/workflow.rs
@@ -3,6 +3,7 @@ use crate::cache::{
     cached_function_guid, insert_cached_function_match, try_cached_function_guid,
     try_cached_function_match,
 };
+use crate::container::{Container, SourceId};
 use crate::convert::{platform_to_target, to_bn_symbol_at_address, to_bn_type};
 use crate::matcher::{Matcher, MatcherSettings};
 use crate::plugin::settings::PluginSettings;
@@ -15,14 +16,13 @@ use binaryninja::function::Function as BNFunction;
 use binaryninja::rc::Ref as BNRef;
 use binaryninja::settings::{QueryOptions, Settings};
 use binaryninja::workflow::{activity, Activity, AnalysisContext, Workflow, WorkflowBuilder};
+use dashmap::DashSet;
 use itertools::Itertools;
 use rayon::iter::IntoParallelIterator;
 use rayon::iter::ParallelIterator;
 use std::cmp::Ordering;
 use std::collections::HashMap;
-use std::sync::atomic::AtomicUsize;
-use std::sync::{Arc, Mutex};
-use std::time::{Duration, Instant};
+use std::time::Instant;
 use warp::r#type::class::function::{Location, RegisterLocation, StackLocation};
 use warp::signature::function::{Function, FunctionGUID};
 use warp::target::Target;
@@ -120,6 +120,23 @@ pub fn run_matcher(view: &BinaryView) {
     let _ = get_warp_ignore_tag_type(view);
     view.file().forget_undo_actions(&undo_id);
 
+    let filter_functions = |functions: &mut Vec<Function>| {
+        // We sort primarily by symbol, then by type, so we can deduplicate in-place.
+        functions.sort_unstable_by(|a, b| match a.symbol.cmp(&b.symbol) {
+            Ordering::Equal => match (&a.ty, &b.ty) {
+                (None, None) => Ordering::Equal,
+                (None, Some(_)) => Ordering::Less,
+                (Some(_), None) => Ordering::Greater,
+                // TODO: We still need to order the types, probably cant do this in place.
+                // TODO: Once Type can be ordered, we can remove this entire explicit match stmt.
+                (Some(_), Some(_)) => Ordering::Equal,
+            },
+            other => other,
+        });
+        // This removes consecutive duplicates efficiently
+        functions.dedup_by(|a, b| a.symbol == b.symbol && a.ty == b.ty);
+    };
+
     // Then we want to actually find matching functions.
     let background_task = BackgroundTask::new("Matching on WARP functions...", true);
     let start = Instant::now();
@@ -135,88 +152,106 @@ pub fn run_matcher(view: &BinaryView) {
         return;
     };
 
-    // TODO: Target gets cloned a lot.
-    // TODO: Containers might both match on the same function. What should we do?
-    let matched_count = AtomicUsize::new(0);
-    for_cached_containers(|container| {
-        if background_task.is_cancelled() {
+    let matcher_results: DashSet<u64> = DashSet::new();
+    let match_for_guid = |target, container: &dyn Container, sources: Vec<SourceId>, guid| {
+        let mut matched_functions: Vec<Function> = sources
+            .iter()
+            .flat_map(|source| {
+                container
+                    .functions_with_guid(target, source, &guid)
+                    .unwrap_or_default()
+            })
+            .collect();
+
+        // 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)
+        {
+            log::warn!(
+                "Skipping {}, too many possible functions: {}",
+                guid,
+                matched_functions.len()
+            );
             return;
         }
 
-        for (target, guids) in &function_set.guids_by_target {
-            let function_guid_with_sources = container
-                .sources_with_function_guids(target, guids)
-                .unwrap_or_default();
+        // Filter out duplicate functions for matching.
+        filter_functions(&mut matched_functions);
 
-            function_guid_with_sources
-                .into_par_iter()
-                .for_each(|(guid, sources)| {
-                    let mut matched_functions: Vec<Function> = sources
-                        .iter()
-                        .flat_map(|source| {
-                            container
-                                .functions_with_guid(target, source, &guid)
-                                .unwrap_or_default()
-                        })
-                        .collect();
+        let functions = function_set
+            .functions_by_target_and_guid
+            .get(&(guid, target.clone()))
+            .expect("Function guid not found");
 
-                    // We sort primarily by symbol, then by type, so we can deduplicate in-place.
-                    matched_functions.sort_unstable_by(|a, b| match a.symbol.cmp(&b.symbol) {
-                        Ordering::Equal => match (&a.ty, &b.ty) {
-                            (None, None) => Ordering::Equal,
-                            (None, Some(_)) => Ordering::Less,
-                            (Some(_), None) => Ordering::Greater,
-                            // TODO: We still need to order the types, probably cant do this in place.
-                            // TODO: Once Type can be ordered, we can remove this entire explicit match stmt.
-                            (Some(_), Some(_)) => Ordering::Equal,
-                        },
-                        other => other,
-                    });
-                    // This removes consecutive duplicates efficiently
-                    matched_functions.dedup_by(|a, b| a.symbol == b.symbol && a.ty == b.ty);
+        for function in functions {
+            // Match on all the possible functions
+            if let Some(matched_function) =
+                matcher.match_function_from_constraints(function, &matched_functions)
+            {
+                // Because we can do multiple rounds of matching at once, we only want to insert a function
+                // match if we have not already done so in a previous round.
+                // TODO: What if the new round changes the matched function metadata? Unlikely but possible.
+                if matcher_results.insert(function.start()) {
+                    // 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));
+                }
+            }
+        }
+    };
 
-                    // 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)
-                    {
-                        log::warn!(
-                            "Skipping {}, too many possible functions: {}",
-                            guid,
-                            matched_functions.len()
-                        );
-                        return;
-                    }
+    // NOTE: Because matching can depend on other functions to have matched, we will run multiple
+    // rounds of matching until it stabilizes (e.g. no more newly matched functions), there are other
+    // ways to have the same behavior that may take less time, such as a work list, and pushing callers
+    // back into the work list on matches of a function, on top of that you could order the functions
+    // matched bottom up, with a reverse post order sort.
+
+    // TODO: Target gets cloned a lot.
+    // TODO: Containers might both match on the same function. What should we do?
+    let maximum_rounds = matcher.settings.maximum_matching_rounds.unwrap_or(100);
+    let mut final_matched_round = 1;
+    for matched_round in 1..=maximum_rounds {
+        let bg_task_text = format!("Matching on WARP functions... ({} rounds)", matched_round);
+        background_task.set_progress_text(&bg_task_text);
+        let matched_count_before = matcher_results.len();
 
-                    let functions = function_set
-                        .functions_by_target_and_guid
-                        .get(&(guid, target.clone()))
-                        .expect("Function guid not found");
+        for_cached_containers(|container| {
+            if background_task.is_cancelled() {
+                return;
+            }
 
-                    for function in functions {
-                        // Match on all the possible functions
-                        if let Some(matched_function) =
-                            matcher.match_function_from_constraints(function, &matched_functions)
-                        {
-                            matched_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-                            // 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));
-                        }
-                    }
-                });
+            for (target, guids) in &function_set.guids_by_target {
+                let function_guid_with_sources = container
+                    .sources_with_function_guids(target, guids)
+                    .unwrap_or_default();
+
+                function_guid_with_sources
+                    .into_par_iter()
+                    .for_each(|(guid, sources)| {
+                        match_for_guid(target, container, sources, guid);
+                    });
+            }
+        });
+
+        final_matched_round = matched_round;
+        // If the number of matches did not increase we can stop matching.
+        let matched_count_after = matcher_results.len();
+        if matched_count_after == 0 || matched_count_after == matched_count_before {
+            break;
         }
-    });
+    }
 
     if background_task.is_cancelled() {
         log::info!("Matcher was cancelled by user, you may run it again by running the 'Run Matcher' command.");
     }
 
     log::info!(
-        "Function matching took {:.3} seconds and matched {} functions",
+        "Function matching took {:.3} seconds and matched {} functions after {} rounds",
         start.elapsed().as_secs_f64(),
-        matched_count.load(std::sync::atomic::Ordering::Relaxed)
+        matcher_results.len(),
+        final_matched_round
     );
     background_task.finish();