[WARP] Enhanced network support

3 files changed, 307 insertions, 4 deletions

diff --git a/plugins/warp/src/plugin/ffi/container.rs b/plugins/warp/src/plugin/ffi/container.rs
index 4de6bab3..72c85366 100644
--- a/plugins/warp/src/plugin/ffi/container.rs
+++ b/plugins/warp/src/plugin/ffi/container.rs
@@ -1,5 +1,7 @@
 use crate::cache::container::cached_containers;
-use crate::container::SourcePath;
+use crate::container::{
+    ContainerSearchItem, ContainerSearchItemKind, ContainerSearchQuery, SourcePath, SourceTag,
+};
 use crate::convert::{from_bn_type, to_bn_type};
 use crate::plugin::ffi::{
     BNWARPContainer, BNWARPFunction, BNWARPFunctionGUID, BNWARPSource, BNWARPTarget, BNWARPTypeGUID,
@@ -12,7 +14,163 @@ use binaryninja::types::Type;
 use binaryninjacore_sys::{BNArchitecture, BNBinaryView, BNType};
 use std::ffi::{c_char, CStr};
 use std::mem::ManuallyDrop;
+use std::ops::Deref;
 use std::sync::Arc;
+use warp::r#type::guid::TypeGUID;
+
+pub type BNWARPContainerSearchQuery = ContainerSearchQuery;
+pub type BNWARPContainerSearchItem = ContainerSearchItem;
+
+#[repr(C)]
+pub enum BNWARPContainerSearchItemKind {
+    Source = 0,
+    Function = 1,
+    Type = 2,
+    Symbol = 3,
+}
+
+#[repr(C)]
+pub struct BNWARPContainerSearchResponse {
+    pub count: usize,
+    pub items: *mut *mut BNWARPContainerSearchItem,
+    pub offset: usize,
+    pub total: usize,
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPNewContainerSearchQuery(
+    query: *mut c_char,
+    offset: *const usize,
+    limit: *const usize,
+    source: *const BNWARPSource,
+    source_tags: *mut *mut c_char,
+    source_tags_count: usize,
+) -> *mut BNWARPContainerSearchQuery {
+    let query_cstr = unsafe { CStr::from_ptr(query) };
+    let Ok(query) = query_cstr.to_str() else {
+        return std::ptr::null_mut();
+    };
+    let mut search_query = ContainerSearchQuery::new(query.to_string());
+    if !offset.is_null() {
+        search_query.offset = Some(*offset);
+    }
+    if !limit.is_null() {
+        search_query.limit = Some(*limit);
+    }
+    if !source.is_null() {
+        search_query.source = Some(*source);
+    }
+    if !source_tags.is_null() {
+        let source_tags_raw = unsafe { std::slice::from_raw_parts(source_tags, source_tags_count) };
+        let source_tags: Vec<SourceTag> = source_tags_raw
+            .iter()
+            .filter_map(|&ptr| CStr::from_ptr(ptr).to_str().ok())
+            .map(|s| s.into())
+            .collect();
+        search_query.tags = source_tags;
+    }
+    Box::into_raw(Box::new(search_query))
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPContainerSearchItemGetKind(
+    item: *mut BNWARPContainerSearchItem,
+) -> BNWARPContainerSearchItemKind {
+    let item = ManuallyDrop::new(Arc::from_raw(item));
+    match &item.kind {
+        ContainerSearchItemKind::Source { .. } => BNWARPContainerSearchItemKind::Source,
+        ContainerSearchItemKind::Function(_) => BNWARPContainerSearchItemKind::Function,
+        ContainerSearchItemKind::Type(_) => BNWARPContainerSearchItemKind::Type,
+        ContainerSearchItemKind::Symbol(_) => BNWARPContainerSearchItemKind::Symbol,
+    }
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPContainerSearchItemGetSource(
+    item: *mut BNWARPContainerSearchItem,
+) -> BNWARPSource {
+    let item = ManuallyDrop::new(Arc::from_raw(item));
+    item.source
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPContainerSearchItemGetType(
+    arch: *mut BNArchitecture,
+    item: *mut BNWARPContainerSearchItem,
+) -> *mut BNType {
+    // NOTE: to convert the type, we must have an architecture.
+    let arch = match !arch.is_null() {
+        true => Some(CoreArchitecture::from_raw(arch)),
+        false => None,
+    };
+
+    let item = ManuallyDrop::new(Arc::from_raw(item));
+    match &item.kind {
+        ContainerSearchItemKind::Source { .. } => std::ptr::null_mut(),
+        ContainerSearchItemKind::Function(func) => {
+            match &func.ty {
+                None => std::ptr::null_mut(),
+                Some(ty) => {
+                    let bn_ty = to_bn_type(arch, &ty);
+                    // NOTE: The type ref has been pre-incremented for the caller.
+                    unsafe { Ref::into_raw(bn_ty) }.handle
+                }
+            }
+        }
+        ContainerSearchItemKind::Type(ty) => {
+            let bn_ty = to_bn_type(arch, &ty);
+            // NOTE: The type ref has been pre-incremented for the caller.
+            unsafe { Ref::into_raw(bn_ty) }.handle
+        }
+        ContainerSearchItemKind::Symbol(_) => std::ptr::null_mut(),
+    }
+}
+
+// NOTE: In the future we should allow for the possibility of this returning a null pointer.
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPContainerSearchItemGetName(
+    item: *mut BNWARPContainerSearchItem,
+) -> *mut c_char {
+    let item = ManuallyDrop::new(Arc::from_raw(item));
+    match &item.kind {
+        ContainerSearchItemKind::Source { path, .. } => {
+            let bn_name = BnString::new(path.to_string());
+            BnString::into_raw(bn_name)
+        }
+        ContainerSearchItemKind::Function(func) => {
+            let bn_name = BnString::new(func.symbol.name.clone());
+            BnString::into_raw(bn_name)
+        }
+        ContainerSearchItemKind::Type(ty) => {
+            // TODO: Maybe un-named types should return std::ptr::null_mut()?
+            let ty_name = ty
+                .name
+                .clone()
+                .unwrap_or_else(|| TypeGUID::from(ty).to_string());
+            let bn_name = BnString::new(ty_name);
+            BnString::into_raw(bn_name)
+        }
+        ContainerSearchItemKind::Symbol(sym) => {
+            let bn_name = BnString::new(sym.name.clone());
+            BnString::into_raw(bn_name)
+        }
+    }
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPContainerSearchItemGetFunction(
+    item: *mut BNWARPContainerSearchItem,
+) -> *mut BNWARPFunction {
+    let item = ManuallyDrop::new(Arc::from_raw(item));
+    match &item.kind {
+        ContainerSearchItemKind::Source { .. } => std::ptr::null_mut(),
+        ContainerSearchItemKind::Function(func) => {
+            Arc::into_raw(Arc::new(func.clone())) as *mut BNWARPFunction
+        }
+        ContainerSearchItemKind::Type(_) => std::ptr::null_mut(),
+        ContainerSearchItemKind::Symbol(_) => std::ptr::null_mut(),
+    }
+}
 
 #[no_mangle]
 pub unsafe extern "C" fn BNWARPGetContainers(count: *mut usize) -> *mut *mut BNWARPContainer {
@@ -39,6 +197,8 @@ pub unsafe extern "C" fn BNWARPContainerGetName(container: *mut BNWARPContainer)
 pub unsafe extern "C" fn BNWARPContainerFetchFunctions(
     container: *mut BNWARPContainer,
     target: *mut BNWARPTarget,
+    source_tags: *mut *mut c_char,
+    source_tags_count: usize,
     guids: *const BNWARPFunctionGUID,
     count: usize,
 ) {
@@ -49,9 +209,16 @@ pub unsafe extern "C" fn BNWARPContainerFetchFunctions(
 
     let target = unsafe { ManuallyDrop::new(Arc::from_raw(target)) };
 
+    let source_tags_raw = unsafe { std::slice::from_raw_parts(source_tags, source_tags_count) };
+    let source_tags: Vec<SourceTag> = source_tags_raw
+        .iter()
+        .filter_map(|&ptr| CStr::from_ptr(ptr).to_str().ok())
+        .map(|s| s.into())
+        .collect();
+
     let guids = unsafe { std::slice::from_raw_parts(guids, count) };
 
-    if let Err(e) = container.fetch_functions(&target, guids) {
+    if let Err(e) = container.fetch_functions(&target, &source_tags, guids) {
         log::error!("Failed to fetch functions: {}", e);
     }
 }
@@ -373,7 +540,7 @@ pub unsafe extern "C" fn BNWARPContainerGetTypeWithGUID(
     let Some(ty) = container.type_with_guid(&source, &guid).unwrap_or_default() else {
         return std::ptr::null_mut();
     };
-    let function_type = to_bn_type(&arch, &ty);
+    let function_type = to_bn_type(Some(arch), &ty);
     // NOTE: The type ref has been pre-incremented for the caller.
     unsafe { Ref::into_raw(function_type) }.handle
 }
@@ -405,6 +572,45 @@ pub unsafe extern "C" fn BNWARPContainerGetTypeGUIDsWithName(
 }
 
 #[no_mangle]
+pub unsafe extern "C" fn BNWARPContainerSearch(
+    container: *mut BNWARPContainer,
+    query: *mut BNWARPContainerSearchQuery,
+) -> *mut BNWARPContainerSearchResponse {
+    let arc_container = ManuallyDrop::new(Arc::from_raw(container));
+    let Ok(container) = arc_container.read() else {
+        return std::ptr::null_mut();
+    };
+
+    let query = unsafe { ManuallyDrop::new(Arc::from_raw(query)) };
+
+    let result = match container.search(&query) {
+        Ok(result) => result,
+        Err(err) => {
+            log::error!("Failed to search container {:?}: {}", query.deref(), err);
+            return std::ptr::null_mut();
+        }
+    };
+
+    let boxed_raw_items: Box<[_]> = result
+        .items
+        .into_iter()
+        .map(Arc::new)
+        .map(Arc::into_raw)
+        .collect();
+    let count = boxed_raw_items.len();
+    // NOTE: Leak the functions to be freed by BNWARPFreeContainerSearchItemList
+    let leaked_raw_items = Box::into_raw(boxed_raw_items) as *mut *mut BNWARPContainerSearchItem;
+    let raw_result = BNWARPContainerSearchResponse {
+        count,
+        items: leaked_raw_items,
+        total: result.total,
+        offset: result.offset,
+    };
+    // NOTE: Leak the result to be freed by BNWARPFreeContainerSearchResult
+    Box::into_raw(Box::new(raw_result))
+}
+
+#[no_mangle]
 pub unsafe extern "C" fn BNWARPNewContainerReference(
     container: *mut BNWARPContainer,
 ) -> *mut BNWARPContainer {
@@ -431,3 +637,62 @@ pub unsafe extern "C" fn BNWARPFreeContainerList(
         BNWARPFreeContainerReference(container);
     }
 }
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPNewContainerSearchQueryReference(
+    query: *mut BNWARPContainerSearchQuery,
+) -> *mut BNWARPContainerSearchQuery {
+    Arc::increment_strong_count(query);
+    query
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPFreeContainerSearchQueryReference(
+    query: *mut BNWARPContainerSearchQuery,
+) {
+    if query.is_null() {
+        return;
+    }
+    Arc::decrement_strong_count(query);
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPNewContainerSearchItemReference(
+    item: *mut BNWARPContainerSearchItem,
+) -> *mut BNWARPContainerSearchItem {
+    Arc::increment_strong_count(item);
+    item
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPFreeContainerSearchItemReference(
+    item: *mut BNWARPContainerSearchItem,
+) {
+    if item.is_null() {
+        return;
+    }
+    Arc::decrement_strong_count(item);
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPFreeContainerSearchItemList(
+    items: *mut *mut BNWARPContainerSearchItem,
+    count: usize,
+) {
+    let items_ptr = std::ptr::slice_from_raw_parts_mut(items, count);
+    let items = unsafe { Box::from_raw(items_ptr) };
+    for item in items {
+        BNWARPFreeContainerSearchItemReference(item);
+    }
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPFreeContainerSearchResponse(
+    response: *mut BNWARPContainerSearchResponse,
+) {
+    if response.is_null() {
+        return;
+    }
+    let response = unsafe { Box::from_raw(response) };
+    BNWARPFreeContainerSearchItemList(response.items, response.count);
+}
diff --git a/plugins/warp/src/plugin/ffi/file.rs b/plugins/warp/src/plugin/ffi/file.rs
new file mode 100644
index 00000000..951b5eb2
--- /dev/null
+++ b/plugins/warp/src/plugin/ffi/file.rs
@@ -0,0 +1,38 @@
+use std::ffi::c_char;
+use std::sync::Arc;
+use warp::WarpFile;
+
+pub type BNWARPFile = WarpFile<'static>;
+
+// TODO: At some point we may want to expose chunks directly. For now we will just enumerate all of them.
+// pub type BNWARPChunk = warp::chunk::Chunk<'static>;
+
+// TODO: From bytes as well.
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPNewFileFromPath(path: *mut c_char) -> *mut BNWARPFile {
+    let path_cstr = unsafe { std::ffi::CStr::from_ptr(path) };
+    let Ok(path) = path_cstr.to_str() else {
+        return std::ptr::null_mut();
+    };
+    let Ok(bytes) = std::fs::read(path) else {
+        return std::ptr::null_mut();
+    };
+    let Some(file) = WarpFile::from_owned_bytes(bytes) else {
+        return std::ptr::null_mut();
+    };
+    Arc::into_raw(Arc::new(file)) as *mut BNWARPFile
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPNewFileReference(file: *mut BNWARPFile) -> *mut BNWARPFile {
+    Arc::increment_strong_count(file);
+    file
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn BNWARPFreeFileReference(file: *mut BNWARPFile) {
+    if file.is_null() {
+        return;
+    }
+    Arc::decrement_strong_count(file);
+}
diff --git a/plugins/warp/src/plugin/ffi/function.rs b/plugins/warp/src/plugin/ffi/function.rs
index b37d613b..e3f6a079 100644
--- a/plugins/warp/src/plugin/ffi/function.rs
+++ b/plugins/warp/src/plugin/ffi/function.rs
@@ -124,7 +124,7 @@ pub unsafe extern "C" fn BNWARPFunctionGetType(
     match &function.ty {
         Some(func_ty) => {
             let arch = analysis_function.arch();
-            let function_type = to_bn_type(&arch, func_ty);
+            let function_type = to_bn_type(Some(arch), func_ty);
             // NOTE: The type ref has been pre-incremented for the caller.
             unsafe { Ref::into_raw(function_type) }.handle
         }