Support relocations in Rust architecture plugins

3 files changed, 585 insertions, 0 deletions

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 {}
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index d9ea2d46..3636eb40 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -153,6 +153,7 @@ pub mod mlil;
 pub mod platform;
 pub mod rc;
 pub mod references;
+pub mod relocation;
 pub mod section;
 pub mod segment;
 pub mod settings;
diff --git a/rust/src/relocation.rs b/rust/src/relocation.rs
new file mode 100644
index 00000000..cb7e4e3d
--- /dev/null
+++ b/rust/src/relocation.rs
@@ -0,0 +1,552 @@
+use crate::string::BnStrCompatible;
+use crate::{
+    architecture::{Architecture, CoreArchitecture},
+    binaryview::BinaryView,
+    llil,
+    rc::{Ref, RefCountable},
+    symbol::Symbol,
+};
+use binaryninjacore_sys::*;
+use std::borrow::Borrow;
+use std::os::raw::c_void;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub enum RelocationType {
+    ELFGlobalRelocationType,
+    ELFCopyRelocationType,
+    ELFJumpSlotRelocationType,
+    StandardRelocationType,
+    IgnoredRelocation,
+    UnhandledRelocation,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub enum RelocationOperand {
+    Operand(usize),
+    AutocoerceExternPtr,
+    NocoerceExternPtr,
+    Invalid,
+}
+
+impl From<BNRelocationType> for RelocationType {
+    fn from(t: BNRelocationType) -> Self {
+        match t {
+            BNRelocationType::ELFGlobalRelocationType => RelocationType::ELFGlobalRelocationType,
+            BNRelocationType::ELFCopyRelocationType => RelocationType::ELFCopyRelocationType,
+            BNRelocationType::ELFJumpSlotRelocationType => {
+                RelocationType::ELFJumpSlotRelocationType
+            }
+            BNRelocationType::StandardRelocationType => RelocationType::StandardRelocationType,
+            BNRelocationType::IgnoredRelocation => RelocationType::IgnoredRelocation,
+            BNRelocationType::UnhandledRelocation => RelocationType::UnhandledRelocation,
+        }
+    }
+}
+
+impl From<RelocationType> for BNRelocationType {
+    fn from(t: RelocationType) -> Self {
+        match t {
+            RelocationType::ELFGlobalRelocationType => BNRelocationType::ELFGlobalRelocationType,
+            RelocationType::ELFCopyRelocationType => BNRelocationType::ELFCopyRelocationType,
+            RelocationType::ELFJumpSlotRelocationType => {
+                BNRelocationType::ELFJumpSlotRelocationType
+            }
+            RelocationType::StandardRelocationType => BNRelocationType::StandardRelocationType,
+            RelocationType::IgnoredRelocation => BNRelocationType::IgnoredRelocation,
+            RelocationType::UnhandledRelocation => BNRelocationType::UnhandledRelocation,
+        }
+    }
+}
+
+impl From<usize> for RelocationOperand {
+    fn from(operand: usize) -> Self {
+        match operand {
+            0xfffffffd => RelocationOperand::AutocoerceExternPtr,
+            0xfffffffe => RelocationOperand::NocoerceExternPtr,
+            0xffffffff => RelocationOperand::Invalid,
+            _ => RelocationOperand::Operand(operand),
+        }
+    }
+}
+
+impl From<RelocationOperand> for usize {
+    fn from(operand: RelocationOperand) -> Self {
+        match operand {
+            RelocationOperand::Operand(operand) => operand,
+            RelocationOperand::AutocoerceExternPtr => 0xfffffffd,
+            RelocationOperand::NocoerceExternPtr => 0xfffffffe,
+            RelocationOperand::Invalid => 0xffffffff,
+        }
+    }
+}
+
+// TODO: How to handle related relocation linked lists?
+#[derive(Clone)]
+pub struct RelocationInfo {
+    pub type_: RelocationType,
+    pub pc_relative: bool,
+    pub base_relative: bool,
+    pub base: u64,
+    pub size: usize,
+    pub truncate_size: usize,
+    pub native_type: u64,
+    pub addend: usize,
+    pub has_sign: bool,
+    pub implicit_addend: bool,
+    pub external: bool,
+    pub symbol_index: usize,
+    pub section_index: usize,
+    pub address: u64,
+    pub target: u64,
+    pub data_relocation: bool,
+    relocation_data_cache: [u8; MAX_RELOCATION_SIZE as usize],
+}
+
+impl RelocationInfo {
+    pub fn new() -> Self {
+        RelocationInfo {
+            type_: RelocationType::UnhandledRelocation,
+            pc_relative: false,
+            base_relative: false,
+            base: 0,
+            size: 0,
+            truncate_size: 0,
+            native_type: 0,
+            addend: 0,
+            has_sign: false,
+            implicit_addend: false,
+            external: false,
+            symbol_index: 0,
+            section_index: 0,
+            address: 0,
+            target: 0,
+            data_relocation: false,
+            relocation_data_cache: [0; MAX_RELOCATION_SIZE as usize],
+        }
+    }
+
+    pub(crate) fn from_raw(reloc: &BNRelocationInfo) -> Self {
+        RelocationInfo {
+            type_: reloc.type_.into(),
+            pc_relative: reloc.pcRelative,
+            base_relative: reloc.baseRelative,
+            base: reloc.base,
+            size: reloc.size,
+            truncate_size: reloc.truncateSize,
+            native_type: reloc.nativeType,
+            addend: reloc.addend,
+            has_sign: reloc.hasSign,
+            implicit_addend: reloc.implicitAddend,
+            external: reloc.external,
+            symbol_index: reloc.symbolIndex,
+            section_index: reloc.sectionIndex,
+            address: reloc.address,
+            target: reloc.target,
+            data_relocation: reloc.dataRelocation,
+            relocation_data_cache: reloc.relocationDataCache,
+        }
+    }
+
+    pub(crate) fn as_raw(&self) -> BNRelocationInfo {
+        BNRelocationInfo {
+            type_: self.type_.into(),
+            pcRelative: self.pc_relative,
+            baseRelative: self.base_relative,
+            base: self.base,
+            size: self.size,
+            truncateSize: self.truncate_size,
+            nativeType: self.native_type,
+            addend: self.addend,
+            hasSign: self.has_sign,
+            implicitAddend: self.implicit_addend,
+            external: self.external,
+            symbolIndex: self.symbol_index,
+            sectionIndex: self.section_index,
+            address: self.address,
+            target: self.target,
+            dataRelocation: self.data_relocation,
+            relocationDataCache: self.relocation_data_cache,
+            prev: core::ptr::null_mut(),
+            next: core::ptr::null_mut(),
+        }
+    }
+}
+
+pub struct Relocation(*mut BNRelocation);
+
+impl Relocation {
+    pub(crate) unsafe fn from_raw(reloc: *mut BNRelocation) -> Self {
+        Relocation(reloc)
+    }
+
+    pub fn info(&self) -> RelocationInfo {
+        RelocationInfo::from_raw(unsafe { &BNRelocationGetInfo(self.0) })
+    }
+
+    pub fn architecture(&self) -> Option<CoreArchitecture> {
+        let raw = unsafe { BNRelocationGetArchitecture(self.0) };
+        if raw.is_null() {
+            return None;
+        }
+
+        Some(unsafe { CoreArchitecture::from_raw(raw) })
+    }
+
+    pub fn target(&self) -> u64 {
+        unsafe { BNRelocationGetTarget(self.0) }
+    }
+
+    pub fn address(&self) -> u64 {
+        unsafe { BNRelocationGetReloc(self.0) }
+    }
+
+    pub fn symbol(&self) -> Option<Ref<Symbol>> {
+        let raw = unsafe { BNRelocationGetSymbol(self.0) };
+        if raw.is_null() {
+            return None;
+        }
+
+        Some(unsafe { Symbol::ref_from_raw(raw) })
+    }
+}
+
+pub trait RelocationHandler: 'static + Sized + AsRef<CoreRelocationHandler> {
+    type Handle: Borrow<Self>;
+
+    fn get_relocation_info(
+        &self,
+        _bv: &BinaryView,
+        _arch: &CoreArchitecture,
+        _info: &mut [RelocationInfo],
+    ) -> bool {
+        false
+    }
+
+    fn apply_relocation(
+        &self,
+        bv: &BinaryView,
+        arch: &CoreArchitecture,
+        reloc: &Relocation,
+        dest: &mut [u8],
+    ) -> bool {
+        self.default_apply_relocation(bv, arch, reloc, dest)
+    }
+
+    fn get_operand_for_external_relocation(
+        &self,
+        _data: &[u8],
+        _addr: u64,
+        _il: &llil::RegularFunction<CoreArchitecture>,
+        _reloc: &Relocation,
+    ) -> RelocationOperand {
+        RelocationOperand::AutocoerceExternPtr
+    }
+
+    fn handle(&self) -> Self::Handle;
+}
+
+pub trait RelocationHandlerExt: RelocationHandler {
+    fn default_apply_relocation(
+        &self,
+        bv: &BinaryView,
+        arch: &CoreArchitecture,
+        reloc: &Relocation,
+        dest: &mut [u8],
+    ) -> bool {
+        unsafe {
+            BNRelocationHandlerDefaultApplyRelocation(
+                self.as_ref().0,
+                bv.handle,
+                arch.handle().as_ref().0,
+                reloc.0,
+                dest.as_mut_ptr(),
+                dest.len(),
+            )
+        }
+    }
+}
+
+impl<T: RelocationHandler> RelocationHandlerExt for T {}
+
+#[derive(Eq, PartialEq, Hash, Debug)]
+pub struct CoreRelocationHandler(*mut BNRelocationHandler);
+
+unsafe impl Send for CoreRelocationHandler {}
+unsafe impl Sync for CoreRelocationHandler {}
+
+impl CoreRelocationHandler {
+    pub(crate) unsafe fn ref_from_raw(raw: *mut BNRelocationHandler) -> Ref<Self> {
+        unsafe { Ref::new(CoreRelocationHandler(raw)) }
+    }
+}
+
+impl AsRef<CoreRelocationHandler> for CoreRelocationHandler {
+    fn as_ref(&self) -> &Self {
+        self
+    }
+}
+
+impl RelocationHandler for CoreRelocationHandler {
+    type Handle = Self;
+
+    fn get_relocation_info(
+        &self,
+        bv: &BinaryView,
+        arch: &CoreArchitecture,
+        info: &mut [RelocationInfo],
+    ) -> bool {
+        let mut raw_info = info.iter().map(|i| i.as_raw()).collect::<Vec<_>>();
+        let res = unsafe {
+            BNRelocationHandlerGetRelocationInfo(
+                self.0,
+                bv.handle,
+                arch.handle().as_ref().0,
+                raw_info.as_mut_ptr(),
+                raw_info.len(),
+            )
+        };
+        for (info, raw) in info.iter_mut().zip(raw_info.iter()) {
+            *info = RelocationInfo::from_raw(raw);
+        }
+        res
+    }
+
+    fn apply_relocation(
+        &self,
+        bv: &BinaryView,
+        arch: &CoreArchitecture,
+        reloc: &Relocation,
+        dest: &mut [u8],
+    ) -> bool {
+        unsafe {
+            BNRelocationHandlerApplyRelocation(
+                self.0,
+                bv.handle,
+                arch.handle().as_ref().0,
+                reloc.0,
+                dest.as_mut_ptr(),
+                dest.len(),
+            )
+        }
+    }
+
+    fn get_operand_for_external_relocation(
+        &self,
+        data: &[u8],
+        addr: u64,
+        il: &llil::RegularFunction<CoreArchitecture>,
+        reloc: &Relocation,
+    ) -> RelocationOperand {
+        unsafe {
+            BNRelocationHandlerGetOperandForExternalRelocation(
+                self.0,
+                data.as_ptr(),
+                addr,
+                data.len(),
+                il.handle,
+                reloc.0,
+            )
+            .into()
+        }
+    }
+
+    fn handle(&self) -> CoreRelocationHandler {
+        CoreRelocationHandler(self.0)
+    }
+}
+
+impl ToOwned for CoreRelocationHandler {
+    type Owned = Ref<Self>;
+
+    fn to_owned(&self) -> Self::Owned {
+        unsafe { RefCountable::inc_ref(self) }
+    }
+}
+
+unsafe impl RefCountable for CoreRelocationHandler {
+    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
+        Ref::new(Self(BNNewRelocationHandlerReference(handle.0)))
+    }
+
+    unsafe fn dec_ref(handle: &Self) {
+        BNFreeRelocationHandler(handle.0);
+    }
+}
+
+pub(crate) fn register_relocation_handler<S, R, F>(arch: &CoreArchitecture, name: S, func: F)
+where
+    S: BnStrCompatible,
+    R: 'static + RelocationHandler<Handle = CustomRelocationHandlerHandle<R>> + Send + Sync + Sized,
+    F: FnOnce(CustomRelocationHandlerHandle<R>, CoreRelocationHandler) -> R,
+{
+    #[repr(C)]
+    struct RelocationHandlerBuilder<R>
+    where
+        R: 'static + RelocationHandler<Handle = CustomRelocationHandlerHandle<R>> + Send + Sync,
+    {
+        handler: R,
+    }
+
+    extern "C" fn cb_free<R>(ctxt: *mut c_void)
+    where
+        R: 'static + RelocationHandler<Handle = CustomRelocationHandlerHandle<R>> + Send + Sync,
+    {
+        unsafe {
+            let _handler = Box::from_raw(ctxt as *mut RelocationHandlerBuilder<R>);
+        }
+    }
+
+    extern "C" fn cb_get_relocation_info<R>(
+        ctxt: *mut c_void,
+        bv: *mut BNBinaryView,
+        arch: *mut BNArchitecture,
+        result: *mut BNRelocationInfo,
+        count: usize,
+    ) -> bool
+    where
+        R: 'static + RelocationHandler<Handle = CustomRelocationHandlerHandle<R>> + Send + Sync,
+    {
+        let custom_handler = unsafe { &*(ctxt as *mut R) };
+        let bv = unsafe { BinaryView::from_raw(BNNewViewReference(bv)) };
+        let arch = unsafe { CoreArchitecture::from_raw(arch) };
+        let result = unsafe { core::slice::from_raw_parts_mut(result, count) };
+        let mut info = result
+            .iter()
+            .map(|i| RelocationInfo::from_raw(i))
+            .collect::<Vec<_>>();
+        let ok =
+            custom_handler.get_relocation_info(bv.as_ref(), arch.as_ref(), info.as_mut_slice());
+        for (result, info) in result.iter_mut().zip(info.iter()) {
+            *result = info.as_raw();
+        }
+        ok
+    }
+
+    extern "C" fn cb_apply_relocation<R>(
+        ctxt: *mut c_void,
+        bv: *mut BNBinaryView,
+        arch: *mut BNArchitecture,
+        reloc: *mut BNRelocation,
+        dest: *mut u8,
+        len: usize,
+    ) -> bool
+    where
+        R: 'static + RelocationHandler<Handle = CustomRelocationHandlerHandle<R>> + Send + Sync,
+    {
+        let custom_handler = unsafe { &*(ctxt as *mut R) };
+        let bv = unsafe { BinaryView::from_raw(BNNewViewReference(bv)) };
+        let arch = unsafe { CoreArchitecture::from_raw(arch) };
+        let reloc = unsafe { Relocation::from_raw(reloc) };
+        let dest = unsafe { core::slice::from_raw_parts_mut(dest, len) };
+        custom_handler.apply_relocation(bv.as_ref(), arch.as_ref(), &reloc, dest)
+    }
+
+    extern "C" fn cb_get_operand_for_external_relocation<R>(
+        ctxt: *mut c_void,
+        data: *const u8,
+        addr: u64,
+        len: usize,
+        il: *mut BNLowLevelILFunction,
+        reloc: *mut BNRelocation,
+    ) -> usize
+    where
+        R: 'static + RelocationHandler<Handle = CustomRelocationHandlerHandle<R>> + Send + Sync,
+    {
+        let custom_handler = unsafe { &*(ctxt as *mut R) };
+        let data = unsafe { core::slice::from_raw_parts(data, len) };
+        let reloc = unsafe { Relocation::from_raw(reloc) };
+
+        let func = unsafe { BNGetLowLevelILOwnerFunction(il) };
+        if func.is_null() {
+            return RelocationOperand::Invalid.into();
+        }
+
+        let arch = unsafe { BNGetFunctionArchitecture(func) };
+        unsafe { BNFreeFunction(func) };
+        if arch.is_null() {
+            return RelocationOperand::Invalid.into();
+        }
+        let arch = unsafe { CoreArchitecture::from_raw(arch) };
+
+        let il = unsafe { llil::RegularFunction::from_raw(arch, il) };
+
+        custom_handler
+            .get_operand_for_external_relocation(data, addr, &il, &reloc)
+            .into()
+    }
+
+    let name = name.into_bytes_with_nul();
+
+    let uninit_handler = RelocationHandlerBuilder {
+        handler: unsafe { std::mem::zeroed() },
+    };
+    let raw = Box::into_raw(Box::new(uninit_handler));
+    let mut custom_handler = BNCustomRelocationHandler {
+        context: raw as *mut _,
+        freeObject: Some(cb_free::<R>),
+        getRelocationInfo: Some(cb_get_relocation_info::<R>),
+        applyRelocation: Some(cb_apply_relocation::<R>),
+        getOperandForExternalRelocation: Some(cb_get_operand_for_external_relocation::<R>),
+    };
+
+    let handle_raw = unsafe { BNCreateRelocationHandler(&mut custom_handler) };
+    assert!(!handle_raw.is_null());
+    let handle = CoreRelocationHandler(handle_raw);
+    let custom_handle = CustomRelocationHandlerHandle {
+        handle: raw as *mut R,
+    };
+    unsafe {
+        core::ptr::write(
+            &mut raw.as_mut().unwrap().handler,
+            func(custom_handle, CoreRelocationHandler(handle.0)),
+        );
+
+        BNArchitectureRegisterRelocationHandler(
+            arch.handle().as_ref().0,
+            name.as_ref().as_ptr() as *const _,
+            handle.handle().as_ref().0,
+        );
+    }
+}
+
+pub struct CustomRelocationHandlerHandle<R>
+where
+    R: 'static + RelocationHandler<Handle = CustomRelocationHandlerHandle<R>> + Send + Sync,
+{
+    handle: *mut R,
+}
+
+unsafe impl<R> Send for CustomRelocationHandlerHandle<R> where
+    R: 'static + RelocationHandler<Handle = CustomRelocationHandlerHandle<R>> + Send + Sync
+{
+}
+
+unsafe impl<R> Sync for CustomRelocationHandlerHandle<R> where
+    R: 'static + RelocationHandler<Handle = CustomRelocationHandlerHandle<R>> + Send + Sync
+{
+}
+
+impl<R> Clone for CustomRelocationHandlerHandle<R>
+where
+    R: 'static + RelocationHandler<Handle = Self> + Send + Sync,
+{
+    fn clone(&self) -> Self {
+        Self {
+            handle: self.handle,
+        }
+    }
+}
+
+impl<R> Copy for CustomRelocationHandlerHandle<R> where
+    R: 'static + RelocationHandler<Handle = Self> + Send + Sync
+{
+}
+
+impl<R> Borrow<R> for CustomRelocationHandlerHandle<R>
+where
+    R: 'static + RelocationHandler<Handle = Self> + Send + Sync,
+{
+    fn borrow(&self) -> &R {
+        unsafe { &*self.handle }
+    }
+}