From 7c98724a614550e37e5a33805e13179c87363b91 Mon Sep 17 00:00:00 2001
From: Rusty Wagner <rusty.wagner@gmail.com>
Date: Thu, 21 Dec 2023 18:54:27 -0700
Subject: Support relocations in Rust architecture plugins

---
 rust/src/architecture.rs | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

(limited to 'rust/src/architecture.rs')

diff --git a/rust/src/architecture.rs b/rust/src/architecture.rs
index 02a1e901..2c9843d9 100644
--- a/rust/src/architecture.rs
+++ b/rust/src/architecture.rs
@@ -36,6 +36,7 @@ use crate::{
     },
     platform::Platform,
     rc::*,
+    relocation::CoreRelocationHandler,
     string::BnStrCompatible,
     string::*,
     types::{Conf, NameAndType, Type},
@@ -174,6 +175,7 @@ impl InstructionInfo {
     }
 }
 
+use crate::relocation::{CustomRelocationHandlerHandle, RelocationHandler};
 pub use binaryninjacore_sys::BNFlagRole as FlagRole;
 pub use binaryninjacore_sys::BNImplicitRegisterExtend as ImplicitRegisterExtend;
 pub use binaryninjacore_sys::BNLowLevelILFlagCondition as FlagCondition;
@@ -1614,6 +1616,36 @@ pub trait ArchitectureExt: Architecture {
             Some(Platform::ref_from_raw(handle))
         }
     }
+
+    fn relocation_handler(&self, view_name: &str) -> Option<Ref<CoreRelocationHandler>> {
+        let view_name = match CString::new(view_name) {
+            Ok(view_name) => view_name,
+            Err(_) => return None,
+        };
+
+        unsafe {
+            let handle = BNArchitectureGetRelocationHandler(self.as_ref().0, view_name.as_ptr());
+
+            if handle.is_null() {
+                return None;
+            }
+
+            Some(CoreRelocationHandler::ref_from_raw(handle))
+        }
+    }
+
+    fn register_relocation_handler<S, R, F>(&self, name: S, func: F)
+    where
+        S: BnStrCompatible,
+        R: 'static
+            + RelocationHandler<Handle = CustomRelocationHandlerHandle<R>>
+            + Send
+            + Sync
+            + Sized,
+        F: FnOnce(CustomRelocationHandlerHandle<R>, CoreRelocationHandler) -> R,
+    {
+        crate::relocation::register_relocation_handler(self.as_ref(), name, func);
+    }
 }
 
 impl<T: Architecture> ArchitectureExt for T {}
-- 
cgit v1.3.1