[WARP] Add optional `analysis.warp.fetcher` activity to run fetch networked functions as apart of initial analysis

This is off by default for the obvious reasons, we may also in the future want to expand upon this with a "matched view" such that the first step a user can take after loading is to validate what functions were matched and decide if they want to keep any of that data.

This fetch happens as apart of the initial analysis, such that a headless script will not return from the load function until it has finished fetching. We may want to add a configurable timeout, and/or run the fetch out of band (let load return before fetching finishes).

2 files changed, 128 insertions, 39 deletions

diff --git a/plugins/warp/src/cache/container.rs b/plugins/warp/src/cache/container.rs
index 59b31bad..0b478ed1 100644
--- a/plugins/warp/src/cache/container.rs
+++ b/plugins/warp/src/cache/container.rs
@@ -15,6 +15,15 @@ pub fn for_cached_containers(f: impl Fn(&dyn Container)) {
     }
 }
 
+pub fn for_cached_containers_mut(f: impl Fn(&mut dyn Container)) {
+    let containers_cache = CONTAINER_CACHE.get_or_init(Default::default);
+    for container in containers_cache.iter() {
+        if let Ok(mut guarded_container) = container.write() {
+            f(guarded_container.as_mut());
+        }
+    }
+}
+
 // TODO: The static lifetime here is a little wierd... (we need it to Box)
 pub fn add_cached_container(container: impl Container + 'static) {
     let containers_cache = CONTAINER_CACHE.get_or_init(Default::default);
diff --git a/plugins/warp/src/plugin/workflow.rs b/plugins/warp/src/plugin/workflow.rs
index 08d07c52..a1b0e202 100644
--- a/plugins/warp/src/plugin/workflow.rs
+++ b/plugins/warp/src/plugin/workflow.rs
@@ -1,15 +1,18 @@
-use crate::cache::container::for_cached_containers;
+use crate::cache::container::{for_cached_containers, for_cached_containers_mut};
 use crate::cache::{
     cached_function_guid, insert_cached_function_match, try_cached_function_guid,
     try_cached_function_match,
 };
 use crate::convert::{platform_to_target, to_bn_type};
 use crate::matcher::{Matcher, MatcherSettings};
+use crate::plugin::settings::PluginSettings;
 use crate::{get_warp_tag_type, relocatable_regions};
 use binaryninja::architecture::RegisterId;
 use binaryninja::background_task::BackgroundTask;
 use binaryninja::binary_view::{BinaryView, BinaryViewExt};
 use binaryninja::command::Command;
+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 itertools::Itertools;
@@ -47,6 +50,62 @@ impl Command for RunMatcher {
     }
 }
 
