[Rust] WIP: Misc cleanup surrounding custom basic block analysis

Author: emesare

arch/riscv/src/lib.rs

@@ -9,17 +9,17 @@ use binaryninja::relocation::{Relocation, RelocationHandlerExt};
 use binaryninja::{
     add_optional_plugin_dependency, architecture,
     architecture::{
-        llvm_assemble, Architecture, ArchitectureExt, CoreArchitecture, CustomArchitectureHandle,
-        ImplicitRegisterExtend, InstructionInfo, LlvmServicesCodeModel, LlvmServicesDialect,
-        LlvmServicesRelocMode, Register as Reg, RegisterInfo, UnusedFlag, UnusedRegisterStack,
-        UnusedRegisterStackInfo,
+        Architecture, ArchitectureExt, CoreArchitecture, CustomArchitectureHandle,
+        ImplicitRegisterExtend, InstructionInfo, Register as Reg, RegisterInfo, UnusedFlag,
+        UnusedRegisterStack, UnusedRegisterStackInfo,
     },
     binary_view::{BinaryView, BinaryViewExt},
     calling_convention::{register_calling_convention, CallingConvention, ConventionBuilder},
     custom_binary_view::{BinaryViewType, BinaryViewTypeExt},
     disassembly::{InstructionTextToken, InstructionTextTokenKind},
     function::Function,
     function_recognizer::FunctionRecognizer,
+    llvm::{llvm_assemble, LlvmServicesCodeModel, LlvmServicesDialect, LlvmServicesRelocMode},
     rc::Ref,
     relocation::{
         CoreRelocationHandler, CustomRelocationHandlerHandle, RelocationHandler, RelocationInfo,

rust/src/architecture.rs

@@ -16,10 +16,9 @@
 
 use std::ops::Deref;
 use std::{
-    borrow::{Borrow, Cow},
-    collections::HashMap,
-    ffi::{c_char, c_int, c_void, CStr, CString},
-    fmt::{Debug, Display, Formatter},
+    borrow::Borrow,
+    ffi::{c_char, c_void, CString},
+    fmt::{Debug, Formatter},
     hash::Hash,
     mem::MaybeUninit,
 };
@@ -38,29 +37,29 @@ use crate::{
     calling_convention::CoreCallingConvention,
     data_buffer::DataBuffer,
     disassembly::InstructionTextToken,
-    function::{ArchAndAddr, Function, NativeBlock},
+    function::Function,
     platform::Platform,
     rc::*,
     relocation::CoreRelocationHandler,
     string::{IntoCStr, *},
     types::{NameAndType, Type},
-    BranchType, Endianness,
+    Endianness,
 };
 
 pub mod basic_block;
+pub mod branches;
 pub mod flag;
 pub mod instruction;
 pub mod intrinsic;
-pub mod llvm;
 pub mod register;
 
 // Re-export all the submodules to keep from breaking everyone's code.
 // We split these out just to clarify each part, not necessarily to enforce an extra namespace.
 pub use basic_block::*;
+pub use branches::*;
 pub use flag::*;
 pub use instruction::*;
 pub use intrinsic::*;
-pub use llvm::*;
 pub use register::*;
 
 #[macro_export]
@@ -861,222 +860,6 @@ impl Architecture for CoreArchitecture {
     }
 }
 
-pub struct BasicBlockAnalysisContext {
-    pub(crate) handle: *mut BNBasicBlockAnalysisContext,
-    contextual_returns_dirty: bool,
-
-    // In
-    pub indirect_branches: Vec<IndirectBranchInfo>,
-    pub indirect_no_return_calls: HashSet<ArchAndAddr>,
-    pub analysis_skip_override: BNFunctionAnalysisSkipOverride,
-    pub translate_tail_calls: bool,
-    pub disallow_branch_to_string: bool,
-    pub max_function_size: u64,
-    pub halt_on_invalid_instruction: bool,
-    pub max_size_reached: bool,
-
-    // In/Out
-    contextual_returns: HashMap<ArchAndAddr, bool>,
-
-    // Out
-    direct_code_references: HashMap<u64, ArchAndAddr>,
-    direct_no_return_calls: HashSet<ArchAndAddr>,
-    halted_disassembly_addresses: HashSet<ArchAndAddr>,
-}
-
-impl BasicBlockAnalysisContext {
-    pub unsafe fn from_raw(handle: *mut BNBasicBlockAnalysisContext) -> Self {
-        debug_assert!(!handle.is_null());
-
-        let ctx_ref = &*handle;
-
-        let indirect_branches = (0..ctx_ref.indirectBranchesCount)
-            .map(|i| {
-                let raw: BNIndirectBranchInfo =
-                    unsafe { std::ptr::read(ctx_ref.indirectBranches.add(i)) };
-                IndirectBranchInfo::from(raw)
-            })
-            .collect::<Vec<_>>();
-
-        let indirect_no_return_calls = (0..ctx_ref.indirectNoReturnCallsCount)
-            .map(|i| {
-                let raw = unsafe { std::ptr::read(ctx_ref.indirectNoReturnCalls.add(i)) };
-                ArchAndAddr::from(raw)
-            })
-            .collect::<HashSet<_>>();
-
-        let contextual_returns = (0..ctx_ref.contextualFunctionReturnCount)
-            .map(|i| {
-                let loc = unsafe {
-                    let raw = std::ptr::read(ctx_ref.contextualFunctionReturnLocations.add(i));
-                    ArchAndAddr::from(raw)
-                };
-                let val = unsafe { *ctx_ref.contextualFunctionReturnValues.add(i) };
-                (loc, val)
-            })
-            .collect::<HashMap<_, _>>();
-
-        let direct_code_references = (0..ctx_ref.directRefCount)
-            .map(|i| {
-                let src = unsafe {
-                    let raw = std::ptr::read(ctx_ref.directRefSources.add(i));
-                    ArchAndAddr::from(raw)
-                };
-                let tgt = unsafe { *ctx_ref.directRefTargets.add(i) };
-                (tgt, src)
-            })
-            .collect::<HashMap<_, _>>();
-
-        let direct_no_return_calls = (0..ctx_ref.directNoReturnCallsCount)
-            .map(|i| {
-                let raw = unsafe { std::ptr::read(ctx_ref.directNoReturnCalls.add(i)) };
-                ArchAndAddr::from(raw)
-            })
-            .collect::<HashSet<_>>();
-
-        let halted_disassembly_addresses = (0..ctx_ref.haltedDisassemblyAddressesCount)
-            .map(|i| {
-                let raw = unsafe { std::ptr::read(ctx_ref.haltedDisassemblyAddresses.add(i)) };
-                ArchAndAddr::from(raw)
-            })
-            .collect::<HashSet<_>>();
-
-        BasicBlockAnalysisContext {
-            handle,
-            contextual_returns_dirty: false,
-            indirect_branches,
-            indirect_no_return_calls,
-            analysis_skip_override: ctx_ref.analysisSkipOverride,
-            translate_tail_calls: ctx_ref.translateTailCalls,
-            disallow_branch_to_string: ctx_ref.disallowBranchToString,
-            max_function_size: ctx_ref.maxFunctionSize,
-            halt_on_invalid_instruction: ctx_ref.haltOnInvalidInstructions,
-            max_size_reached: ctx_ref.maxSizeReached,
-            contextual_returns,
-            direct_code_references,
-            direct_no_return_calls,
-            halted_disassembly_addresses,
-        }
-    }
-
-    pub fn add_contextual_return(&mut self, loc: ArchAndAddr, value: bool) {
-        if !self.contextual_returns.contains_key(&loc) {
-            self.contextual_returns_dirty = true;
-        }
-
-        self.contextual_returns.insert(loc, value);
-    }
-
-    pub fn add_direct_code_reference(&mut self, target: u64, src: ArchAndAddr) {
-        self.direct_code_references.entry(target).or_insert(src);
-    }
-
-    pub fn add_direct_no_return_call(&mut self, loc: ArchAndAddr) {
-        self.direct_no_return_calls.insert(loc);
-    }
-
-    pub fn add_halted_disassembly_address(&mut self, loc: ArchAndAddr) {
-        self.halted_disassembly_addresses.insert(loc);
-    }
-
-    pub fn create_basic_block(
-        &self,
-        arch: CoreArchitecture,
-        start: u64,
-    ) -> Option<Ref<BasicBlock<NativeBlock>>> {
-        let raw_block =
-            unsafe { BNAnalyzeBasicBlocksContextCreateBasicBlock(self.handle, arch.handle, start) };
-
-        if raw_block.is_null() {
-            return None;
-        }
-
-        unsafe { Some(BasicBlock::ref_from_raw(raw_block, NativeBlock::new())) }
-    }
-
-    pub fn add_basic_block(&self, block: Ref<BasicBlock<NativeBlock>>) {
-        unsafe {
-            BNAnalyzeBasicBlocksContextAddBasicBlockToFunction(self.handle, block.handle);
-        }
-    }
-
-    pub fn add_temp_outgoing_reference(&self, target: &Function) {
-        unsafe {
-            BNAnalyzeBasicBlocksContextAddTempReference(self.handle, target.handle);
-        }
-    }
-
-    pub fn finalize(&mut self) {
-        if !self.direct_code_references.is_empty() {
-            let total = self.direct_code_references.len();
-            let mut sources: Vec<BNArchitectureAndAddress> = Vec::with_capacity(total);
-            let mut targets: Vec<u64> = Vec::with_capacity(total);
-            for (target, src) in &self.direct_code_references {
-                sources.push(src.into_raw());
-                targets.push(*target);
-            }
-            unsafe {
-                BNAnalyzeBasicBlocksContextSetDirectCodeReferences(
-                    self.handle,
-                    sources.as_mut_ptr(),
-                    targets.as_mut_ptr(),
-                    total,
-                );
-            }
-        }
-
-        if !self.direct_no_return_calls.is_empty() {
-            let total = self.direct_no_return_calls.len();
-            let mut locations: Vec<BNArchitectureAndAddress> = Vec::with_capacity(total);
-            for loc in &self.direct_no_return_calls {
-                locations.push(loc.into_raw());
-            }
-            unsafe {
-                BNAnalyzeBasicBlocksContextSetDirectNoReturnCalls(
-                    self.handle,
-                    locations.as_mut_ptr(),
-                    total,
-                );
-            }
-        }
-
-        if !self.halted_disassembly_addresses.is_empty() {
-            let total = self.halted_disassembly_addresses.len();
-            let mut locations: Vec<BNArchitectureAndAddress> = Vec::with_capacity(total);
-            for loc in &self.halted_disassembly_addresses {
-                locations.push(loc.into_raw());
-            }
-            unsafe {
-                BNAnalyzeBasicBlocksContextSetHaltedDisassemblyAddresses(
-                    self.handle,
-                    locations.as_mut_ptr(),
-                    total,
-                );
-            }
-        }
-
-        if self.contextual_returns_dirty {
-            let total = self.contextual_returns.len();
-            let mut locations: Vec<BNArchitectureAndAddress> = Vec::with_capacity(total);
-            let mut values: Vec<bool> = Vec::with_capacity(total);
-            for (loc, value) in &self.contextual_returns {
-                locations.push(loc.into_raw());
-                values.push(*value);
-            }
-            unsafe {
-                BNAnalyzeBasicBlocksContextSetContextualFunctionReturns(
-                    self.handle,
-                    locations.as_mut_ptr(),
-                    values.as_mut_ptr(),
-                    total,
-                );
-            }
-        }
-
-        unsafe { BNAnalyzeBasicBlocksContextFinalize(self.handle) };
-    }
-}
-
 impl Debug for CoreArchitecture {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("CoreArchitecture")