Implement function level metadata

Author: zznop

binaryninjaapi.h

@@ -11458,6 +11458,12 @@ namespace BinaryNinja {
 		void CollapseRegion(uint64_t hash);
 		void ExpandRegion(uint64_t hash);
 		void ExpandAll();
+
+		void StoreMetadata(const std::string& key, Ref<Metadata> value, bool isAuto = false);
+		Ref<Metadata> QueryMetadata(const std::string& key);
+		Ref<Metadata> GetMetadata();
+		Ref<Metadata> GetAutoMetadata();
+		void RemoveMetadata(const std::string& key);
 	};
 
 	/*!

binaryninjacore.h

@@ -37,7 +37,7 @@
 // Current ABI version for linking to the core. This is incremented any time
 // there are changes to the API that affect linking, including new functions,
 // new types, or modifications to existing functions or types.
-#define BN_CURRENT_CORE_ABI_VERSION 105
+#define BN_CURRENT_CORE_ABI_VERSION 106
 
 // Minimum ABI version that is supported for loading of plugins. Plugins that
 // are linked to an ABI version less than this will not be able to load and
@@ -5011,6 +5011,13 @@ extern "C"
 	BINARYNINJACOREAPI void BNFunctionCollapseRegion(BNFunction* func,  uint64_t hash);
 	BINARYNINJACOREAPI void BNFunctionExpandRegion(BNFunction* func, uint64_t hash);
 
+	BINARYNINJACOREAPI void BNFunctionStoreMetadata(
+	    BNFunction* func, const char* key, BNMetadata* value, bool isAuto);
+	BINARYNINJACOREAPI BNMetadata* BNFunctionQueryMetadata(BNFunction* func, const char* key);
+	BINARYNINJACOREAPI void BNFunctionRemoveMetadata(BNFunction* func, const char* key);
+	BINARYNINJACOREAPI BNMetadata* BNFunctionGetMetadata(BNFunction* func);
+	BINARYNINJACOREAPI BNMetadata* BNFunctionGetAutoMetadata(BNFunction* func);
+
 	BINARYNINJACOREAPI void BNSetAutoCallTypeAdjustment(
 	    BNFunction* func, BNArchitecture* arch, uint64_t addr, BNTypeWithConfidence* type);
 	BINARYNINJACOREAPI void BNSetUserCallTypeAdjustment(

function.cpp

@@ -3369,6 +3369,42 @@ void Function::ExpandAll()
 	BNFunctionExpandAll(m_object);
 }
 
+
+void Function::StoreMetadata(const std::string& key, Ref<Metadata> value, bool isAuto)
+{
+	if (!value)
+		return;
+	BNFunctionStoreMetadata(m_object, key.c_str(), value->GetObject(), isAuto);
+}
+
+
+Ref<Metadata> Function::QueryMetadata(const std::string& key)
+{
+	BNMetadata* value = BNFunctionQueryMetadata(m_object, key.c_str());
+	if (!value)
+		return nullptr;
+	return new Metadata(value);
+}
+
+
+Ref<Metadata> Function::GetMetadata()
+{
+	return new Metadata(BNFunctionGetMetadata(m_object));
+}
+
+
+Ref<Metadata> Function::GetAutoMetadata()
+{
+	return new Metadata(BNFunctionGetAutoMetadata(m_object));
+}
+
+
+void Function::RemoveMetadata(const std::string& key)
+{
+	BNFunctionRemoveMetadata(m_object, key.c_str());
+}
+
+
 AdvancedFunctionAnalysisDataRequestor::AdvancedFunctionAnalysisDataRequestor(Function* func) : m_func(func)
 {
 	if (m_func)

python/function.py

@@ -47,6 +47,7 @@
 from . import workflow
 from . import languagerepresentation
 from . import deprecation
+from . import metadata
 from . import __version__
 
 # we define the following as such so the linter doesn't confuse 'highlight' the module with the
@@ -3390,6 +3391,68 @@ def is_region_collapsed(self, hash) -> bool:
 		"""
 		return core.BNFunctionIsRegionCollapsed(self.handle, hash)
 
+	def store_metadata(self, key: str, md: metadata.MetadataValueType, isAuto: bool = False) -> None:
+		"""
+		`store_metadata` stores an object for the given key in the current Function. Objects stored using
+		`store_metadata` can be retrieved when the database is reopened unless isAuto is set to True.
+
+		:param str key: key value to associate the Metadata object with
+		:param Varies md: object to store
+		:param bool isAuto: whether the metadata is an auto metadata
+		:rtype: None
+        """
+		_md = md
+		if not isinstance(_md, metadata.Metadata):
+			_md = metadata.Metadata(_md)
+		core.BNFunctionStoreMetadata(self.handle, key, _md.handle, isAuto)
+
+	def query_metadata(self, key: str) -> 'metadata.MetadataValueType':
+		"""
+		`query_metadata` retrieves metadata associated with the given key stored in the current function.
+
+		:param str key: key to query
+		:rtype: metadata associated with the key
+		"""
+		md_handle = core.BNFunctionQueryMetadata(self.handle, key)
+		if md_handle is None:
+			raise KeyError(key)
+		return metadata.Metadata(handle=md_handle).value
+
+	def remove_metadata(self, key: str) -> None:
+		"""
+		`remove_metadata` removes the metadata associated with key from the current function.
+
+		:param str key: key associated with metadata to remove from the function
+		:rtype: None
+		"""
+		core.BNFunctionRemoveMetadata(self.handle, key)
+
+	@property
+	def metadata(self) -> Dict[str, 'metadata.MetadataValueType']:
+		"""
+		`metadata` retrieves the metadata associated with the current function.
+
+		:rtype: metadata associated with the function
+		"""
+		md_handle = core.BNFunctionGetMetadata(self.handle)
+		assert md_handle is not None, "core.BNFunctionGetMetadata returned None"
+		value = metadata.Metadata(handle=md_handle).value
+		assert isinstance(value, dict), "core.BNFunctionGetMetadata did not return a dict"
+		return value
+
+	@property
+	def auto_metadata(self) -> Dict[str, 'metadata.MetadataValueType']:
+		"""
+		`metadata` retrieves the metadata associated with the current function.
+
+		:rtype: metadata associated with the function
+		"""
+		md_handle = core.BNFunctionGetAutoMetadata(self.handle)
+		assert md_handle is not None, "core.BNFunctionGetAutoMetadata returned None"
+		value = metadata.Metadata(handle=md_handle).value
+		assert isinstance(value, dict), "core.BNFunctionGetAutoMetadata did not return a dict"
+		return value
+
 
 class AdvancedFunctionAnalysisDataRequestor:
 	def __init__(self, func: Optional['Function'] = None):

rust/src/function.rs

@@ -48,6 +48,7 @@ use crate::variable::{
     StackVariableReference, Variable,
 };
 use crate::workflow::Workflow;
+use crate::metadata::Metadata;
 use std::ffi::CStr;
 use std::fmt::{Debug, Formatter};
 use std::ptr::NonNull;
@@ -2428,6 +2429,58 @@ 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 {

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));
+}