1 files changed, 18 insertions, 7 deletions

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