[WARP] Server-side constraint matching

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

Author: emesare

plugins/warp/api/python/warp.py

@@ -426,19 +426,25 @@ def remove_types(self, source: Source, guids: List[TypeGUID]) -> bool:
             core_guids[i] = guids[i].uuid
         return warpcore.BNWARPContainerRemoveTypes(self.handle, source.uuid, core_guids, count)
 
-    def fetch_functions(self, target: WarpTarget, guids: List[FunctionGUID], source_tags: Optional[List[str]] = None):
+    def fetch_functions(self, target: WarpTarget, guids: List[FunctionGUID], source_tags: Optional[List[str]] = None, constraints: Optional[List[ConstraintGUID]] = None):
         count = len(guids)
         core_guids = (warpcore.BNWARPFunctionGUID * count)()
         for i in range(count):
             core_guids[i] = guids[i].uuid
+        if constraints is None:
+            constraints = []
+        constraints_count = len(constraints)
+        core_constraints = (warpcore.BNWARPConstraintGUID * constraints_count)()
+        for i in range(constraints_count):
+            core_constraints[i] = constraints[i].uuid
         if source_tags is None:
             source_tags = []
         source_tags_ptr = (ctypes.c_char_p * len(source_tags))()
         source_tags_len = len(source_tags)
         for i in range(len(source_tags)):
             source_tags_ptr[i] = source_tags[i].encode('utf-8')
         source_tags_array_ptr = ctypes.cast(source_tags_ptr, ctypes.POINTER(ctypes.c_char_p))
-        warpcore.BNWARPContainerFetchFunctions(self.handle, target.handle, source_tags_array_ptr, source_tags_len, core_guids, count)
+        warpcore.BNWARPContainerFetchFunctions(self.handle, target.handle, source_tags_array_ptr, source_tags_len, core_guids, count, core_constraints, constraints_count)
 
     def get_sources_with_function_guid(self, target: WarpTarget, guid: FunctionGUID) -> List[Source]:
         count = ctypes.c_size_t()

plugins/warp/api/warp.cpp

@@ -353,7 +353,7 @@ bool Container::RemoveTypes(const Source &source, const std::vector<TypeGUID> &g
     return result;
 }
 
-void Container::FetchFunctions(const Target &target, const std::vector<FunctionGUID> &guids, const std::vector<SourceTag> &tags) const
+void Container::FetchFunctions(const Target &target, const std::vector<FunctionGUID> &guids, const std::vector<SourceTag> &tags, const std::vector<ConstraintGUID> &constraints) const
 {
     size_t count = guids.size();
     BNWARPFunctionGUID *apiGuids = new BNWARPFunctionGUID[count];
@@ -363,9 +363,14 @@ void Container::FetchFunctions(const Target &target, const std::vector<FunctionG
     const char** rawTags = new const char*[tagCount];
     for (size_t i = 0; i < tagCount; i++)
         rawTags[i] = tags[i].c_str();
-    BNWARPContainerFetchFunctions(m_object, target.m_object, rawTags, tagCount, apiGuids, count);
+    size_t constraintCount = constraints.size();
+    BNWARPConstraintGUID *apiConstraints = new BNWARPConstraintGUID[constraintCount];
+    for (size_t i = 0; i < constraintCount; i++)
+        apiConstraints[i] = *constraints[i].Raw();
+    BNWARPContainerFetchFunctions(m_object, target.m_object, rawTags, tagCount, apiGuids, count, apiConstraints, constraintCount);
     delete[] apiGuids;
     delete[] rawTags;
+    delete[] apiConstraints;
 }
 
 std::vector<Source> Container::GetSourcesWithFunctionGUID(const Target& target, const FunctionGUID &guid) const

plugins/warp/api/warp.h

@@ -409,7 +409,7 @@ namespace Warp {
 
 		bool RemoveTypes(const Source &source, const std::vector<TypeGUID> &guids) const;
 
-		void FetchFunctions(const Target &target, const std::vector<FunctionGUID> &guids, const std::vector<SourceTag> &tags = {}) const;
+		void FetchFunctions(const Target &target, const std::vector<FunctionGUID> &guids, const std::vector<SourceTag> &tags = {}, const std::vector<ConstraintGUID> &constraints = {}) const;
 
 		std::vector<Source> GetSourcesWithFunctionGUID(const Target &target, const FunctionGUID &guid) const;
 

plugins/warp/api/warpcore.h

@@ -128,7 +128,7 @@ extern "C"
     WARP_FFI_API bool BNWARPContainerRemoveFunctions(BNWARPContainer* container, const BNWARPTarget* target, const BNWARPSource* source, BNWARPFunction** functions, size_t count);
     WARP_FFI_API bool BNWARPContainerRemoveTypes(BNWARPContainer* container, const BNWARPSource* source, BNWARPTypeGUID* types, size_t count);
 
-    WARP_FFI_API void BNWARPContainerFetchFunctions(BNWARPContainer* container, BNWARPTarget* target, const char** sourceTags, size_t sourceTagCount, const BNWARPTypeGUID* guids, size_t count);
+    WARP_FFI_API void BNWARPContainerFetchFunctions(BNWARPContainer* container, BNWARPTarget* target, const char** sourceTags, size_t sourceTagCount, const BNWARPFunctionGUID* guids, size_t count, const BNWARPConstraintGUID* constraints, size_t constraintCount);
 
     WARP_FFI_API BNWARPSource* BNWARPContainerGetSourcesWithFunctionGUID(BNWARPContainer* container, const BNWARPTarget* target, const BNWARPFunctionGUID* guid, size_t* count);
     WARP_FFI_API BNWARPSource* BNWARPContainerGetSourcesWithTypeGUID(BNWARPContainer* container, const BNWARPTypeGUID* guid, size_t* count);

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(())
     }

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.
     ///
@@ -110,7 +111,7 @@ impl NetworkContainer {
         guids: &[FunctionGUID],
     ) -> HashMap<SourceId, Vec<FunctionGUID>> {
         let Some(target_id) = target.and_then(|t| self.get_target_id(t)) else {
-            tracing::debug!("Cannot query functions source without a target, skipping...");
+            tracing::error!("Cannot query functions source without a target, skipping...");
             return HashMap::new();
         };
 
@@ -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(())

plugins/warp/src/container/network/client.rs

@@ -7,12 +7,14 @@ 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 std::time::Instant;
 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 +32,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 +217,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 +235,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 +253,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 +284,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)?;

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);
     }
 }

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" : [],

plugins/warp/src/plugin/workflow.rs

@@ -21,9 +21,10 @@ use itertools::Itertools;
 use rayon::iter::IntoParallelIterator;
 use rayon::iter::ParallelIterator;
 use std::cmp::Ordering;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 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,
+                );
             }
         }
     });