[WARP] Enhanced network support

1 files changed, 383 insertions, 22 deletions

diff --git a/plugins/warp/src/container/network/client.rs b/plugins/warp/src/container/network/client.rs
index f77f1118..39e7640a 100644
--- a/plugins/warp/src/container/network/client.rs
+++ b/plugins/warp/src/container/network/client.rs
@@ -1,12 +1,20 @@
 use crate::container::network::NetworkTargetId;
-use crate::container::SourceId;
+use crate::container::{
+    ContainerSearchItem, ContainerSearchItemKind, ContainerSearchQuery, ContainerSearchResponse,
+    SourceId, SourcePath, SourceTag,
+};
 use reqwest::blocking::Client;
 use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION};
 use reqwest::StatusCode;
+use serde::Deserialize;
 use serde_json::json;
 use std::collections::HashMap;
 use std::str::FromStr;
-use warp::signature::function::FunctionGUID;
+use uuid::Uuid;
+use warp::chunk::ChunkKind;
+use warp::r#type::guid::TypeGUID;
+use warp::r#type::{ComputedType, Type};
+use warp::signature::function::{Function, FunctionGUID};
 use warp::target::Target;
 use warp::WarpFile;
 
@@ -16,7 +24,7 @@ use warp::WarpFile;
 #[derive(Clone, Debug)]
 pub struct NetworkClient {
     client: Client,
-    server_url: String,
+    pub server_url: String,
 }
 
 impl NetworkClient {
@@ -65,6 +73,116 @@ impl NetworkClient {
         Ok(resp.status())
     }
 
+    /// Query the logged in user.
+    ///
+    /// NOTE: **THIS IS BLOCKING**
+    ///
+    /// Route: `api/v1/users/me` (TODO: Comment about the query)
+    pub fn current_user(&self) -> reqwest::Result<(i32, String)> {
+        let current_user_url = format!("{}/api/v1/users/me", self.server_url);
+
+        #[derive(Deserialize)]
+        struct CurrentUser {
+            username: String,
+            id: i32,
+        }
+
+        let resp = self
+            .client
+            .get(&current_user_url)
+            .send()?
+            .error_for_status()?;
+        let user: CurrentUser = resp.json()?;
+        Ok((user.id, user.username))
+    }
+
+    /// Query the logged in user.
+    ///
+    /// NOTE: **THIS IS BLOCKING**
+    ///
+    /// Route: `api/v1/users/me` (TODO: Comment about the query)
+    pub fn source_name(&self, id: SourceId) -> reqwest::Result<String> {
+        let source_url = format!("{}/api/v1/sources/{}", self.server_url, id);
+
+        #[derive(Deserialize)]
+        struct Source {
+            name: String,
+        }
+
+        let resp = self.client.get(&source_url).send()?.error_for_status()?;
+        let src: Source = resp.json()?;
+        Ok(src.name)
+    }
+
+    /// Create a new source with the given name.
+    ///
+    /// The current user will be added to the source.
+    ///
+    /// NOTE: You must be logged in to create a source.
+    ///
+    /// NOTE: **THIS IS BLOCKING**
+    ///
+    /// Route: `api/v1/sources/`
+    pub fn create_source(&self, name: &str) -> reqwest::Result<SourceId> {
+        let source_url = format!("{}/api/v1/sources", self.server_url);
+
+        let body = json!({
+            "name": name,
+            // Passing nothing here will add the current user to the source.
+            "user_ids": []
+        });
+
+        #[derive(Deserialize)]
+        struct CreateSourceResponse {
+            id: Uuid,
+        }
+
+        let resp = self
+            .client
+            .post(&source_url)
+            .json(&body)
+            .send()?
+            .error_for_status()?;
+
+        let parsed: CreateSourceResponse = resp.json()?;
+        Ok(SourceId(parsed.id))
+    }
+
+    /// Query the [`SourceId`]s for the given user.
+    ///
+    /// NOTE: **THIS IS BLOCKING**
+    ///
+    /// Route: `api/v1/sources/query` (TODO: Comment about the query)
+    pub fn query_sources(&self, user_id: Option<i32>) -> reqwest::Result<Vec<SourceId>> {
+        let sources_url = format!("{}/api/v1/sources/query", self.server_url);
+
+        #[derive(Deserialize)]
+        struct SourceItem {
+            id: Uuid,
+        }
+
+        #[derive(Deserialize)]
+        struct SourcesQueryResponse {
+            items: Vec<SourceItem>,
+        }
+
+        let mut query = HashMap::new();
+        if let Some(user_id) = user_id {
+            query.insert("user_id", user_id);
+        }
+        let query_str = json!(query).to_string();
+        let resp = self
+            .client
+            .post(&sources_url)
+            .body(query_str)
+            .header("Content-Type", "application/json")
+            .send()?
+            .error_for_status()?;
+
+        let parsed: SourcesQueryResponse = resp.json()?;
+        Ok(parsed.items.into_iter().map(|it| SourceId(it.id)).collect())
+    }
+
     /// Query the [`NetworkTargetId`] for the given [`Target`].
     ///
     /// NOTE: **THIS IS BLOCKING**
@@ -78,25 +196,36 @@ impl NetworkClient {
             query.insert("platform", platform);
         }
         if let Some(architecture) = &target.architecture {
-            query.insert("architecture", architecture);
+            query.insert("arch", architecture);
+        }
+        let query_str = json!(query).to_string();
+
+        #[derive(Deserialize)]
+        struct TargetQueryResponse {
+            id: NetworkTargetId,
         }
 
         // NOTE: This is blocking.
-        let target_id: NetworkTargetId = self
+        let response = self
             .client
-            .get(query_target_url)
-            .query(&query)
+            .post(query_target_url)
+            .body(query_str)
+            .header("Content-Type", "application/json")
             .send()
-            .ok()?
-            .json::<NetworkTargetId>()
             .ok()?;
 
-        Some(target_id)
+        // Assuming the first response is the one we want.
+        // TODO: Handle multiple responses, or error out.
+        let json_response: Vec<TargetQueryResponse> = response.json().ok()?;
+        let first_response = json_response.first()?;
+
+        Some(first_response.id)
     }
 
     fn query_functions_body(
         target: Option<NetworkTargetId>,
         source: Option<SourceId>,
+        source_tags: &[SourceTag],
         guids: &[FunctionGUID],
     ) -> serde_json::Value {
         let guids_str: Vec<String> = guids.iter().map(|g| g.to_string()).collect();
@@ -112,6 +241,9 @@ impl NetworkClient {
         if let Some(source_id) = source {
             body["source_id"] = json!(source_id.to_string());
         }
+        if !source_tags.is_empty() {
+            body["source_tags"] = json!(source_tags);
+        }
         body
     }
 
@@ -127,7 +259,9 @@ impl NetworkClient {
         guids: &[FunctionGUID],
     ) -> Option<WarpFile<'static>> {
         let query_functions_url = format!("{}/api/v1/functions/query", self.server_url);
-        let payload = Self::query_functions_body(target, source, guids);
+        // 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);
 
         // Make the POST request
         let response = self
@@ -154,11 +288,12 @@ impl NetworkClient {
     pub fn query_functions_source(
         &self,
         target: Option<NetworkTargetId>,
+        tags: &[SourceTag],
         guids: &[FunctionGUID],
     ) -> Option<HashMap<SourceId, Vec<FunctionGUID>>> {
         let query_functions_source_url =
             format!("{}/api/v1/functions/query/source", self.server_url);
-        let payload = Self::query_functions_body(target, None, guids);
+        let payload = Self::query_functions_body(target, None, tags, guids);
 
         // Make the POST request
         let response = self
@@ -194,20 +329,24 @@ impl NetworkClient {
     /// NOTE: **THIS IS BLOCKING**
     ///
     /// Route: `api/v1/files/{source}`
-    pub fn push_file(&self, source_id: SourceId, file: &WarpFile) -> bool {
-        let push_file_url = format!("{}/api/v1/files/{}", self.server_url, source_id.to_string());
+    pub fn push_file(&self, source_id: SourceId, file: &WarpFile, name: &str) -> bool {
+        let push_file_url = format!("{}/api/v1/files", self.server_url);
 
         // Convert WarpFile to bytes
         let file_bytes = file.to_bytes();
 
-        // Create the form part with the file
-        let form = reqwest::blocking::multipart::Form::new().part(
-            "file",
-            reqwest::blocking::multipart::Part::bytes(file_bytes)
-                .file_name("data.warp")
-                .mime_str("application/octet-stream")
-                .unwrap(),
-        );
+        let Ok(file_part) = reqwest::blocking::multipart::Part::bytes(file_bytes)
+            .file_name("data.warp")
+            .mime_str("application/octet-stream")
+        else {
+            log::error!("Failed to create file part");
+            return false;
+        };
+
+        let form = reqwest::blocking::multipart::Form::new()
+            .part("file", file_part)
+            .text("name", name.to_string())
+            .text("source", source_id.to_string());
 
         // Send the request
         match self.client.post(&push_file_url).multipart(form).send() {
@@ -225,4 +364,226 @@ impl NetworkClient {
             }
         }
     }
+
+    pub fn function_data(&self, id: i32) -> Option<Function> {
+        let function_data_url = format!("{}/api/v1/functions/{}/data", self.server_url, id);
+        let response = self.client.get(&function_data_url).send().ok()?;
+        if !response.status().is_success() {
+            log::error!(
+                "Failed to fetch function data for {}: {}",
+                id,
+                response.status()
+            );
+            return None;
+        }
+        let bytes = response.bytes().ok()?;
+        Function::from_bytes(bytes.as_ref())
+    }
+
+    pub fn function_datas(&self, ids: &[i32]) -> Option<Vec<Function>> {
+        if ids.is_empty() {
+            return Some(Vec::new());
+        }
+        let function_data_url = format!("{}/api/v1/functions/data", self.server_url);
+        let body = json!({
+            "ids": ids,
+        });
+        let response = self
+            .client
+            .post(&function_data_url)
+            .json(&body)
+            .send()
+            .ok()?;
+        if !response.status().is_success() {
+            log::error!("Failed to fetch function data: {}", response.status());
+            return None;
+        }
+        let bytes = response.bytes().ok()?;
+        let file = WarpFile::from_bytes(bytes.as_ref())?;
+        let mut functions = Vec::with_capacity(ids.len());
+        for chunk in file.chunks {
+            let ChunkKind::Signature(sc) = chunk.kind else {
+                continue;
+            };
+            functions.extend(sc.functions());
+        }
+        Some(functions)
+    }
+
+    pub fn type_data(&self, guid: TypeGUID) -> Option<Type> {
+        let type_data_url = format!("{}/api/v1/types/{}/data", self.server_url, guid.to_string());
+        let response = self.client.get(&type_data_url).send().ok()?;
+        if !response.status().is_success() {
+            log::error!(
+                "Failed to fetch type data for {}: {}",
+                guid.to_string(),
+                response.status()
+            );
+            return None;
+        }
+        let bytes = response.bytes().ok()?;
+        Type::from_bytes(bytes.as_ref())
+    }
+
+    pub fn type_datas(&self, guids: &[TypeGUID]) -> Option<Vec<ComputedType>> {
+        if guids.is_empty() {
+            return Some(Vec::new());
+        }
+        let type_data_url = format!("{}/api/v1/types/data", self.server_url);
+        let body = json!({
+            "ids": guids.iter().map(|g| g.to_string()).collect::<Vec<_>>(),
+        });
+        let response = self.client.post(&type_data_url).json(&body).send().ok()?;
+        if !response.status().is_success() {
+            log::error!("Failed to fetch type data: {}", response.status());
+            return None;
+        }
+        let bytes = response.bytes().ok()?;
+        let file = WarpFile::from_bytes(bytes.as_ref())?;
+        let mut types = Vec::with_capacity(guids.len());
+        for chunk in file.chunks {
+            let ChunkKind::Type(tc) = chunk.kind else {
+                continue;
+            };
+            types.extend(tc.types());
+        }
+        Some(types)
+    }
+
+    pub fn search(&self, query: &ContainerSearchQuery) -> Option<ContainerSearchResponse> {
+        let search_url = format!("{}/api/v1/search", self.server_url);
+
+        #[derive(serde::Serialize)]
+        struct SearchRequest<'a> {
+            #[serde(rename = "q")]
+            q: &'a str,
+            #[serde(skip_serializing_if = "Option::is_none")]
+            limit: Option<usize>,
+            #[serde(skip_serializing_if = "Option::is_none")]
+            offset: Option<usize>,
+            #[serde(rename = "source_id", skip_serializing_if = "Option::is_none")]
+            source_id: Option<String>,
+            #[serde(skip_serializing_if = "Option::is_none")]
+            source_tags: Option<Vec<SourceTag>>,
+            #[serde(skip_serializing_if = "Option::is_none")]
+            retrieve_data: Option<bool>,
+        }
+
+        #[derive(serde::Deserialize)]
+        struct SearchResponse {
+            items: Vec<SearchItem>,
+            offset: usize,
+            total: usize,
+        }
+
+        #[derive(serde::Deserialize)]
+        struct SearchItem {
+            id: String,
+            kind: String,
+            #[serde(default)]
+            name: Option<String>,
+            #[serde(default)]
+            source_id: Option<Uuid>,
+            #[serde(default)]
+            data: Option<Vec<u8>>,
+        }
+
+        let source_id_str = query.source.map(|s| s.to_string());
+        let request = SearchRequest {
+            q: &query.query,
+            limit: query.limit,
+            offset: query.offset,
+            source_id: source_id_str,
+            source_tags: match query.tags.is_empty() {
+                true => None,
+                false => Some(query.tags.clone()),
+            },
+            // This must be passed to retrieve the function and type data.
+            retrieve_data: Some(true),
+        };
+
+        let resp = match self.client.get(search_url).query(&request).send() {
+            Ok(r) => r,
+            Err(err) => {
+                log::error!("Failed to send search request: {}", err);
+                return None;
+            }
+        };
+
+        let Ok(parsed) = resp.json::<SearchResponse>() else {
+            log::error!("Failed to parse search response");
+            return None;
+        };
+
+        // TODO: This is quite scuffed, but it works for now. (Mostly just that it looks bad and queries a lot)
+        // TODO: Here I think would be a good place to sort it so sources always come first.
+        // TODO: Users searching will want to get to the source first, likely to whitelist or blacklist.
+        let mut items = Vec::with_capacity(parsed.items.len());
+        for item in parsed.items {
+            let Some(source_uuid) = item.source_id else {
+                // Currently not interested in items without a source id.
+                // Things like symbols do not have a source id.
+                continue;
+            };
+
+            let kind = match item.kind.as_str() {
+                "function" => {
+                    let Some(data) = &item.data else {
+                        log::warn!(
+                            "Function item {} has no data from network, skipping...",
+                            item.id
+                        );
+                        continue;
+                    };
+                    let Some(func) = Function::from_bytes(&data) else {
+                        log::warn!(
+                            "Function item {} has invalid data from network, skipping...",
+                            item.id
+                        );
+                        continue;
+                    };
+                    ContainerSearchItemKind::Function(func)
+                }
+                "source" => ContainerSearchItemKind::Source {
+                    path: match item.name {
+                        None => {
+                            log::warn!("Source item {} has no name", item.id);
+                            continue;
+                        }
+                        Some(name) => SourcePath(format!("{}/{}", self.server_url, name).into()),
+                    },
+                    id: SourceId(source_uuid),
+                },
+                "type" => {
+                    let Some(data) = &item.data else {
+                        log::warn!(
+                            "Type item {} has no data from network, skipping...",
+                            item.id
+                        );
+                        continue;
+                    };
+                    let Some(ty) = Type::from_bytes(&data) else {
+                        log::warn!(
+                            "Type item {} has invalid data from network, skipping...",
+                            item.id
+                        );
+                        continue;
+                    };
+                    ContainerSearchItemKind::Type(ty)
+                }
+                _ => continue,
+            };
+
+            items.push(ContainerSearchItem {
+                source: SourceId(source_uuid),
+                kind,
+            });
+        }
+
+        Some(ContainerSearchResponse {
+            items,
+            total: parsed.total,
+            offset: parsed.offset,
+        })
+    }
 }