warp wip

This is a partial commit that builds, this is so that mason can go do something else.

Author: emesare

Cargo.lock

@@ -17,6 +17,7 @@ dashmap = "6.1"
 walkdir = "2.5"
 rfd = "0.15"
 serde_json = "1.0"
+uuid = { version = "1.12.0", features = ["v4"] }
 
 # For sigem
 env_logger = { version = "0.11", optional = true }

plugins/warp/Cargo.toml

@@ -4,33 +4,36 @@ use criterion::{criterion_group, criterion_main, Criterion};
 use std::path::PathBuf;
 use warp_ninja::convert::from_bn_type;
 
+// These are the target files present in OUT_DIR
+// Add the files to fixtures/bin
+static TARGET_FILES: [&str; 1] = ["atox.obj"];
+
 pub fn type_conversion_benchmark(c: &mut Criterion) {
     let session = Session::new().expect("Failed to initialize session");
     let out_dir = env!("OUT_DIR").parse::<PathBuf>().unwrap();
-    for entry in std::fs::read_dir(out_dir).expect("Failed to read OUT_DIR") {
-        let entry = entry.expect("Failed to read directory entry");
-        let path = entry.path();
-        if path.is_file() {
-            if let Some(bv) = session.load(path.to_str().unwrap()) {
-                let functions = bv.functions();
-                c.bench_function("type conversion all functions", |b| {
-                    b.iter(|| {
-                        for func in &functions {
-                            from_bn_type(&bv, &func.function_type(), u8::MAX);
-                        }
-                    })
-                });
+    for file_name in TARGET_FILES {
+        let path = out_dir.join(file_name);
+        let view = session.load(&path).expect("Failed to load view");
+
+        // Bench function types.
+        let functions: Vec<_> = view.functions().iter().map(|f| f.to_owned()).collect();
+        c.bench_function("type conversion all functions", |b| {
+            b.iter(|| {
+                for func in &functions {
+                    from_bn_type(&view, &func.function_type(), u8::MAX);
+                }
+            })
+        });
 
-                let types = bv.types();
-                c.bench_function("type conversion all types", |b| {
-                    b.iter(|| {
-                        for ty in &types {
-                            from_bn_type(&bv, &ty.ty, u8::MAX);
-                        }
-                    })
-                });
-            }
-        }
+        // Bench view types.
+        let types = view.types().to_vec();
+        c.bench_function("type conversion all types", |b| {
+            b.iter(|| {
+                for ty in &types {
+                    from_bn_type(&view, &ty.ty, u8::MAX);
+                }
+            })
+        });
     }
 }
 

plugins/warp/benches/convert.rs

@@ -2,43 +2,45 @@ use binaryninja::binary_view::BinaryViewExt;
 use binaryninja::headless::Session;
 use criterion::{criterion_group, criterion_main, Criterion};
 use rayon::prelude::*;
+use std::path::PathBuf;
 use warp_ninja::build_function;
 use warp_ninja::cache::FunctionCache;
 
+// These are the target files present in OUT_DIR
+// Add the files to fixtures/bin
+static TARGET_FILES: [&str; 1] = ["atox.obj"];
+
 pub fn function_benchmark(c: &mut Criterion) {
     let session = Session::new().expect("Failed to initialize session");
-    let bv = session.load(env!("TEST_BIN_LIBRARY_OBJ")).unwrap();
-    let functions = bv.functions();
-    assert_eq!(functions.len(), 6);
-    let mut function_iter = functions.into_iter();
-    let first_function = function_iter.next().unwrap();
-
-    c.bench_function("signature first function", |b| {
-        b.iter(|| {
-            let _ = build_function(&first_function, &first_function.low_level_il().unwrap());
-        })
-    });
+    let out_dir = env!("OUT_DIR").parse::<PathBuf>().unwrap();
+    for file_name in TARGET_FILES {
+        let path = out_dir.join(file_name);
+        let view = session.load(&path).expect("Failed to load view");
 
-    c.bench_function("signature all functions", |b| {
-        b.iter(|| {
-            for func in &functions {
-                let _ = build_function(&func, &func.low_level_il().unwrap());
-            }
-        })
-    });
+        let functions: Vec<_> = view.functions().iter().map(|f| f.to_owned()).collect();
+        // Bench all functions sequentially
+        c.bench_function("signature all functions", |b| {
+            b.iter(|| {
+                for func in &functions {
+                    let _ = build_function(&func, &func.low_level_il().unwrap());
+                }
+            })
+        });
 
-    let cache = FunctionCache::default();
-    c.bench_function("signature all functions rayon", |b| {
-        b.iter(|| {
-            functions
-                .par_iter()
-                .map_with(cache.clone(), |par_cache, func| {
-                    let llil = func.low_level_il().ok()?;
-                    Some(par_cache.function(func.as_ref(), llil.as_ref()))
-                })
-                .collect::<Vec<_>>()
-        })
-    });
+        // Bench all functions in parallel
+        let cache = FunctionCache::default();
+        c.bench_function("signature all functions rayon", |b| {
+            b.iter(|| {
+                functions
+                    .par_iter()
+                    .map_with(cache.clone(), |par_cache, func| {
+                        let llil = func.low_level_il().ok()?;
+                        Some(par_cache.function(func.as_ref(), llil.as_ref()))
+                    })
+                    .collect::<Vec<_>>()
+            })
+        });
+    }
 }
 
 criterion_group!(benches, function_benchmark);

plugins/warp/benches/function.rs

@@ -1,22 +1,31 @@
 use binaryninja::binary_view::BinaryViewExt;
 use binaryninja::headless::Session;
 use criterion::{criterion_group, criterion_main, Criterion};
+use std::path::PathBuf;
 use warp_ninja::function_guid;
 
+// These are the target files present in OUT_DIR
+// Add the files to fixtures/bin
+static TARGET_FILES: [&str; 1] = ["atox.obj"];
+
 pub fn guid_benchmark(c: &mut Criterion) {
     let session = Session::new().expect("Failed to initialize session");
-    let bv = session.load(env!("TEST_BIN_LIBRARY_OBJ")).unwrap();
-    let functions = bv.functions();
-    assert_eq!(functions.len(), 6);
-    let mut function_iter = functions.into_iter();
-    let first_function = function_iter.next().unwrap();
+    let out_dir = env!("OUT_DIR").parse::<PathBuf>().unwrap();
+    for file_name in TARGET_FILES {
+        let path = out_dir.join(file_name);
+        let view = session.load(&path).expect("Failed to load view");
 
-    c.bench_function("function guid", |b| {
-        b.iter(|| {
-            let llil = first_function.low_level_il().unwrap();
-            function_guid(&first_function, &llil);
-        })
-    });
+        let functions: Vec<_> = view.functions().iter().map(|f| f.to_owned()).collect();
+        // Bench all functions sequentially
+        c.bench_function("guid all functions", |b| {
+            b.iter(|| {
+                for func in &functions {
+                    let llil = func.low_level_il().unwrap();
+                    function_guid(&func, &llil);
+                }
+            })
+        });
+    }
 }
 
 criterion_group!(benches, guid_benchmark);

plugins/warp/benches/guid.rs

@@ -1,3 +1,4 @@
+pub mod container;
 pub mod function;
 pub mod guid;
 pub mod type_reference;

plugins/warp/fixtures/bin/random.sbin

@@ -0,0 +1,44 @@
+use crate::cache::ViewID;
+use crate::container::Container;
+use binaryninja::binary_view::BinaryView;
+use dashmap::DashMap;
+use std::sync::OnceLock;
+
+pub static CONTAINER_CACHE: OnceLock<DashMap<ViewID, ContainerCache>> = OnceLock::new();
+
+pub fn cached_containers(view: &BinaryView, f: impl Fn(&dyn Container)) {
+    let view_id = ViewID::from(view.as_ref());
+    let containers_cache = CONTAINER_CACHE.get_or_init(Default::default);
+    if let Some(cache) = containers_cache.get(&view_id) {
+        for container in cache.containers() {
+            f(container);
+        }
+    }
+}
+
+// TODO: The static lifetime here is a little wierd... (we need it to Box)
+pub fn add_cached_container(view: &BinaryView, container: impl Container + 'static) {
+    let view_id = ViewID::from(view.as_ref());
+    let containers_cache = CONTAINER_CACHE.get_or_init(Default::default);
+    let mut cache = containers_cache
+        .entry(view_id)
+        .or_insert_with(Default::default);
+    // TODO: What happens if we have already added the container?
+    // TODO: We need to replace container.
+    cache.add_container(Box::new(container));
+}
+
+#[derive(Default)]
+pub struct ContainerCache {
+    pub cache: Vec<Box<dyn Container>>,
+}
+
+impl ContainerCache {
+    pub fn add_container(&mut self, container: Box<dyn Container>) {
+        self.cache.push(container);
+    }
+
+    pub fn containers(&self) -> impl Iterator<Item = &dyn Container> {
+        self.cache.iter().map(|c| c.as_ref())
+    }
+}

plugins/warp/src/cache.rs

@@ -29,7 +29,6 @@ pub fn clear_function_cache(view: &BinaryView) {
     }
 }
 
