From d52991073d6fa4e9fc8359a34fa9b28626633240 Mon Sep 17 00:00:00 2001 From: Mason Reed Date: Fri, 27 Feb 2026 14:10:03 -0800 Subject: [PATCH 1/4] [WARP] Demote locking surrounding container function fetching By demoting the containers lock to read only for fetching we can prevent blocking the main ui thread while waiting for the network requests to finish --- plugins/warp/src/container.rs | 5 +- plugins/warp/src/container/network.rs | 147 ++++++++++++++++------- plugins/warp/src/plugin/ffi/container.rs | 4 +- 3 files changed, 108 insertions(+), 48 deletions(-) diff --git a/plugins/warp/src/container.rs b/plugins/warp/src/container.rs index b5e92bc7d1..79c243b3ee 100644 --- a/plugins/warp/src/container.rs +++ b/plugins/warp/src/container.rs @@ -39,6 +39,8 @@ pub enum ContainerError { SearchFailed(String), #[error("failed to commit source '{0}': {1}")] CommitFailed(SourceId, String), + #[error("container error encountered: {0}")] + Custom(String), } /// Represents the ID for a single container source. @@ -224,6 +226,7 @@ pub trait Container: Send + Sync + Display + Debug { /// to verify the permissions of the source. fn add_source(&mut self, path: SourcePath) -> ContainerResult; + // TODO: Make interior mutable. /// Flush changes made to a source. /// /// Because writing to a source can require file or network operations, we let the container @@ -293,7 +296,7 @@ pub trait Container: Send + Sync + Display + Debug { /// 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. fn fetch_functions( - &mut self, + &self, _target: &Target, _tags: &[SourceTag], _functions: &[FunctionGUID], diff --git a/plugins/warp/src/container/network.rs b/plugins/warp/src/container/network.rs index 0ba515f9d9..307b44fa1f 100644 --- a/plugins/warp/src/container/network.rs +++ b/plugins/warp/src/container/network.rs @@ -3,10 +3,12 @@ use crate::container::{ Container, ContainerError, ContainerResult, ContainerSearchQuery, ContainerSearchResponse, SourceId, SourcePath, SourceTag, }; +use dashmap::DashMap; use directories::ProjectDirs; use std::collections::{HashMap, HashSet}; use std::fmt::{Debug, Display, Formatter}; use std::path::PathBuf; +use std::sync::RwLock; use warp::chunk::{Chunk, ChunkKind, CompressionType}; use warp::r#type::chunk::TypeChunk; use warp::r#type::guid::TypeGUID; @@ -28,15 +30,21 @@ pub struct NetworkContainer { client: NetworkClient, /// This is the store that the interface will write to; then we have special functions for pulling /// and pushing to the network source. - cache: DiskContainer, + cache: RwLock, /// Where to place newly created sources. /// /// This is typically a directory inside [`NetworkContainer::root_cache_location`]. cache_path: PathBuf, /// Populated when targets are queried. - known_targets: HashMap>, + /// + /// 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_targets: DashMap>, /// Populated when function sources are queried. - known_function_sources: HashMap>, + /// + /// 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>, /// Populated when user adds function, this is used for writing back to the server. added_chunks: HashMap>>, /// Populated when connecting to the server, this is used to determine which sources are writable. @@ -47,12 +55,12 @@ pub struct NetworkContainer { impl NetworkContainer { pub fn new(client: NetworkClient, cache_path: PathBuf, writable_sources: &[SourceId]) -> Self { - let mut container = Self { - cache: DiskContainer::new_from_dir(cache_path.clone()), + let container = Self { + cache: RwLock::new(DiskContainer::new_from_dir(cache_path.clone())), cache_path, client, - known_targets: HashMap::new(), - known_function_sources: HashMap::new(), + known_targets: DashMap::new(), + known_function_sources: DashMap::new(), added_chunks: HashMap::new(), writable_sources: writable_sources.into_iter().copied().collect(), }; @@ -74,7 +82,7 @@ impl NetworkContainer { /// # Caching policy /// /// The [`NetworkTargetId`] is unique and immutable, so they will be persisted indefinitely. - pub fn get_target_id(&mut self, target: &Target) -> Option { + pub fn get_target_id(&self, target: &Target) -> Option { // It's highly probable we have previously queried the target, check that first. if let Some(target_id) = self.known_targets.get(target) { return target_id.clone(); @@ -96,7 +104,7 @@ impl NetworkContainer { /// for now as the requests for functions come at the request of some user interaction. Any guid /// with no sources will still be cached. pub fn get_unseen_functions_source( - &mut self, + &self, target: Option<&Target>, tags: &[SourceTag], guids: &[FunctionGUID], @@ -157,12 +165,7 @@ 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( - &mut self, - target: &Target, - source: &SourceId, - functions: &[FunctionGUID], - ) { + pub fn pull_functions(&self, target: &Target, source: &SourceId, functions: &[FunctionGUID]) { let target_id = self.get_target_id(target); let file = match self .client @@ -182,18 +185,25 @@ impl NetworkContainer { let functions: Vec<_> = sc.functions().collect(); // Probe the source before attempting to access it, as it might not exist locally. self.probe_source(*source); - match self.cache.add_functions(target, source, &functions) { - Ok(_) => tracing::debug!( - "Added {} functions into cached source '{}'", - functions.len(), - source - ), - Err(err) => tracing::error!( - "Failed to add {} function into cached source '{}': {}", - functions.len(), - source, - err - ), + + match self.cache.write() { + Ok(mut cache) => match cache.add_functions(target, source, &functions) { + Ok(_) => tracing::debug!( + "Added {} functions into cached source '{}'", + functions.len(), + source + ), + Err(err) => tracing::error!( + "Failed to add {} function into cached source '{}': {}", + functions.len(), + source, + err + ), + }, + Err(err) => { + tracing::error!("Failed to write to cache: {}", err); + return; + } } } // TODO; Probably want to pull type in with this. @@ -214,8 +224,13 @@ impl NetworkContainer { /// Probe the source to make sure it exists in the cache. Retrieving the name from the server. /// /// **This is blocking** - pub fn probe_source(&mut self, source_id: SourceId) { - if !self.cache.source_path(&source_id).is_ok() { + pub fn probe_source(&self, source_id: SourceId) { + let Ok(mut cache) = self.cache.write() else { + tracing::error!("Cannot probe source '{}', cache is poisoned", source_id); + return; + }; + + if !cache.source_path(&source_id).is_ok() { // Add the source to the cache. Using the source id and source name as the source path. match self.client.source_name(source_id) { Ok(source_name) => { @@ -224,7 +239,7 @@ impl NetworkContainer { .cache_path .join(source_id.to_string()) .join(source_name); - let _ = self.cache.insert_source(source_id, SourcePath(source_path)); + let _ = cache.insert_source(source_id, SourcePath(source_path)); } Err(e) => { tracing::error!("Failed to probe source '{}': {}", source_id, e); @@ -251,7 +266,10 @@ impl NetworkContainer { impl Container for NetworkContainer { fn sources(&self) -> ContainerResult> { - self.cache.sources() + self.cache + .read() + .map_err(|e| ContainerError::Custom(format!("Cache read error: {}", e)))? + .sources() } fn add_source(&mut self, path: SourcePath) -> ContainerResult { @@ -295,11 +313,17 @@ impl Container for NetworkContainer { } fn source_tags(&self, source: &SourceId) -> ContainerResult> { - self.cache.source_tags(source) + self.cache + .read() + .map_err(|e| ContainerError::Custom(format!("Cache read error: {}", e)))? + .source_tags(source) } fn source_path(&self, source: &SourceId) -> ContainerResult { - self.cache.source_path(source) + self.cache + .read() + .map_err(|e| ContainerError::Custom(format!("Cache read error: {}", e)))? + .source_path(source) } fn add_computed_types( @@ -310,7 +334,10 @@ impl Container for NetworkContainer { // NOTE: We must `add_computed_types` to the cache before we add the chunk, as `added_chunks` is // not consulted when retrieving types from the cache, if we fail to add the types to // the cache, we will not see them show up in the UI or when matching. - self.cache.add_computed_types(source, types)?; + self.cache + .write() + .map_err(|e| ContainerError::Custom(format!("Cache write error: {}", e)))? + .add_computed_types(source, types)?; let type_chunk = TypeChunk::new_with_computed(types).ok_or( ContainerError::CorruptedData("signature chunk failed to validate"), )?; @@ -320,7 +347,10 @@ impl Container for NetworkContainer { } fn remove_types(&mut self, source: &SourceId, guids: &[TypeGUID]) -> ContainerResult<()> { - self.cache.remove_types(source, guids) + self.cache + .write() + .map_err(|e| ContainerError::Custom(format!("Cache write error: {}", e)))? + .remove_types(source, guids) } fn add_functions( @@ -332,7 +362,10 @@ impl Container for NetworkContainer { // NOTE: We must `add_functions` to the cache before we add the chunk, as `added_chunks` is // not consulted when retrieving functions from the cache, if we fail to add the functions to // the cache, we will not see them show up in the UI or when matching. - self.cache.add_functions(target, source, functions)?; + self.cache + .write() + .map_err(|e| ContainerError::Custom(format!("Cache write error: {}", e)))? + .add_functions(target, source, functions)?; let signature_chunk = SignatureChunk::new(functions).ok_or( ContainerError::CorruptedData("signature chunk failed to validate"), )?; @@ -352,11 +385,14 @@ impl Container for NetworkContainer { functions: &[Function], ) -> ContainerResult<()> { // TODO: Wont persist, need to add remote removal. - self.cache.remove_functions(target, source, functions) + self.cache + .write() + .map_err(|e| ContainerError::Custom(format!("Cache write error: {}", e)))? + .remove_functions(target, source, functions) } fn fetch_functions( - &mut self, + &self, target: &Target, tags: &[SourceTag], functions: &[FunctionGUID], @@ -376,14 +412,20 @@ impl Container for NetworkContainer { } fn sources_with_type_guid(&self, guid: &TypeGUID) -> ContainerResult> { - self.cache.sources_with_type_guid(guid) + self.cache + .read() + .map_err(|e| ContainerError::Custom(format!("Cache read error: {}", e)))? + .sources_with_type_guid(guid) } fn sources_with_type_guids( &self, guids: &[TypeGUID], ) -> ContainerResult>> { - self.cache.sources_with_type_guids(guids) + self.cache + .read() + .map_err(|e| ContainerError::Custom(format!("Cache read error: {}", e)))? + .sources_with_type_guids(guids) } fn type_guids_with_name( @@ -391,11 +433,17 @@ impl Container for NetworkContainer { source: &SourceId, name: &str, ) -> ContainerResult> { - self.cache.type_guids_with_name(source, name) + self.cache + .read() + .map_err(|e| ContainerError::Custom(format!("Cache read error: {}", e)))? + .type_guids_with_name(source, name) } fn type_with_guid(&self, source: &SourceId, guid: &TypeGUID) -> ContainerResult> { - self.cache.type_with_guid(source, guid) + self.cache + .read() + .map_err(|e| ContainerError::Custom(format!("Cache read error: {}", e)))? + .type_with_guid(source, guid) } fn sources_with_function_guid( @@ -403,7 +451,10 @@ impl Container for NetworkContainer { target: &Target, guid: &FunctionGUID, ) -> ContainerResult> { - self.cache.sources_with_function_guid(target, guid) + self.cache + .read() + .map_err(|e| ContainerError::Custom(format!("Cache read error: {}", e)))? + .sources_with_function_guid(target, guid) } fn sources_with_function_guids( @@ -411,7 +462,10 @@ impl Container for NetworkContainer { target: &Target, guids: &[FunctionGUID], ) -> ContainerResult>> { - self.cache.sources_with_function_guids(target, guids) + self.cache + .read() + .map_err(|e| ContainerError::Custom(format!("Cache read error: {}", e)))? + .sources_with_function_guids(target, guids) } fn functions_with_guid( @@ -420,7 +474,10 @@ impl Container for NetworkContainer { source: &SourceId, guid: &FunctionGUID, ) -> ContainerResult> { - self.cache.functions_with_guid(target, source, guid) + self.cache + .read() + .map_err(|e| ContainerError::Custom(format!("Cache read error: {}", e)))? + .functions_with_guid(target, source, guid) } fn search(&self, query: &ContainerSearchQuery) -> ContainerResult { diff --git a/plugins/warp/src/plugin/ffi/container.rs b/plugins/warp/src/plugin/ffi/container.rs index a107a78538..79f45dcb4c 100644 --- a/plugins/warp/src/plugin/ffi/container.rs +++ b/plugins/warp/src/plugin/ffi/container.rs @@ -220,7 +220,7 @@ pub unsafe extern "C" fn BNWARPContainerFetchFunctions( count: usize, ) { let arc_container = ManuallyDrop::new(Arc::from_raw(container)); - let Ok(mut container) = arc_container.write() else { + let Ok(container) = arc_container.read() else { return; }; @@ -246,7 +246,7 @@ pub unsafe extern "C" fn BNWARPContainerGetSources( count: *mut usize, ) -> *mut BNWARPSource { let arc_container = ManuallyDrop::new(Arc::from_raw(container)); - let Ok(container) = arc_container.write() else { + let Ok(container) = arc_container.read() else { return std::ptr::null_mut(); }; From c092a37f88dd83a027e6ce9ca998b8ddde0c87d5 Mon Sep 17 00:00:00 2001 From: Mason Reed Date: Tue, 3 Mar 2026 07:57:33 -0800 Subject: [PATCH 2/4] [WARP] Server-side constraint matching Reduce networked functions by constraining on the returned set of functions on the server --- plugins/warp/api/python/warp.py | 10 ++++-- plugins/warp/api/warp.cpp | 9 +++-- plugins/warp/api/warp.h | 2 +- plugins/warp/api/warpcore.h | 2 +- plugins/warp/src/container.rs | 5 +++ plugins/warp/src/container/network.rs | 38 ++++++++++++-------- plugins/warp/src/container/network/client.rs | 26 ++++++++++---- plugins/warp/src/plugin/ffi/container.rs | 8 +++-- plugins/warp/src/plugin/settings.rs | 6 ++-- plugins/warp/src/plugin/workflow.rs | 27 +++++++++++--- plugins/warp/ui/shared/fetchdialog.cpp | 34 +++++------------- plugins/warp/ui/shared/fetchdialog.h | 3 +- plugins/warp/ui/shared/fetcher.cpp | 28 ++++++++++++--- plugins/warp/ui/shared/misc.h | 4 +-- 14 files changed, 132 insertions(+), 70 deletions(-) diff --git a/plugins/warp/api/python/warp.py b/plugins/warp/api/python/warp.py index a699ebbcb7..2ba1c68189 100644 --- a/plugins/warp/api/python/warp.py +++ b/plugins/warp/api/python/warp.py @@ -426,11 +426,17 @@ 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))() @@ -438,7 +444,7 @@ def fetch_functions(self, target: WarpTarget, guids: List[FunctionGUID], source_ 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() diff --git a/plugins/warp/api/warp.cpp b/plugins/warp/api/warp.cpp index c051e3c78f..debbbc00a7 100644 --- a/plugins/warp/api/warp.cpp +++ b/plugins/warp/api/warp.cpp @@ -353,7 +353,7 @@ bool Container::RemoveTypes(const Source &source, const std::vector &g return result; } -void Container::FetchFunctions(const Target &target, const std::vector &guids, const std::vector &tags) const +void Container::FetchFunctions(const Target &target, const std::vector &guids, const std::vector &tags, const std::vector &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 Container::GetSourcesWithFunctionGUID(const Target& target, const FunctionGUID &guid) const diff --git a/plugins/warp/api/warp.h b/plugins/warp/api/warp.h index 01f4aaee8e..ccb4da9706 100644 --- a/plugins/warp/api/warp.h +++ b/plugins/warp/api/warp.h @@ -409,7 +409,7 @@ namespace Warp { bool RemoveTypes(const Source &source, const std::vector &guids) const; - void FetchFunctions(const Target &target, const std::vector &guids, const std::vector &tags = {}) const; + void FetchFunctions(const Target &target, const std::vector &guids, const std::vector &tags = {}, const std::vector &constraints = {}) const; std::vector GetSourcesWithFunctionGUID(const Target &target, const FunctionGUID &guid) const; diff --git a/plugins/warp/api/warpcore.h b/plugins/warp/api/warpcore.h index c52485aace..21ef105ba5 100644 --- a/plugins/warp/api/warpcore.h +++ b/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); diff --git a/plugins/warp/src/container.rs b/plugins/warp/src/container.rs index 79c243b3ee..4feed24cac 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 307b44fa1f..ba2db10c19 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>, - /// 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>>, /// Populated when connecting to the server, this is used to determine which sources are writable. /// @@ -110,7 +111,7 @@ impl NetworkContainer { guids: &[FunctionGUID], ) -> HashMap> { 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(()) diff --git a/plugins/warp/src/container/network/client.rs b/plugins/warp/src/container/network/client.rs index 0b772dcc44..f0c3979125 100644 --- a/plugins/warp/src/container/network/client.rs +++ b/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) -> 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, source_tags: &[SourceTag], guids: &[FunctionGUID], + constraints: &[ConstraintGUID], ) -> serde_json::Value { - let guids_str: Vec = guids.iter().map(|g| g.to_string()).collect(); + let guids_str: HashSet = 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 = + constraints.iter().map(|g| g.to_string()).collect(); + body["constraints"] = json!(constraint_guids_str); + } body } @@ -244,13 +253,13 @@ impl NetworkClient { target: Option, source: Option, guids: &[FunctionGUID], + constraints: &[ConstraintGUID], ) -> Result, 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>, 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 79f45dcb4c..aa0c37d39a 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 6469be0f18..896bee7f8a 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 1f4ef1016f..2b318edf5b 100644 --- a/plugins/warp/src/plugin/workflow.rs +++ b/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 = 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, + ); } } }); diff --git a/plugins/warp/ui/shared/fetchdialog.cpp b/plugins/warp/ui/shared/fetchdialog.cpp index 328f00c253..44d7c2375d 100644 --- a/plugins/warp/ui/shared/fetchdialog.cpp +++ b/plugins/warp/ui/shared/fetchdialog.cpp @@ -67,12 +67,6 @@ WarpFetchDialog::WarpFetchDialog(BinaryViewRef bv, std::shared_ptr for (const auto& t : GetAllowedTagsFromView(m_bv)) AddListItem(m_tagsList, QString::fromStdString(t)); - // Batch size and matcher checkbox - m_batchSize = new QSpinBox(this); - m_batchSize->setRange(10, 1000); - m_batchSize->setValue(GetBatchSizeFromView(m_bv)); - m_batchSize->setToolTip("Number of functions to fetch in each batch"); - m_rerunMatcher = new QCheckBox("Re-run matcher after fetch", this); m_rerunMatcher->setChecked(true); @@ -83,7 +77,6 @@ WarpFetchDialog::WarpFetchDialog(BinaryViewRef bv, std::shared_ptr form->addRow(new QLabel("Container: "), m_containerCombo); form->addRow(new QLabel("Allowed Tags: "), tagWrapper); - form->addRow(new QLabel("Batch Size: "), m_batchSize); form->addRow(m_rerunMatcher); form->addRow(m_clearProcessed); @@ -144,7 +137,6 @@ void WarpFetchDialog::onAccept() if (idx > 0) // 0 == All Containers containerIndex = static_cast(idx - 1); - const auto batch = static_cast(m_batchSize->value()); const bool rerun = m_rerunMatcher->isChecked(); const auto tags = collectTags(); @@ -155,7 +147,7 @@ void WarpFetchDialog::onAccept() m_fetchProcessor->ClearProcessed(); // Execute the network fetch in batches - runBatchedFetch(containerIndex, tags, batch, rerun); + runBatchedFetch(containerIndex, tags, rerun); accept(); } @@ -169,7 +161,7 @@ void WarpFetchDialog::onReject() } void WarpFetchDialog::runBatchedFetch(const std::optional& containerIndex, - const std::vector& allowedTags, size_t batchSize, bool rerunMatcher) + const std::vector& allowedTags, bool rerunMatcher) { if (!m_bv) return; @@ -177,42 +169,34 @@ void WarpFetchDialog::runBatchedFetch(const std::optional& containerInde std::vector> funcs = m_bv->GetAnalysisFunctionList(); if (funcs.empty()) return; - const size_t totalFuncs = funcs.size(); - const size_t totalBatches = (totalFuncs + batchSize - 1) / batchSize; // Create a background task to show progress in the UI Ref task = - new BackgroundTask("Fetching WARP functions (0 / " + std::to_string(totalBatches) + ")", false); + new BackgroundTask("Fetching WARP functions (0 / " + std::to_string(funcs.size()) + ")", true); auto fetcher = m_fetchProcessor; auto bv = m_bv; // TODO: Too many captures in this thing lol. WorkerInteractiveEnqueue( - [fetcher, bv, funcs = std::move(funcs), batchSize, rerunMatcher, task, allowedTags]() mutable { + [fetcher, bv, funcs = std::move(funcs), rerunMatcher, task, allowedTags]() mutable { + const auto batchSize = GetBatchSizeFromView(bv); size_t processed = 0; - size_t batchIndex = 0; - while (processed < funcs.size()) { + if (task->IsCancelled()) + break; const size_t remaining = funcs.size() - processed; const size_t thisBatchCount = std::min(batchSize, remaining); - for (size_t i = 0; i < thisBatchCount; ++i) fetcher->AddPendingFunction(funcs[processed + i]); - fetcher->FetchPendingFunctions(allowedTags); - - ++batchIndex; processed += thisBatchCount; - - task->SetProgressText("Fetching WARP functions (" + std::to_string(batchIndex) + " / " - + std::to_string((funcs.size() + batchSize - 1) / batchSize) + ")"); + task->SetProgressText("Fetching WARP functions (" + std::to_string(processed) + " / " + std::to_string(funcs.size()) + ")"); } task->Finish(); - // TODO: Print how long it took? - Logger("WARP Fetcher").LogInfo("Finished fetching WARP functions..."); + Logger("WARP Fetcher").LogInfo("Finished fetching WARP functions in %d seconds...", task->GetRuntimeSeconds()); if (rerunMatcher && bv) Warp::RunMatcher(*bv); diff --git a/plugins/warp/ui/shared/fetchdialog.h b/plugins/warp/ui/shared/fetchdialog.h index c642d31e40..72b8c57e42 100644 --- a/plugins/warp/ui/shared/fetchdialog.h +++ b/plugins/warp/ui/shared/fetchdialog.h @@ -22,7 +22,6 @@ class WarpFetchDialog : public QDialog QPushButton* m_removeTagBtn; QPushButton* m_resetTagBtn; - QSpinBox* m_batchSize; QCheckBox* m_rerunMatcher; QCheckBox* m_clearProcessed; @@ -51,7 +50,7 @@ private slots: std::vector collectTags() const; void runBatchedFetch(const std::optional& containerIndex, const std::vector& allowedTags, - size_t batchSize, bool rerunMatcher); + bool rerunMatcher); }; void RegisterWarpFetchFunctionsCommand(); diff --git a/plugins/warp/ui/shared/fetcher.cpp b/plugins/warp/ui/shared/fetcher.cpp index 51ab018c65..2eab910a9f 100644 --- a/plugins/warp/ui/shared/fetcher.cpp +++ b/plugins/warp/ui/shared/fetcher.cpp @@ -63,22 +63,40 @@ void WarpFetcher::FetchPendingFunctions(const std::vector& allo // Because we must fetch for a single target we map the function guids to the associated platform to perform fetches // for each. - std::map> platformMappedGuids; + std::map> platformMappedGuidSet; + std::map> platformMappedConstraintSet; for (const auto& func : requests) { - const auto guid = Warp::GetAnalysisFunctionGUID(*func); - if (!guid.has_value()) + const auto warpFunc = Warp::Function::Get(*func); + if (!warpFunc) continue; auto platform = func->GetPlatform(); - platformMappedGuids[platform].push_back(guid.value()); + platformMappedGuidSet[platform].insert(warpFunc->GetGUID()); + + // We want to keep track of the guids so we can constrain the server response to only return functions with any of them. + const auto constraints = warpFunc->GetConstraints(); + std::vector constraintGuids; + constraintGuids.reserve(constraints.size()); + for (const auto& constraint : constraints) + constraintGuids.push_back(constraint.guid); + platformMappedConstraintSet[platform].insert(constraintGuids.begin(), constraintGuids.end()); } + std::map> platformMappedGuids; + for (const auto& [platform, guids] : platformMappedGuidSet) + platformMappedGuids[platform] = std::vector(guids.begin(), guids.end()); + + // We keep them in the set above so we don't duplicate a bunch for functions with the same set of constraint guids. + std::map> platformMappedConstraints; + for (const auto& [platform, guids] : platformMappedConstraintSet) + platformMappedConstraints[platform] = std::vector(guids.begin(), guids.end()); + for (const auto& [platform, guids] : platformMappedGuids) { m_logger->LogDebugF("Fetching {} functions for platform {}", guids.size(), platform->GetName()); auto target = Warp::Target::FromPlatform(*platform); for (const auto& container : Warp::Container::All()) - container->FetchFunctions(*target, guids, allowedTags); + container->FetchFunctions(*target, guids, allowedTags, platformMappedConstraints[platform]); std::lock_guard lock(m_requestMutex); for (const auto& guid : guids) diff --git a/plugins/warp/ui/shared/misc.h b/plugins/warp/ui/shared/misc.h index 97adcab00e..92e27fd88a 100644 --- a/plugins/warp/ui/shared/misc.h +++ b/plugins/warp/ui/shared/misc.h @@ -138,10 +138,10 @@ inline void SetTagsToView(const BinaryViewRef& view, const std::vectorSet(ALLOWED_TAGS_SETTING, tags, view); } -inline int GetBatchSizeFromView(const BinaryViewRef& view) +inline size_t GetBatchSizeFromView(const BinaryViewRef& view) { auto settings = BinaryNinja::Settings::Instance(); if (!settings->Contains(BATCH_SIZE_SETTING)) - return 100; + return 10000; return settings->Get(BATCH_SIZE_SETTING, view); } \ No newline at end of file From 5f5c91e5c6d8a44a54a6f8bedbdcfbc7fa6915c8 Mon Sep 17 00:00:00 2001 From: Mason Reed Date: Tue, 3 Mar 2026 09:15:02 -0800 Subject: [PATCH 3/4] [WARP] Update the selected sidebar function when refocusing Should fix issue where opening the sidebar for the first time will not show anything in the selected function until the user clicks in the view frame --- plugins/warp/ui/plugin.cpp | 7 ++++++- plugins/warp/ui/plugin.h | 2 ++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/plugins/warp/ui/plugin.cpp b/plugins/warp/ui/plugin.cpp index 66c80ba394..9421cc02a4 100644 --- a/plugins/warp/ui/plugin.cpp +++ b/plugins/warp/ui/plugin.cpp @@ -227,7 +227,6 @@ void WarpSidebarWidget::notifyViewChanged(ViewFrame* view) if (view == m_currentFrame) return; m_currentFrame = view; - // TODO: We need to set some stuff here prolly. } void WarpSidebarWidget::notifyViewLocationChanged(View* view, const ViewLocation& location) @@ -241,6 +240,12 @@ void WarpSidebarWidget::notifyViewLocationChanged(View* view, const ViewLocation m_currentFunctionWidget->SetCurrentFunction(function); } +void WarpSidebarWidget::focus() +{ + m_currentFunctionWidget->SetCurrentFunction(m_currentFrame->getViewLocation().getFunction()); + SidebarWidget::focus(); +} + WarpSidebarWidgetType::WarpSidebarWidgetType() : SidebarWidgetType(QImage(":/icons/images/warp.png"), "WARP") {} diff --git a/plugins/warp/ui/plugin.h b/plugins/warp/ui/plugin.h index e70b7afafc..0a4c630c04 100644 --- a/plugins/warp/ui/plugin.h +++ b/plugins/warp/ui/plugin.h @@ -39,6 +39,8 @@ class WarpSidebarWidget : public SidebarWidget void notifyViewChanged(ViewFrame*) override; void notifyViewLocationChanged(View*, const ViewLocation&) override; + + void focus() override; }; class WarpSidebarWidgetType : public SidebarWidgetType From 02b5693b967ca9c4cc21d7997d794e232995fdc8 Mon Sep 17 00:00:00 2001 From: Mason Reed Date: Tue, 3 Mar 2026 09:42:14 -0800 Subject: [PATCH 4/4] [WARP] Add a spinner to the possible matches widget while fetching from network A little extra pizzaz --- plugins/warp/ui/matches.cpp | 25 ++++++++++++++++++++++--- plugins/warp/ui/matches.h | 4 +++- plugins/warp/ui/plugin.cpp | 2 +- 3 files changed, 26 insertions(+), 5 deletions(-) diff --git a/plugins/warp/ui/matches.cpp b/plugins/warp/ui/matches.cpp index 289c110244..cede06ab8b 100644 --- a/plugins/warp/ui/matches.cpp +++ b/plugins/warp/ui/matches.cpp @@ -11,7 +11,7 @@ #include "warp.h" #include "shared/misc.h" -WarpCurrentFunctionWidget::WarpCurrentFunctionWidget() +WarpCurrentFunctionWidget::WarpCurrentFunctionWidget(QWidget* parent) : QWidget(parent) { // We must explicitly support no current function. m_current = nullptr; @@ -31,10 +31,27 @@ WarpCurrentFunctionWidget::WarpCurrentFunctionWidget() m_splitter = new QSplitter(Qt::Vertical); m_splitter->setContentsMargins(0, 0, 0, 0); + // Wrap the table and the spinner so that we can overlay the spinner on the table. + QWidget* tableWrapper = new QWidget(m_splitter); + QGridLayout* wrapperLayout = new QGridLayout(tableWrapper); + wrapperLayout->setContentsMargins(0, 0, 0, 0); + // Add a widget to display the matches. - m_tableWidget = new WarpFunctionTableWidget(this); + m_tableWidget = new WarpFunctionTableWidget(tableWrapper); m_tableWidget->setContentsMargins(0, 0, 0, 0); - m_splitter->addWidget(m_tableWidget); + + // Spinner for when we are fetching functions over the network. + m_spinner = new QProgressBar(tableWrapper); + m_spinner->setRange(0, 0); + m_spinner->setTextVisible(false); + m_spinner->setFixedHeight(6); + m_spinner->hide(); + + // The table has no alignment, so it expands to fill the entire cell. + wrapperLayout->addWidget(m_tableWidget, 0, 0); + wrapperLayout->addWidget(m_spinner, 0, 0, Qt::AlignBottom); + + m_splitter->addWidget(tableWrapper); // Add a widget to display the info about the selected function match. m_infoWidget = new WarpFunctionInfoWidget(this); @@ -145,10 +162,12 @@ void WarpCurrentFunctionWidget::SetCurrentFunction(FunctionRef current) if (!m_fetcher->m_requestInProgress.exchange(true)) { BinaryNinja::WorkerPriorityEnqueue([this]() { + QMetaObject::invokeMethod(this, [this] { m_spinner->show(); }, Qt::QueuedConnection); BinaryNinja::Ref bgTask = new BinaryNinja::BackgroundTask("Fetching WARP Functions...", true); const auto allowedTags = GetAllowedTagsFromView(m_current->GetView()); m_fetcher->FetchPendingFunctions(allowedTags); bgTask->Finish(); + QMetaObject::invokeMethod(this, [this] { m_spinner->hide(); }, Qt::QueuedConnection); }); } } diff --git a/plugins/warp/ui/matches.h b/plugins/warp/ui/matches.h index 6a2061efec..1599dd541b 100644 --- a/plugins/warp/ui/matches.h +++ b/plugins/warp/ui/matches.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include "filter.h" #include "render.h" @@ -13,6 +14,7 @@ class WarpCurrentFunctionWidget : public QWidget FunctionRef m_current; QSplitter* m_splitter; + QProgressBar* m_spinner; WarpFunctionTableWidget* m_tableWidget; WarpFunctionInfoWidget* m_infoWidget; @@ -22,7 +24,7 @@ class WarpCurrentFunctionWidget : public QWidget std::shared_ptr m_fetcher; public: - explicit WarpCurrentFunctionWidget(); + explicit WarpCurrentFunctionWidget(QWidget* parent = nullptr); ~WarpCurrentFunctionWidget() override = default; diff --git a/plugins/warp/ui/plugin.cpp b/plugins/warp/ui/plugin.cpp index 9421cc02a4..e99059b98a 100644 --- a/plugins/warp/ui/plugin.cpp +++ b/plugins/warp/ui/plugin.cpp @@ -136,7 +136,7 @@ WarpSidebarWidget::WarpSidebarWidget(BinaryViewRef data) : SidebarWidget("WARP") m_headerWidget->setLayout(headerLayout); QFrame* currentFunctionFrame = new QFrame(this); - m_currentFunctionWidget = new WarpCurrentFunctionWidget(); + m_currentFunctionWidget = new WarpCurrentFunctionWidget(this); QVBoxLayout* currentFunctionLayout = new QVBoxLayout(); currentFunctionLayout->setContentsMargins(0, 0, 0, 0); currentFunctionLayout->setSpacing(0);