+/// This is a helper struct that contains all the functions that match on a given target.
+///
+/// We build this when matching and fetching so that we can relate different functions to another.
+#[derive(Debug, Clone, PartialEq, Eq, Default)]
+struct FunctionSet {
+    functions_by_target_and_guid: HashMap<(FunctionGUID, Target), Vec<BNRef<BNFunction>>>,
+    guids_by_target: HashMap<Target, Vec<FunctionGUID>>,
+}
+
+impl FunctionSet {
+    fn from_view(view: &BinaryView) -> Option<Self> {
+        let mut set = FunctionSet::default();
+
+        // TODO: Par iter this? Using dashmap
+        set.functions_by_target_and_guid = view
+            .functions()
+            .iter()
+            .filter_map(|f| {
+                let guid = try_cached_function_guid(&f)?;
+                let target = platform_to_target(&f.platform());
+                Some(((guid, target), f.to_owned()))
+            })
+            .into_group_map();
+
+        if set.functions_by_target_and_guid.is_empty() && !view.functions().is_empty() {
+            // The user is likely trying to run the matcher on a database before guids were automatically
+            // generated, we should alert them and ask if they would like to reanalyze.
+            // NOTE: We only alert if we actually have the GUID activity enabled.
+            if let Some(sample_function) = view.functions().iter().next() {
+                let function_workflow = sample_function
+                    .workflow()
+                    .expect("Function has no workflow");
+                if function_workflow.contains(GUID_ACTIVITY_NAME) {
+                    log::error!("No function guids in database, please reanalyze the database.");
+                } else {
+                    log::error!(
+                    "Activity '{}' is not in workflow '{}', create function guids manually to run matcher...",
+                    GUID_ACTIVITY_NAME,
+                    function_workflow.name()
+                )
+                }
+            }
+            return None;
+        }
+
+        // TODO: Par iter this? Using dashmap
+        set.guids_by_target = set
+            .functions_by_target_and_guid
+            .keys()
+            .map(|(guid, target)| (target.clone(), *guid))
+            .into_group_map();
+
+        Some(set)
+    }
+}
+
 pub fn run_matcher(view: &BinaryView) {
     // TODO: Create the tag type so we dont have UB in the apply function workflow.
     let undo_id = view.file().begin_undo_actions(false);
@@ -63,44 +122,10 @@ pub fn run_matcher(view: &BinaryView) {
     let matcher_settings = MatcherSettings::from_settings(&view_settings, &mut query_opts);
     let matcher = Matcher::new(matcher_settings);
 
-    // TODO: Par iter this? Using dashmap
-    let functions_by_target_and_guid: HashMap<(FunctionGUID, Target), Vec<_>> = view
-        .functions()
-        .iter()
-        .filter_map(|f| {
-            let guid = try_cached_function_guid(&f)?;
-            let target = platform_to_target(&f.platform());
-            Some(((guid, target), f.to_owned()))
-        })
-        .into_group_map();
-
-    if functions_by_target_and_guid.is_empty() && !view.functions().is_empty() {
-        // The user is likely trying to run the matcher on a database before guids were automatically
-        // generated, we should alert them and ask if they would like to reanalyze.
-        // NOTE: We only alert if we actually have the GUID activity enabled.
-        if let Some(sample_function) = view.functions().iter().next() {
-            let function_workflow = sample_function
-                .workflow()
-                .expect("Function has no workflow");
-            if function_workflow.contains(GUID_ACTIVITY_NAME) {
-                log::error!("No function guids in database, please reanalyze the database.");
-            } else {
-                log::error!(
-                    "Activity '{}' is not in workflow '{}', create function guids manually to run matcher...",
-                    GUID_ACTIVITY_NAME,
-                    function_workflow.name()
-                )
-            }
-        }
+    let Some(function_set) = FunctionSet::from_view(view) else {
         background_task.finish();
         return;
-    }
-
-    // TODO: Par iter this? Using dashmap
-    let guids_by_target: HashMap<Target, Vec<FunctionGUID>> = functions_by_target_and_guid
-        .keys()
-        .map(|(guid, target)| (target.clone(), *guid))
-        .into_group_map();
+    };
 
     // TODO: Target gets cloned a lot.
     // TODO: Containers might both match on the same function. What should we do?
@@ -109,7 +134,7 @@ pub fn run_matcher(view: &BinaryView) {
             return;
         }
 
-        for (target, guids) in &guids_by_target {
+        for (target, guids) in &function_set.guids_by_target {
             let function_guid_with_sources = container
                 .sources_with_function_guids(target, guids)
                 .unwrap_or_default();
@@ -140,7 +165,8 @@ pub fn run_matcher(view: &BinaryView) {
                         return;
                     }
 
-                    let functions = functions_by_target_and_guid
+                    let functions = function_set
+                        .functions_by_target_and_guid
                         .get(&(guid, target.clone()))
                         .expect("Function guid not found");
 
@@ -169,6 +195,46 @@ pub fn run_matcher(view: &BinaryView) {
     view.update_analysis();
 }
 
+pub fn run_fetcher(view: &BinaryView) {
+    let background_task = BackgroundTask::new("Fetching WARP functions...", true);
+    let start = Instant::now();
+
+    // Build matcher
+    let view_settings = Settings::new();
+    let mut query_opts = QueryOptions::new_with_view(view);
+    let plugin_settings = PluginSettings::from_settings(&view_settings, &mut query_opts);
+
+    let Some(function_set) = FunctionSet::from_view(view) else {
+        background_task.finish();
+        return;
+    };
+
+    for_cached_containers_mut(|container| {
+        if background_task.is_cancelled() {
+            return;
+        }
+
+        for (target, guids) in &function_set.guids_by_target {
+            for batch in guids.chunks(plugin_settings.fetch_batch_size) {
+                if background_task.is_cancelled() {
+                    break;
+                }
+                let _ =
+                    container.fetch_functions(target, &plugin_settings.allowed_source_tags, batch);
+            }
+        }
+    });
+
+    if background_task.is_cancelled() {
+        log::info!(
+            "Fetcher was cancelled by user, you may run it again by running the 'Fetch' command."
+        );
+    }
+
+    log::info!("Fetching took {:?}", start.elapsed());
+    background_task.finish();
+}
+
 pub fn insert_workflow() -> Result<(), ()> {
     // TODO: Note: because of symbol persistence function symbol is applied in `insert_cached_function_match`.
     // TODO: Comments are also applied there, they are "user" like, persisted and make undo actions.
@@ -232,6 +298,11 @@ pub fn insert_workflow() -> Result<(), ()> {
         run_matcher(&view);
     };
 
+    let fetcher_activity = |ctx: &AnalysisContext| {
+        let view = ctx.view();
+        run_fetcher(&view);
+    };
+
     let guid_activity = |ctx: &AnalysisContext| {
         let function = ctx.function();
         cached_function_guid(&function, || unsafe { ctx.lifted_il_function() });
@@ -269,6 +340,14 @@ pub fn insert_workflow() -> Result<(), ()> {
     // TODO: Remove this once the objectivec workflow is registered on the meta workflow.
     add_function_activities(Workflow::cloned("core.function.objectiveC"))?;
 
+    let fetcher_config = activity::Config::action(
+        "analysis.warp.fetcher",
+        "WARP Fetcher",
+        "This analysis step attempts to fetch WARP functions from network containers after the initial analysis is complete...",
+    )
+        .eligibility(activity::Eligibility::auto_with_default(false).run_once(true));
+    let fetcher_activity = Activity::new_with_action(&fetcher_config, fetcher_activity);
+
     let matcher_config = activity::Config::action(
         "analysis.warp.matcher",
         "WARP Matcher",
@@ -281,6 +360,7 @@ pub fn insert_workflow() -> Result<(), ()> {
     Workflow::cloned("core.module.metaAnalysis")
         .ok_or(())?
         .activity_before(&matcher_activity, "core.module.finishUpdate")?
+        .activity_before(&fetcher_activity, "analysis.warp.matcher")?
         .register()?;
 
     Ok(())