-// TODO: The cached function match here might need to remove the function callback.
 pub fn cached_function_match<F>(function: &BNFunction, f: F) -> Option<Function>
 where
     F: Fn() -> Option<Function>,

plugins/warp/src/cache/container.rs

@@ -0,0 +1,72 @@
+use std::collections::HashMap;
+use std::fmt::{Debug, Display};
+use std::hash::Hash;
+use uuid::Uuid;
+use warp::r#type::guid::TypeGUID;
+use warp::r#type::Type;
+use warp::signature::function::{Function, FunctionGUID};
+
+pub mod disk;
+pub mod memory;
+pub mod network;
+
+/// Represents the ID for a single container source.
+///
+/// A source is used to relate types and functions separate from the container. This allows
+/// type name lookups and for containers which are bandwidth sensitive to exist.
+///
+/// An example of a bandwidth sensitive container would be a container which pulls functions over
+/// the network instead of from memory or disk.
+#[derive(Clone, Debug, Eq, PartialEq, Hash, Copy)]
+pub struct SourceId(Uuid);
+
+impl SourceId {
+    pub fn new() -> Self {
+        Self(Uuid::new_v4())
+    }
+}
+
+/// Storage for WARP functions.
+///
+/// Containers are made up of sources, see [`SourceId`] for more details.
+pub trait Container: Send + Sync + Display + Debug {
+    /// Get the sources that contain a type with the given [`TypeGUID`].
+    fn sources_with_type_guid(&self, guid: &TypeGUID) -> Vec<&SourceId>;
+
+    /// Plural version of [`Container::sources_with_type_guid`].
+    ///
+    /// Each source will have a list of the containing GUID's so that when looking up a source you give
+    /// it only the GUID's that it knows about, for networking this means cutting down traffic significantly.
+    fn sources_with_type_guids<'a>(
+        &'a self,
+        guids: &'a [TypeGUID],
+    ) -> HashMap<&'a TypeGUID, Vec<&'a SourceId>>;
+
+    /// Retrieve all [`TypeGUID`]'s with the given name.
+    fn type_guids_with_name(&self, source: &SourceId, name: &str) -> Vec<TypeGUID>;
+
+    fn type_with_guid(&self, source: &SourceId, guid: &TypeGUID) -> Option<Type>;
+
+    fn has_type_with_guid(&self, source: &SourceId, guid: &TypeGUID) -> bool {
+        self.type_with_guid(source, guid).is_some()
+    }
+
+    /// Get the sources that contain functions with the given [`FunctionGUID`].
+    fn sources_with_function_guid(&self, guid: &FunctionGUID) -> Vec<&SourceId>;
+
+    // TODO: Allocating with Vec is not good.
+    /// Plural version of [`Container::sources_with_function_guid`].
+    ///
+    /// Each source will have a list of the containing GUID's so that when looking up a source you give
+    /// it only the GUID's that it knows about, for networking this means cutting down traffic significantly.
+    fn sources_with_function_guids<'a>(
+        &'a self,
+        guids: &'a [FunctionGUID],
+    ) -> HashMap<&'a FunctionGUID, Vec<&'a SourceId>>;
+
+    fn functions_with_guid(&self, source: &SourceId, guid: &FunctionGUID) -> Vec<Function>;
+
+    fn has_function_with_guid(&self, source: &SourceId, guid: &FunctionGUID) -> bool {
+        !self.functions_with_guid(source, guid).is_empty()
+    }
+}