diff --git a/rust/src/architecture.rs b/rust/src/architecture.rs index b38cc52a4..713e8bf1e 100644 --- a/rust/src/architecture.rs +++ b/rust/src/architecture.rs @@ -715,21 +715,6 @@ impl Register for CoreRegister { } } -impl CoreArrayProvider for CoreRegister { - type Raw = u32; - type Context = CoreArchitecture; - type Wrapped<'a> = Self; -} - -unsafe impl CoreArrayProviderInner for CoreRegister { - unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) { - BNFreeRegisterList(raw) - } - unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> { - context.register_from_id(*raw).unwrap() - } -} - pub struct CoreRegisterStackInfo(*mut BNArchitecture, BNRegisterStackInfo); impl RegisterStackInfo for CoreRegisterStackInfo { diff --git a/rust/src/function.rs b/rust/src/function.rs index 295b420ad..9a098c0fa 100644 --- a/rust/src/function.rs +++ b/rust/src/function.rs @@ -15,7 +15,7 @@ use binaryninjacore_sys::*; use crate::{ - architecture::{CoreArchitecture, CoreRegister}, + architecture::{Architecture, CoreArchitecture, CoreRegister, Register}, basicblock::{BasicBlock, BlockContext}, binaryview::{BinaryView, BinaryViewExt}, callingconvention::CallingConvention, @@ -28,8 +28,8 @@ use crate::{ tags::{Tag, TagReference, TagType}, types::{ Conf, ConstantReference, HighlightColor, IndirectBranchInfo, IntegerDisplayType, - NamedTypedVariable, QualifiedName, RegisterValue, RegisterValueType, - StackVariableReference, Type, Variable, + NamedTypedVariable, QualifiedName, RegisterList, RegisterSet, RegisterValue, + RegisterValueType, StackVariableReference, Type, Variable, }, }; use crate::{databuffer::DataBuffer, disassembly::InstructionTextToken, rc::*}; @@ -1034,7 +1034,11 @@ impl Function { Conf::new(register.value.into(), register.confidence) } - pub fn regs_read_by(&self, addr: u64, arch: Option) -> Array { + pub fn regs_read_by( + &self, + addr: u64, + arch: Option, + ) -> Array> { let arch = arch.unwrap_or_else(|| self.arch()); let mut count = 0; let regs = @@ -1047,7 +1051,7 @@ impl Function { &self, addr: u64, arch: Option, - ) -> Array { + ) -> Array> { let arch = arch.unwrap_or_else(|| self.arch()); let mut count = 0; let regs = @@ -1056,6 +1060,48 @@ impl Function { unsafe { Array::new(regs, count, arch) } } + pub fn set_auto_clobbered_regs(&self, reg_ids: I, confidence: u8) + where + I: IntoIterator, + { + let mut regs: Box<[_]> = reg_ids.into_iter().collect(); + unsafe { + BNSetAutoFunctionClobberedRegisters( + self.handle, + &mut BNRegisterSetWithConfidence { + regs: regs.as_mut_ptr(), + count: regs.len(), + confidence, + }, + ) + } + } + + /// Registers that are modified by this function + pub fn clobbered_regs(&self) -> Conf>> { + let result = unsafe { BNGetFunctionClobberedRegisters(self.handle) }; + + let reg_set = unsafe { Array::new(result.regs, result.count, self.arch().handle()) }; + Conf::new(reg_set, result.confidence) + } + + pub fn set_clobbered_regs(&self, _value: I, confidence: u8) + where + I: IntoIterator, + { + let mut regs: Box<[u32]> = _value.into_iter().map(|reg| reg.id()).collect(); + unsafe { + BNSetUserFunctionClobberedRegisters( + self.handle, + &mut BNRegisterSetWithConfidence { + regs: regs.as_mut_ptr(), + count: regs.len(), + confidence, + }, + ) + } + } + pub fn stack_contents_at( &self, addr: u64, diff --git a/rust/src/types.rs b/rust/src/types.rs index 51ea11ff1..6f8496c24 100644 --- a/rust/src/types.rs +++ b/rust/src/types.rs @@ -31,7 +31,7 @@ use crate::{ use lazy_static::lazy_static; use std::{ - borrow::Cow, + borrow::{Borrow, Cow}, collections::{HashMap, HashSet}, ffi::CStr, fmt::{self, Debug, Display, Formatter}, @@ -3582,3 +3582,74 @@ unsafe impl CoreArrayProviderInner for StackVariableReference { Guard::new(Self::from_raw(*raw), context) } } + +///////////////////////// +// RegisterList + +// NOTE only exists as an array element +pub struct RegisterList { + arch: A::Handle, + register: u32, +} + +impl RegisterList { + pub fn register(&self) -> A::Register { + self.arch.borrow().register_from_id(self.register).unwrap() + } +} + +impl CoreArrayProvider for RegisterList { + type Raw = u32; + type Context = A::Handle; + type Wrapped<'a> = Self; +} + +unsafe impl CoreArrayProviderInner for RegisterList { + unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) { + BNFreeRegisterList(raw) + } + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> { + Self { + arch: context.clone(), + register: *raw, + } + } +} + +///////////////////////// +// RegisterSet + +// NOTE only exists as an array element +pub struct RegisterSet { + arch: A::Handle, + register: u32, +} + +impl RegisterSet { + pub fn register(&self) -> A::Register { + self.arch.borrow().register_from_id(self.register).unwrap() + } +} + +impl CoreArrayProvider for RegisterSet { + type Raw = u32; + type Context = A::Handle; + type Wrapped<'a> = Self; +} + +unsafe impl CoreArrayProviderInner for RegisterSet { + unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) { + // NOTE don't confuse this with BNFreeRegisterList + BNFreeRegisterSet(&mut BNRegisterSetWithConfidence { + regs: raw, + count, + confidence: 0, + }) + } + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> { + Self { + arch: context.clone(), + register: *raw, + } + } +}