Implement function level metadata

2 files changed, 112 insertions, 0 deletions

diff --git a/rust/src/function.rs b/rust/src/function.rs
index 60c8970c..89ade8f8 100644
--- a/rust/src/function.rs
+++ b/rust/src/function.rs
@@ -43,6 +43,7 @@ use crate::high_level_il::HighLevelILFunction;
 use crate::language_representation::CoreLanguageRepresentationFunction;
 use crate::low_level_il::LowLevelILRegularFunction;
 use crate::medium_level_il::MediumLevelILFunction;
+use crate::metadata::Metadata;
 use crate::variable::{
     IndirectBranchInfo, MergedVariable, NamedVariableWithType, RegisterValue, RegisterValueType,
     StackVariableReference, Variable,
@@ -2428,6 +2429,48 @@ impl Function {
         assert!(!result.is_null());
         unsafe { Array::new(result, count, ()) }
     }
+
+    pub fn query_metadata(&self, key: &str) -> Option<Ref<Metadata>> {
+        let key = key.to_cstr();
+        let value: *mut BNMetadata = unsafe { BNFunctionQueryMetadata(self.handle, key.as_ptr()) };
+        if value.is_null() {
+            None
+        } else {
+            Some(unsafe { Metadata::ref_from_raw(value) })
+        }
+    }
+
+    pub fn get_metadata(&self) -> Option<Ref<Metadata>> {
+        let value: *mut BNMetadata = unsafe { BNFunctionGetMetadata(self.handle) };
+        if value.is_null() {
+            None
+        } else {
+            Some(unsafe { Metadata::ref_from_raw(value) })
+        }
+    }
+
+    pub fn get_auto_metadata(&self) -> Option<Ref<Metadata>> {
+        let value: *mut BNMetadata = unsafe { BNFunctionGetAutoMetadata(self.handle) };
+        if value.is_null() {
+            None
+        } else {
+            Some(unsafe { Metadata::ref_from_raw(value) })
+        }
+    }
+
+    pub fn store_metadata<V>(&self, key: &str, value: V, is_auto: bool)
+    where
+        V: Into<Ref<Metadata>>,
+    {
+        let md = value.into();
+        let key = key.to_cstr();
+        unsafe { BNFunctionStoreMetadata(self.handle, key.as_ptr(), md.as_ref().handle, is_auto) };
+    }
+
+    pub fn remove_metadata(&self, key: &str) {
+        let key = key.to_cstr();
+        unsafe { BNFunctionRemoveMetadata(self.handle, key.as_ptr()) };
+    }
 }
 
 impl Debug for Function {
diff --git a/rust/tests/function.rs b/rust/tests/function.rs
new file mode 100644
index 00000000..455ba752
--- /dev/null
+++ b/rust/tests/function.rs
@@ -0,0 +1,69 @@
+use binaryninja::binary_view::BinaryViewExt;
+use binaryninja::headless::Session;
+use binaryninja::metadata::Metadata;
+use binaryninja::rc::Ref;
+use std::path::PathBuf;
+
+#[test]
+fn store_and_query_function_metadata() {
+    let _session = Session::new().expect("Failed to initialize session");
+    let out_dir = env!("OUT_DIR").parse::<PathBuf>().unwrap();
+    let view = binaryninja::load(out_dir.join("atox.obj")).expect("Failed to create view");
+    let func = view
+        .entry_point_function()
+        .expect("Failed to get entry point function");
+
+    // Store key/value pairs to user and auto metadata
+    func.store_metadata("one", "one", false);
+    func.store_metadata("two", 2 as u64, true);
+    func.store_metadata("three", "three", true);
+    func.remove_metadata("three");
+
+    // Assert that we can query from both user and auto metadata
+    assert_eq!(
+        func.query_metadata("one")
+            .expect("Failed to query key \"one\"")
+            .get_string()
+            .unwrap()
+            .to_string_lossy(),
+        "one"
+    );
+    assert_eq!(
+        func.query_metadata("two")
+            .expect("Failed to query key \"two\"")
+            .get_unsigned_integer()
+            .unwrap(),
+        2
+    );
+    assert!(
+        func.query_metadata("three") == None,
+        "Query for key \"three\" returned a value"
+    );
+
+    // Assert that user metadata only includes key/values from user data (not auto) and vice-versa
+    let user_metadata = func.get_metadata().expect("Failed to query user metadata");
+    assert_eq!(
+        user_metadata
+            .get("one")
+            .expect("Failed to query key \"one\" from user metadata")
+            .expect("User metadata ref is None")
+            .get_string()
+            .unwrap()
+            .to_string_lossy(),
+        "one"
+    );
+    assert_eq!(user_metadata.get("two"), Ok(None));
+    let auto_metadata = func
+        .get_auto_metadata()
+        .expect("Failed to query auto metadata");
+    assert_eq!(
+        auto_metadata
+            .get("two")
+            .expect("Failed to query key \"two\" from auto metadata")
+            .expect("Auto metadata ref is None")
+            .get_unsigned_integer()
+            .unwrap(),
+        2
+    );
+    assert_eq!(auto_metadata.get("one"), Ok(None));
+}