2 files changed, 41 insertions, 20 deletions

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