Rust bindings for custom basic block analysis

6 files changed, 480 insertions, 14 deletions

diff --git a/rust/src/architecture.rs b/rust/src/architecture.rs
index 6b33b6bb..87182a56 100644
--- a/rust/src/architecture.rs
+++ b/rust/src/architecture.rs
@@ -23,14 +23,13 @@ use crate::{
     calling_convention::CoreCallingConvention,
     data_buffer::DataBuffer,
     disassembly::InstructionTextToken,
-    function::Function,
+    function::{ArchAndAddr, Function, NativeBlock},
     platform::Platform,
     rc::*,
     relocation::CoreRelocationHandler,
-    string::IntoCStr,
-    string::*,
+    string::{IntoCStr, *},
     types::{NameAndType, Type},
-    Endianness,
+    BranchType, Endianness,
 };
 use std::ops::Deref;
 use std::{
@@ -42,8 +41,10 @@ use std::{
     mem::MaybeUninit,
 };
 
+use crate::basic_block::BasicBlock;
 use crate::function_recognizer::FunctionRecognizer;
 use crate::relocation::{CustomRelocationHandlerHandle, RelocationHandler};
+use crate::variable::IndirectBranchInfo;
 
 use crate::confidence::Conf;
 use crate::low_level_il::expression::ValueExpr;
@@ -54,6 +55,7 @@ use crate::low_level_il::{LowLevelILMutableExpression, LowLevelILMutableFunction
 pub use binaryninjacore_sys::BNFlagRole as FlagRole;
 pub use binaryninjacore_sys::BNImplicitRegisterExtend as ImplicitRegisterExtend;
 pub use binaryninjacore_sys::BNLowLevelILFlagCondition as FlagCondition;
+use std::collections::HashSet;
 
 macro_rules! newtype {
     ($name:ident, $inner_type:ty) => {
@@ -171,6 +173,23 @@ impl From<BranchKind> for BranchInfo {
     }
 }
 
+impl From<BranchKind> for BranchType {
+    fn from(value: BranchKind) -> Self {
+        match value {
+            BranchKind::Unresolved => BranchType::UnresolvedBranch,
+            BranchKind::Unconditional(_) => BranchType::UnconditionalBranch,
+            BranchKind::True(_) => BranchType::TrueBranch,
+            BranchKind::False(_) => BranchType::FalseBranch,
+            BranchKind::Call(_) => BranchType::CallDestination,
+            BranchKind::FunctionReturn => BranchType::FunctionReturn,
+            BranchKind::SystemCall => BranchType::SystemCall,
+            BranchKind::Indirect => BranchType::IndirectBranch,
+            BranchKind::Exception => BranchType::ExceptionBranch,
+            BranchKind::UserDefined => BranchType::UserDefinedBranch,
+        }
+    }
+}
+
 /// This is the number of branches that can be specified in an [`InstructionInfo`].
 pub const NUM_BRANCH_INFO: usize = 3;
 
@@ -471,13 +490,13 @@ pub trait Architecture: 'static + Sized + AsRef<CoreArchitecture> {
         il: &LowLevelILMutableFunction,
     ) -> Option<(usize, bool)>;
 
-    unsafe fn analyze_basic_blocks(
+    fn analyze_basic_blocks(
         &self,
         function: &mut Function,
-        context: *mut BNBasicBlockAnalysisContext,
+        context: &mut BasicBlockAnalysisContext,
     ) {
         unsafe {
-            BNArchitectureDefaultAnalyzeBasicBlocks(function.handle, context);
+            BNArchitectureDefaultAnalyzeBasicBlocks(function.handle, context.handle);
         }
     }
 
@@ -1544,13 +1563,13 @@ impl Architecture for CoreArchitecture {
         }
     }
 
-    unsafe fn analyze_basic_blocks(
+    fn analyze_basic_blocks(
         &self,
         function: &mut Function,
-        context: *mut BNBasicBlockAnalysisContext,
+        context: &mut BasicBlockAnalysisContext,
     ) {
         unsafe {
-            BNArchitectureAnalyzeBasicBlocks(self.handle, function.handle, context);
+            BNArchitectureAnalyzeBasicBlocks(self.handle, function.handle, context.handle);
         }
     }
 
@@ -1922,6 +1941,250 @@ 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 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>,
+    inlined_unresolved_indirect_branches: 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<_>>();
+
+        let inlined_unresolved_indirect_branches = (0..ctx_ref
+            .inlinedUnresolvedIndirectBranchCount)
+            .map(|i| {
+                let raw =
+                    unsafe { std::ptr::read(ctx_ref.inlinedUnresolvedIndirectBranches.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,
+            max_size_reached: ctx_ref.maxSizeReached,
+            contextual_returns,
+            direct_code_references,
+            direct_no_return_calls,
+            halted_disassembly_addresses,
+            inlined_unresolved_indirect_branches,
+        }
+    }
+
+    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 add_inlined_unresolved_indirect_branch(&mut self, loc: ArchAndAddr) {
+        self.inlined_unresolved_indirect_branches.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.inlined_unresolved_indirect_branches.is_empty() {
+            let total = self.inlined_unresolved_indirect_branches.len();
+            let mut locations: Vec<BNArchitectureAndAddress> = Vec::with_capacity(total);
+            for loc in &self.inlined_unresolved_indirect_branches {
+                locations.push(loc.into_raw());
+            }
+            unsafe {
+                BNAnalyzeBasicBlocksContextSetInlinedUnresolvedIndirectBranches(
+                    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")
@@ -2264,9 +2527,9 @@ where
     {
         let custom_arch = unsafe { &*(ctxt as *mut A) };
         let mut function = unsafe { Function::from_raw(function) };
-        unsafe {
-            custom_arch.analyze_basic_blocks(&mut function, context);
-        }
+        let mut context: BasicBlockAnalysisContext =
+            unsafe { BasicBlockAnalysisContext::from_raw(context) };
+        custom_arch.analyze_basic_blocks(&mut function, &mut context);
     }
 
     extern "C" fn cb_reg_name<A>(ctxt: *mut c_void, reg: u32) -> *mut c_char
diff --git a/rust/src/basic_block.rs b/rust/src/basic_block.rs
index 3e21094a..da41e712 100644
--- a/rust/src/basic_block.rs
+++ b/rust/src/basic_block.rs
@@ -26,6 +26,13 @@ enum EdgeDirection {
     Outgoing,
 }
 
+pub struct PendingBasicBlockEdge {
+    pub branch_type: BranchType,
+    pub target: u64,
+    pub arch: CoreArchitecture,
+    pub fallthrough: bool,
+}
+
 pub struct Edge<'a, C: 'a + BlockContext> {
     pub branch: BranchType,
     pub back_edge: bool,
@@ -160,6 +167,44 @@ impl<C: BlockContext> BasicBlock<C> {
         C::InstructionIndex::from(unsafe { BNGetBasicBlockEnd(self.handle) })
     }
 
+    pub fn start(&self) -> u64 {
+        unsafe { BNGetBasicBlockStart(self.handle) }
+    }
+
+    pub fn end(&self) -> u64 {
+        unsafe { BNGetBasicBlockEnd(self.handle) }
+    }
+
+    pub fn set_end(&self, end: u64) {
+        unsafe {
+            BNSetBasicBlockEnd(self.handle, end);
+        }
+    }
+
+    pub fn add_instruction_data(&self, data: &[u8]) {
+        unsafe {
+            BNBasicBlockAddInstructionData(self.handle, data.as_ptr() as *const _, data.len());
+        }
+    }
+
+    pub fn instruction_data(&self, addr: u64) -> &[u8] {
+        unsafe {
+            let mut size: usize = 0;
+            let data = BNBasicBlockGetInstructionData(self.handle, addr, &mut size);
+            if data.is_null() {
+                return &[];
+            }
+
+            std::slice::from_raw_parts(data, size)
+        }
+    }
+
+    pub fn set_has_invalid_instructions(&self, value: bool) {
+        unsafe {
+            BNBasicBlockSetHasInvalidInstructions(self.handle, value);
+        }
+    }
+
     pub fn raw_length(&self) -> u64 {
         unsafe { BNGetBasicBlockLength(self.handle) }
     }
@@ -194,15 +239,76 @@ impl<C: BlockContext> BasicBlock<C> {
         }
     }
 
+    pub fn pending_outgoing_edges(&self) -> Vec<PendingBasicBlockEdge> {
+        unsafe {
+            let mut count = 0;
+            let edges_ptr = BNGetBasicBlockPendingOutgoingEdges(self.handle, &mut count);
+            let edges = std::slice::from_raw_parts(edges_ptr, count);
+
+            let mut result = Vec::with_capacity(count);
+            for edge in edges {
+                result.push(PendingBasicBlockEdge {
+                    branch_type: edge.type_,
+                    target: edge.target,
+                    arch: CoreArchitecture::from_raw(edge.arch),
+                    fallthrough: edge.fallThrough,
+                });
+            }
+
+            BNFreePendingBasicBlockEdgeList(edges_ptr);
+            result
+        }
+    }
+
+    pub fn add_pending_outgoing_edge(
+        &self,
+        typ: BranchType,
+        addr: u64,
+        arch: CoreArchitecture,
+        fallthrough: bool,
+    ) {
+        unsafe {
+            BNBasicBlockAddPendingOutgoingEdge(self.handle, typ, addr, arch.handle, fallthrough);
+        }
+    }
+
+    pub fn clear_pending_outgoing_edges(&self) {
+        unsafe {
+            BNClearBasicBlockPendingOutgoingEdges(self.handle);
+        }
+    }
+
+    pub fn set_fallthrough_to_function(&self, value: bool) {
+        unsafe {
+            BNBasicBlockSetFallThroughToFunction(self.handle, value);
+        }
+    }
+
+    pub fn is_fallthrough_to_function(&self) -> bool {
+        unsafe { BNBasicBlockIsFallThroughToFunction(self.handle) }
+    }
+
     // is this valid for il blocks? (it looks like up to MLIL it is)
     pub fn has_undetermined_outgoing_edges(&self) -> bool {
         unsafe { BNBasicBlockHasUndeterminedOutgoingEdges(self.handle) }
     }
 
+    pub fn set_undetermined_outgoing_edges(&self, value: bool) {
+        unsafe {
+            BNBasicBlockSetUndeterminedOutgoingEdges(self.handle, value);
+        }
+    }
+
     pub fn can_exit(&self) -> bool {
         unsafe { BNBasicBlockCanExit(self.handle) }
     }
 
+    pub fn set_can_exit(&self, value: bool) {
+        unsafe {
+            BNBasicBlockSetCanExit(self.handle, value);
+        }
+    }
+
     // TODO: Should we new type this? I just cant tell where the consumers of this are.
     pub fn index(&self) -> usize {
         unsafe { BNGetBasicBlockIndex(self.handle) }
diff --git a/rust/src/binary_view.rs b/rust/src/binary_view.rs
index 8eb822f7..5d8f22cf 100644
--- a/rust/src/binary_view.rs
+++ b/rust/src/binary_view.rs
@@ -34,7 +34,7 @@ use crate::external_library::{ExternalLibrary, ExternalLocation};
 use crate::file_accessor::{Accessor, FileAccessor};
 use crate::file_metadata::FileMetadata;
 use crate::flowgraph::FlowGraph;
-use crate::function::{Function, FunctionViewType, NativeBlock};
+use crate::function::{ArchAndAddr, Function, FunctionViewType, NativeBlock};
 use crate::linear_view::{LinearDisassemblyLine, LinearViewCursor};
 use crate::metadata::Metadata;
 use crate::platform::Platform;
@@ -528,6 +528,10 @@ pub trait BinaryViewExt: BinaryViewBase {
         unsafe { BNAbortAnalysis(self.as_ref().handle) }
     }
 
+    fn analysis_is_aborted(&self) -> bool {
+        unsafe { BNAnalysisIsAborted(self.as_ref().handle) }
+    }
+
     fn workflow(&self) -> Ref<Workflow> {
         unsafe {
             let raw_ptr = BNGetWorkflowForBinaryView(self.as_ref().handle);
@@ -1380,6 +1384,33 @@ pub trait BinaryViewExt: BinaryViewBase {
         }
     }
 
+    fn should_skip_target_analysis(
+        &self,
+        source: &ArchAndAddr,
+        srcfunc: &Function,
+        srcend: u64,
+        target: &ArchAndAddr,
+    ) -> bool {
+        let mut srccopy = BNArchitectureAndAddress {
+            arch: source.arch.handle,
+            address: source.addr,
+        };
+        let mut targetcopy = BNArchitectureAndAddress {
+            arch: target.arch.handle,
+            address: target.addr,
+        };
+
+        unsafe {
+            BNShouldSkipTargetAnalysis(
+                self.as_ref().handle,
+                &mut srccopy,
+                srcfunc.handle,
+                srcend,
+                &mut targetcopy,
+            )
+        }
+    }
+
     fn read_buffer(&self, offset: u64, len: usize) -> Result<DataBuffer> {
         let read_buffer = unsafe { BNReadViewBuffer(self.as_ref().handle, offset, len) };
         if read_buffer.is_null() {
@@ -2149,6 +2180,17 @@ pub trait BinaryViewExt: BinaryViewBase {
         }
     }
 
+    fn string_at(&self, addr: u64) -> Option<BNStringReference> {
+        let mut str_ref = BNStringReference::default();
+        let success = unsafe { BNGetStringAtAddress(self.as_ref().handle, addr, &mut str_ref) };
+
+        if success {
+            Some(str_ref)
+        } else {
+            None
+        }
+    }
+
     /// Retrieve all known strings within the provided `range`.
     ///
     /// NOTE: This returns a list of [`StringReference`] as strings may not be representable
diff --git a/rust/src/function.rs b/rust/src/function.rs
index c4061328..692eb817 100644
--- a/rust/src/function.rs
+++ b/rust/src/function.rs
@@ -2992,3 +2992,34 @@ unsafe impl CoreArrayProviderInner for Comment {
         }
     }
 }
+
+#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
+pub struct ArchAndAddr {
+    pub arch: CoreArchitecture,
+    pub addr: u64,
+}
+
+impl ArchAndAddr {
+    pub fn new(arch: CoreArchitecture, addr: u64) -> Self {
+        Self { arch, addr }
+    }
+}
+
+impl From<BNArchitectureAndAddress> for ArchAndAddr {
+    fn from(raw: BNArchitectureAndAddress) -> Self {
+        unsafe {
+            let arch = CoreArchitecture::from_raw(raw.arch);
+            let addr = raw.address;
+            ArchAndAddr { arch, addr }
+        }
+    }
+}
+
+impl ArchAndAddr {
+    pub fn into_raw(self) -> BNArchitectureAndAddress {
+        BNArchitectureAndAddress {
+            arch: self.arch.handle,
+            address: self.addr,
+        }
+    }
+}
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 3194ae14..9d34f6bc 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -57,6 +57,7 @@ pub mod interaction;
 pub mod language_representation;
 pub mod line_formatter;
 pub mod linear_view;
+pub mod llvm;
 pub mod logger;
 pub mod low_level_il;
 pub mod main_thread;
diff --git a/rust/src/llvm.rs b/rust/src/llvm.rs
new file mode 100644
index 00000000..f74d00e3
--- /dev/null
+++ b/rust/src/llvm.rs
@@ -0,0 +1,23 @@
+use binaryninjacore_sys::BNLlvmServicesDisasmInstruction;
+use std::ffi::CStr;
+
+pub fn disas_instruction(triplet: &str, data: &[u8], address64: u64) -> Option<(usize, String)> {
+    unsafe {
+        let mut buf = vec![0; 256];
+        let instr_len = BNLlvmServicesDisasmInstruction(
+            triplet.as_ptr() as *const i8,
+            data.as_ptr() as *mut u8,
+            data.len() as i32,
+            address64,
+            buf.as_mut_ptr() as *mut i8,
+            buf.len(),
+        );
+        if instr_len > 0 {
+            let cstr = CStr::from_ptr(buf.as_ptr() as *const i8);
+            let string = cstr.to_str().unwrap().to_string();
+            Some((instr_len as usize, string))
+        } else {
+            None
+        }
+    }
+}