[WARP] Server-side constraint matching

Reduce networked functions by constraining on the returned set of functions on the server

6 files changed, 77 insertions, 28 deletions

diff --git a/plugins/warp/src/container.rs b/plugins/warp/src/container.rs
index 79c243b3..4feed24c 100644
--- a/plugins/warp/src/container.rs
+++ b/plugins/warp/src/container.rs
@@ -9,6 +9,7 @@ use thiserror::Error;
 use uuid::Uuid;
 use warp::r#type::guid::TypeGUID;
 use warp::r#type::{ComputedType, Type};
+use warp::signature::constraint::ConstraintGUID;
 use warp::signature::function::{Function, FunctionGUID};
 use warp::symbol::Symbol;
 use warp::target::Target;
@@ -295,11 +296,15 @@ pub trait Container: Send + Sync + Display + Debug {
     /// Typically, a container that resides only in memory has nothing to fetch, so the default implementation
     /// will do nothing. This function is blocking, so assume it will take a few seconds for a container
     /// that intends to fetch over the network.
+    ///
+    /// To constrain on the fetched functions, pass a list of [`ConstraintGUID`]s that will be
+    /// used to filter the fetched functions which do not contain at least one of the constraints.
     fn fetch_functions(
         &self,
         _target: &Target,
         _tags: &[SourceTag],
         _functions: &[FunctionGUID],
+        _constraints: &[ConstraintGUID],
     ) -> ContainerResult<()> {
         Ok(())
     }
diff --git a/plugins/warp/src/container/network.rs b/plugins/warp/src/container/network.rs
index 307b44fa..88192844 100644
--- a/plugins/warp/src/container/network.rs
+++ b/plugins/warp/src/container/network.rs
@@ -14,6 +14,7 @@ use warp::r#type::chunk::TypeChunk;
 use warp::r#type::guid::TypeGUID;
 use warp::r#type::{ComputedType, Type};
 use warp::signature::chunk::SignatureChunk;
+use warp::signature::constraint::ConstraintGUID;
 use warp::signature::function::{Function, FunctionGUID};
 use warp::target::Target;
 use warp::{WarpFile, WarpFileHeader};
@@ -45,7 +46,7 @@ pub struct NetworkContainer {
     /// NOTE: This is a [`DashMap`] purely for the sake of interior mutability as we do not wish to hold
     /// a write lock on the entire container while performing network operations.
     known_function_sources: DashMap<FunctionGUID, Vec<SourceId>>,
-    /// Populated when user adds function, this is used for writing back to the server.
+    /// Populated when the user adds a function, this is used for writing back to the server.
     added_chunks: HashMap<SourceId, Vec<Chunk<'static>>>,
     /// Populated when connecting to the server, this is used to determine which sources are writable.
     ///
@@ -165,18 +166,25 @@ impl NetworkContainer {
     /// Every request we store the returned objects on disk, this means that users will first
     /// query against the disk objects, then the server. This also means we need to cache functions f
     /// or which we have not received any functions for, as otherwise we would keep trying to query it.
-    pub fn pull_functions(&self, target: &Target, source: &SourceId, functions: &[FunctionGUID]) {
+    pub fn pull_functions(
+        &self,
+        target: &Target,
+        source: &SourceId,
+        functions: &[FunctionGUID],
+        constraints: &[ConstraintGUID],
+    ) {
         let target_id = self.get_target_id(target);
-        let file = match self
-            .client
-            .query_functions(target_id, Some(*source), functions)
-        {
-            Ok(file) => file,
-            Err(e) => {
-                tracing::error!("Failed to query functions: {}", e);
-                return;
-            }
-        };
+        let file =
+            match self
+                .client
+                .query_functions(target_id, Some(*source), functions, constraints)
+            {
+                Ok(file) => file,
+                Err(e) => {
+                    tracing::error!("Failed to query functions: {}", e);
+                    return;
+                }
+            };
 
         tracing::debug!("Got {} chunks from server", file.chunks.len());
         for chunk in &file.chunks {
@@ -396,16 +404,18 @@ impl Container for NetworkContainer {
         target: &Target,
         tags: &[SourceTag],
         functions: &[FunctionGUID],
+        constraints: &[ConstraintGUID],
     ) -> ContainerResult<()> {
         // NOTE: Blocking request to get the mapped function sources.
         let mapped_unseen_functions =
             self.get_unseen_functions_source(Some(&target), tags, functions);
 
+        // TODO: It would be nice to have a way to not have to pull through each source individually.
         // Actually get the function data for the unseen guids, we really only want to do this once per
         // session, anymore, and this is annoying!
         for (source, unseen_guids) in mapped_unseen_functions {
             // NOTE: Blocking request to get the function data in the container cache.
-            self.pull_functions(&target, &source, &unseen_guids);
+            self.pull_functions(&target, &source, &unseen_guids, constraints);
         }
 
         Ok(())
diff --git a/plugins/warp/src/container/network/client.rs b/plugins/warp/src/container/network/client.rs
index 0b772dcc..3d0af057 100644
--- a/plugins/warp/src/container/network/client.rs
+++ b/plugins/warp/src/container/network/client.rs
@@ -7,12 +7,13 @@ use base64::Engine;
 use binaryninja::download::DownloadProvider;
 use serde::Deserialize;
 use serde_json::json;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::str::FromStr;
 use uuid::Uuid;
 use warp::chunk::ChunkKind;
 use warp::r#type::guid::TypeGUID;
 use warp::r#type::{ComputedType, Type};
+use warp::signature::constraint::ConstraintGUID;
 use warp::signature::function::{Function, FunctionGUID};
 use warp::target::Target;
 use warp::WarpFile;
@@ -30,7 +31,8 @@ pub struct NetworkClient {
 impl NetworkClient {
     pub fn new(server_url: String, server_token: Option<String>) -> Self {
         // TODO: This might want to be kept for the request header?
-        let mut headers: Vec<(String, String)> = vec![];
+        let mut headers: Vec<(String, String)> =
+            vec![("Content-Encoding".to_string(), "gzip".to_string())];
         if let Some(token) = &server_token {
             headers.push(("authorization".to_string(), format!("Bearer {}", token)));
         }
@@ -214,13 +216,14 @@ impl NetworkClient {
         source: Option<SourceId>,
         source_tags: &[SourceTag],
         guids: &[FunctionGUID],
+        constraints: &[ConstraintGUID],
     ) -> serde_json::Value {
-        let guids_str: Vec<String> = guids.iter().map(|g| g.to_string()).collect();
+        let guids_str: HashSet<String> = guids.iter().map(|g| g.to_string()).collect();
         // TODO: The limit here needs to be somewhat flexible. But 1000 will do for now.
         let mut body = json!({
             "format": "flatbuffer",
             "guids": guids_str,
-            "limit": 1000
+            "limit": 10000,
         });
         if let Some(target_id) = target {
             body["target_id"] = json!(target_id);
@@ -231,6 +234,11 @@ impl NetworkClient {
         if !source_tags.is_empty() {
             body["source_tags"] = json!(source_tags);
         }
+        if !constraints.is_empty() {
+            let constraint_guids_str: HashSet<String> =
+                constraints.iter().map(|g| g.to_string()).collect();
+            body["constraints"] = json!(constraint_guids_str);
+        }
         body
     }
 
@@ -244,13 +252,13 @@ impl NetworkClient {
         target: Option<NetworkTargetId>,
         source: Option<SourceId>,
         guids: &[FunctionGUID],
+        constraints: &[ConstraintGUID],
     ) -> Result<WarpFile<'static>, String> {
         let query_functions_url = format!("{}/api/v1/functions/query", self.server_url);
         // TODO: Allow for source tags? We really only need this in query_functions_source as that
         // TODO: is what prevents a undesired source from being "known" to the container.
-        let payload = Self::query_functions_body(target, source, &[], guids);
+        let payload = Self::query_functions_body(target, source, &[], guids, constraints);
         let mut inst = self.provider.create_instance().unwrap();
-
         let resp = inst.post_json(&query_functions_url, self.headers.clone(), &payload)?;
         if !resp.is_success() {
             return Err(format!(
@@ -275,7 +283,10 @@ impl NetworkClient {
     ) -> Result<HashMap<SourceId, Vec<FunctionGUID>>, String> {
         let query_functions_source_url =
             format!("{}/api/v1/functions/query/source", self.server_url);
-        let payload = Self::query_functions_body(target, None, tags, guids);
+        // NOTE: We do not filter by constraint guids here since this pass is only responsible for
+        // returning the source ids, not the actual function data, see [`NetworkClient::query_functions`]
+        // for the place where the constraints are applied, and _do_ matter.
+        let payload = Self::query_functions_body(target, None, tags, guids, &[]);
         let mut inst = self.provider.create_instance().unwrap();
 
         let resp = inst.post_json(&query_functions_source_url, self.headers.clone(), &payload)?;
diff --git a/plugins/warp/src/plugin/ffi/container.rs b/plugins/warp/src/plugin/ffi/container.rs
index 79f45dcb..aa0c37d3 100644
--- a/plugins/warp/src/plugin/ffi/container.rs
+++ b/plugins/warp/src/plugin/ffi/container.rs
@@ -5,7 +5,8 @@ use crate::container::{
 };
 use crate::convert::{from_bn_type, to_bn_type};
 use crate::plugin::ffi::{
-    BNWARPContainer, BNWARPFunction, BNWARPFunctionGUID, BNWARPSource, BNWARPTarget, BNWARPTypeGUID,
+    BNWARPConstraintGUID, BNWARPContainer, BNWARPFunction, BNWARPFunctionGUID, BNWARPSource,
+    BNWARPTarget, BNWARPTypeGUID,
 };
 use binaryninja::architecture::CoreArchitecture;
 use binaryninja::binary_view::BinaryView;
@@ -218,6 +219,8 @@ pub unsafe extern "C" fn BNWARPContainerFetchFunctions(
     source_tags_count: usize,
     guids: *const BNWARPFunctionGUID,
     count: usize,
+    constraints: *const BNWARPConstraintGUID,
+    constraints_count: usize,
 ) {
     let arc_container = ManuallyDrop::new(Arc::from_raw(container));
     let Ok(container) = arc_container.read() else {
@@ -234,8 +237,9 @@ pub unsafe extern "C" fn BNWARPContainerFetchFunctions(
         .collect();
 
     let guids = unsafe { std::slice::from_raw_parts(guids, count) };
+    let constraints = unsafe { std::slice::from_raw_parts(constraints, constraints_count) };
 
-    if let Err(e) = container.fetch_functions(&target, &source_tags, guids) {
+    if let Err(e) = container.fetch_functions(&target, &source_tags, guids, constraints) {
         tracing::error!("Failed to fetch functions: {}", e);
     }
 }
diff --git a/plugins/warp/src/plugin/settings.rs b/plugins/warp/src/plugin/settings.rs
index 6469be0f..896bee7f 100644
--- a/plugins/warp/src/plugin/settings.rs
+++ b/plugins/warp/src/plugin/settings.rs
@@ -40,7 +40,7 @@ pub struct PluginSettings {
 impl PluginSettings {
     pub const ALLOWED_SOURCE_TAGS_DEFAULT: [&'static str; 2] = ["official", "trusted"];
     pub const ALLOWED_SOURCE_TAGS_SETTING: &'static str = "warp.fetcher.allowedSourceTags";
-    pub const FETCH_BATCH_SIZE_DEFAULT: usize = 100;
+    pub const FETCH_BATCH_SIZE_DEFAULT: usize = 10000;
     pub const FETCH_BATCH_SIZE_SETTING: &'static str = "warp.fetcher.fetchBatchSize";
     pub const LOAD_BUNDLED_FILES_DEFAULT: bool = true;
     pub const LOAD_BUNDLED_FILES_SETTING: &'static str = "warp.container.loadBundledFiles";
@@ -81,8 +81,8 @@ impl PluginSettings {
         let fetch_size_props = json!({
             "title" : "Fetch Batch Limit",
             "type" : "number",
-            "minValue" : 1,
-            "maxValue" : 1000,
+            "minValue" : 100,
+            "maxValue" : 20000,
             "default" : Self::FETCH_BATCH_SIZE_DEFAULT,
             "description" : "The maximum number of functions to fetch in a single batch. This is used to limit the amount of functions to fetch at once, lowering this value will make the fetch process more comprehensive at the cost of more network requests.",
             "ignore" : [],
diff --git a/plugins/warp/src/plugin/workflow.rs b/plugins/warp/src/plugin/workflow.rs
index 1f4ef101..52c8e249 100644
--- a/plugins/warp/src/plugin/workflow.rs
+++ b/plugins/warp/src/plugin/workflow.rs
@@ -24,6 +24,7 @@ use std::cmp::Ordering;
 use std::collections::HashMap;
 use std::time::Instant;
 use warp::r#type::class::function::{Location, RegisterLocation, StackLocation};
+use warp::signature::constraint::ConstraintGUID;
 use warp::signature::function::{Function, FunctionGUID};
 use warp::target::Target;
 
@@ -171,7 +172,7 @@ pub fn run_matcher(view: &BinaryView) {
             .maximum_possible_functions
             .is_some_and(|max| max < matched_functions.len() as u64)
         {
-            tracing::warn!(
+            tracing::debug!(
                 "Skipping {}, too many possible functions: {}",
                 guid,
                 matched_functions.len()
@@ -270,6 +271,20 @@ pub fn run_fetcher(view: &BinaryView) {
     let mut query_opts = QueryOptions::new_with_view(view);
     let plugin_settings = PluginSettings::from_settings(&view_settings, &mut query_opts);
 
+    let is_ignored_func = |f: &BNFunction| !f.function_tags(None, Some(IGNORE_TAG_NAME)).is_empty();
+
+    let constraints: Vec<ConstraintGUID> = view
+        .functions()
+        .iter()
+        // Skip functions that have the ignored tag! Otherwise, we will store their constraints.
+        .filter(|f| !is_ignored_func(f))
+        .filter_map(|f| {
+            let function = try_cached_function_match(&f)?;
+            Some(function.constraints.into_iter().map(|c| c.guid))
+        })
+        .flatten()
+        .collect();
+
     let Some(function_set) = FunctionSet::from_view(view) else {
         background_task.finish();
         return;
@@ -285,8 +300,12 @@ pub fn run_fetcher(view: &BinaryView) {
                 if background_task.is_cancelled() {
                     break;
                 }
-                let _ =
-                    container.fetch_functions(target, &plugin_settings.allowed_source_tags, batch);
+                let _ = container.fetch_functions(
+                    target,
+                    &plugin_settings.allowed_source_tags,
+                    batch,
+                    &constraints,
+                );
             }
         }
     });