From 2a4c7d5d89907497e029337bbaf6f7e467bcde98 Mon Sep 17 00:00:00 2001 From: Brandon Miller Date: Mon, 16 Feb 2026 14:16:25 -0500 Subject: Share context between BB recovery, lifing, and disassembly Adds a GetInstructionTextWithContext callback to the architecture class that can be used to pass data from AnalyzeBasicBlocks. This same context is also supplied to LiftFunction and allows for supplying shared function and/or binary view level information across basic block analysis, function lifting, and disassembly text rendering --- rust/src/architecture.rs | 211 +++++++++++++++++++++++++++++++++++ rust/src/architecture/basic_block.rs | 32 +++++- 2 files changed, 240 insertions(+), 3 deletions(-) (limited to 'rust/src') diff --git a/rust/src/architecture.rs b/rust/src/architecture.rs index 6b5d44e5..6b644e88 100644 --- a/rust/src/architecture.rs +++ b/rust/src/architecture.rs @@ -42,6 +42,8 @@ use std::{ mem::MaybeUninit, }; +use std::ptr::NonNull; + use crate::function_recognizer::FunctionRecognizer; use crate::relocation::{CustomRelocationHandlerHandle, RelocationHandler}; @@ -192,6 +194,30 @@ pub trait Architecture: 'static + Sized + AsRef { addr: u64, ) -> Option<(usize, Vec)>; + /// Disassembles a raw byte sequence into a human-readable list of text tokens. + /// + /// This function is responsible for the visual representation of assembly instructions. + /// It does *not* define semantics (use [`Architecture::instruction_llil`] for that); + /// it simply tells the UI how to print the instruction. This variant includes contextual data, which + /// can be produced by analyze_basic_blocks + /// + /// # Returns + /// + /// An `Option` containing a tuple: + /// + /// * `usize`: The size of the decoded instruction in bytes. Is used to advance to the next instruction. + /// * `Vec`: A list of text tokens representing the instruction. + /// + /// Returns `None` if the bytes do not form a valid instruction. + fn instruction_text_with_context( + &self, + data: &[u8], + addr: u64, + _context: Option>, + ) -> Option<(usize, Vec)> { + self.instruction_text(data, addr) + } + // TODO: Why do we need to return a boolean here? Does `None` not represent the same thing? /// Appends arbitrary low-level il instructions to `il`. /// @@ -550,6 +576,19 @@ pub trait Architecture: 'static + Sized + AsRef { fn handle(&self) -> Self::Handle; } +pub trait ArchitectureWithFunctionContext: Architecture { + type FunctionArchContext: Send + Sync + 'static; + + fn instruction_text_with_typed_context( + &self, + data: &[u8], + addr: u64, + _context: Option<&Self::FunctionArchContext>, + ) -> Option<(usize, Vec)> { + self.instruction_text(data, addr) + } +} + pub struct FunctionLifterContext { pub(crate) handle: *mut BNFunctionLifterContext, } @@ -560,6 +599,20 @@ impl FunctionLifterContext { FunctionLifterContext { handle } } + + pub fn get_function_arch_context( + &self, + _arch: &A, + ) -> Option<&A::FunctionArchContext> { + unsafe { + let ptr = (*self.handle).functionArchContext; + if ptr.is_null() { + None + } else { + Some(&*(ptr as *const A::FunctionArchContext)) + } + } + } } // TODO: WTF?!?!?!? @@ -705,6 +758,38 @@ impl Architecture for CoreArchitecture { } } + fn instruction_text_with_context( + &self, + data: &[u8], + addr: u64, + context: Option>, + ) -> Option<(usize, Vec)> { + let mut consumed = data.len(); + let mut count: usize = 0; + let mut result: *mut BNInstructionTextToken = std::ptr::null_mut(); + let ctx_ptr: *mut c_void = context.map_or(std::ptr::null_mut(), |p| p.as_ptr()); + unsafe { + if BNGetInstructionTextWithContext( + self.handle, + data.as_ptr(), + addr, + &mut consumed, + ctx_ptr, + &mut result, + &mut count, + ) { + let instr_text_tokens = std::slice::from_raw_parts(result, count) + .iter() + .map(InstructionTextToken::from_raw) + .collect(); + BNFreeInstructionText(result, count); + Some((consumed, instr_text_tokens)) + } else { + None + } + } + } + fn instruction_llil( &self, data: &[u8], @@ -1273,6 +1358,15 @@ pub fn register_architecture(name: &str, func: F) -> &'static A where A: 'static + Architecture> + Send + Sync + Sized, F: FnOnce(CustomArchitectureHandle, CoreArchitecture) -> A, +{ + register_architecture_impl(name, func, |_| {}) +} + +fn register_architecture_impl(name: &str, func: F, customize: C) -> &'static A +where + A: 'static + Architecture> + Send + Sync + Sized, + F: FnOnce(CustomArchitectureHandle, CoreArchitecture) -> A, + C: FnOnce(&mut BNCustomArchitecture), { #[repr(C)] struct ArchitectureBuilder @@ -1419,6 +1513,43 @@ where true } + pub unsafe extern "C" fn cb_get_instruction_text_with_context( + ctxt: *mut c_void, + data: *const u8, + addr: u64, + len: *mut usize, + context: *mut c_void, + result: *mut *mut BNInstructionTextToken, + count: *mut usize, + ) -> bool + where + A: 'static + Architecture> + Send + Sync, + { + let custom_arch = unsafe { &*(ctxt as *mut A) }; + let data = unsafe { std::slice::from_raw_parts(data, *len) }; + let result = unsafe { &mut *result }; + let context = NonNull::new(context); + + let Some((res_size, res_tokens)) = + custom_arch.instruction_text_with_context(data, addr, context) + else { + return false; + }; + + let res_tokens: Box<[BNInstructionTextToken]> = res_tokens + .into_iter() + .map(InstructionTextToken::into_raw) + .collect(); + unsafe { + // NOTE: Freed with `cb_free_instruction_text` + let res_tokens = Box::leak(res_tokens); + *result = res_tokens.as_mut_ptr(); + *count = res_tokens.len(); + *len = res_size; + } + true + } + extern "C" fn cb_free_instruction_text(tokens: *mut BNInstructionTextToken, count: usize) { unsafe { let raw_tokens = std::slice::from_raw_parts_mut(tokens, count); @@ -2401,10 +2532,12 @@ where getAssociatedArchitectureByAddress: Some(cb_associated_arch_by_addr::), getInstructionInfo: Some(cb_instruction_info::), getInstructionText: Some(cb_get_instruction_text::), + getInstructionTextWithContext: Some(cb_get_instruction_text_with_context::), freeInstructionText: Some(cb_free_instruction_text), getInstructionLowLevelIL: Some(cb_instruction_llil::), analyzeBasicBlocks: Some(cb_analyze_basic_blocks::), liftFunction: Some(cb_lift_function::), + freeFunctionArchContext: None, getRegisterName: Some(cb_reg_name::), getFlagName: Some(cb_flag_name::), @@ -2471,6 +2604,8 @@ where skipAndReturnValue: Some(cb_skip_and_return_value::), }; + customize(&mut custom_arch); + unsafe { let res = BNRegisterArchitecture(name.as_ptr(), &mut custom_arch as *mut _); @@ -2480,6 +2615,82 @@ where } } +pub fn register_architecture_with_function_context(name: &str, func: F) -> &'static A +where + A: 'static + + ArchitectureWithFunctionContext> + + Send + + Sync + + Sized, + F: FnOnce(CustomArchitectureHandle, CoreArchitecture) -> A, +{ + unsafe extern "C" fn cb_free_function_arch_context_typed( + _ctxt: *mut c_void, + context: *mut c_void, + ) where + A: 'static + + ArchitectureWithFunctionContext> + + Send + + Sync, + { + if context.is_null() { + return; + } + // The context was allocated via Box::into_raw in set_function_arch_context, + // so we reconstruct the Box here and let it drop. + let _ = unsafe { Box::from_raw(context as *mut A::FunctionArchContext) }; + } + + unsafe extern "C" fn cb_get_instruction_text_with_context_typed( + ctxt: *mut c_void, + data: *const u8, + addr: u64, + len: *mut usize, + context: *mut c_void, + result: *mut *mut BNInstructionTextToken, + count: *mut usize, + ) -> bool + where + A: 'static + + ArchitectureWithFunctionContext> + + Send + + Sync, + { + let custom_arch = unsafe { &*(ctxt as *mut A) }; + let data = unsafe { std::slice::from_raw_parts(data, *len) }; + let result = unsafe { &mut *result }; + let typed_context: Option<&A::FunctionArchContext> = if context.is_null() { + None + } else { + Some(unsafe { &*(context as *const A::FunctionArchContext) }) + }; + + let Some((res_size, res_tokens)) = + custom_arch.instruction_text_with_typed_context(data, addr, typed_context) + else { + return false; + }; + + let res_tokens: Box<[BNInstructionTextToken]> = res_tokens + .into_iter() + .map(InstructionTextToken::into_raw) + .collect(); + unsafe { + let res_tokens = Box::leak(res_tokens); + *result = res_tokens.as_mut_ptr(); + *count = res_tokens.len(); + *len = res_size; + } + true + } + + register_architecture_impl(name, func, |custom_arch| { + custom_arch.freeFunctionArchContext = Some(cb_free_function_arch_context_typed::); + custom_arch.getInstructionTextWithContext = + Some(cb_get_instruction_text_with_context_typed::); + }) +} + #[derive(Debug)] pub struct CustomArchitectureHandle where diff --git a/rust/src/architecture/basic_block.rs b/rust/src/architecture/basic_block.rs index 428a97a1..ac5ce774 100644 --- a/rust/src/architecture/basic_block.rs +++ b/rust/src/architecture/basic_block.rs @@ -1,4 +1,4 @@ -use crate::architecture::{CoreArchitecture, IndirectBranchInfo}; +use crate::architecture::{ArchitectureWithFunctionContext, CoreArchitecture, IndirectBranchInfo}; use crate::basic_block::BasicBlock; use crate::function::{Function, Location, NativeBlock}; use crate::rc::Ref; @@ -180,6 +180,34 @@ impl BasicBlockAnalysisContext { self.inlined_unresolved_indirect_branches.insert(loc.into()); } + pub fn set_function_arch_context( + &mut self, + _arch: &A, + context: Box, + ) -> bool { + unsafe { + if !(*self.handle).functionArchContext.is_null() { + return false; + } + (*self.handle).functionArchContext = Box::into_raw(context) as *mut std::ffi::c_void; + } + true + } + + pub fn get_function_arch_context( + &self, + _arch: &A, + ) -> Option<&A::FunctionArchContext> { + unsafe { + let ptr = (*self.handle).functionArchContext; + if ptr.is_null() { + None + } else { + Some(&*(ptr as *const A::FunctionArchContext)) + } + } + } + /// Creates a new [`BasicBlock`] at the specified address for the given [`CoreArchitecture`]. /// /// After creating, you can add using [`BasicBlockAnalysisContext::add_basic_block`]. @@ -328,8 +356,6 @@ impl BasicBlockAnalysisContext { if self.contextual_returns_dirty { self.update_contextual_returns(); } - - unsafe { BNAnalyzeBasicBlocksContextFinalize(self.handle) }; } } -- cgit v1.3.1