Skip to content

Commit

Permalink
change UserVariableValue to the new standard
Browse files Browse the repository at this point in the history
  • Loading branch information
rbran committed May 31, 2024
1 parent 0c9ecbe commit 37f168e
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 39 deletions.
12 changes: 5 additions & 7 deletions rust/src/mlil/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ use crate::function::{Function, Location};
use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Ref, RefCountable};
use crate::string::BnStrCompatible;
use crate::types::{
Conf, PossibleValueSet, RegisterValue, SSAVariable, Type, UserVariableValues, Variable,
Conf, PossibleValueSet, RegisterValue, SSAVariable, Type, UserVariableValue, Variable,
};

use super::{MediumLevelILBlock, MediumLevelILInstruction, MediumLevelILLiftedInstruction};
Expand Down Expand Up @@ -242,20 +242,18 @@ impl MediumLevelILFunction {

/// Returns a map of current defined user variable values.
/// Returns a Map of user current defined user variable values and their definition sites.
pub fn user_var_values(&self) -> UserVariableValues {
pub fn user_var_values(&self) -> Array<UserVariableValue> {
let mut count = 0;
let function = self.get_function();
let var_values = unsafe { BNGetAllUserVariableValues(function.handle, &mut count) };
assert!(!var_values.is_null());
UserVariableValues {
vars: core::ptr::slice_from_raw_parts(var_values, count),
}
unsafe { Array::new(var_values, count, ()) }
}

/// Clear all user defined variable values.
pub fn clear_user_var_values(&self) -> Result<(), ()> {
for (var, arch_and_addr, _value) in self.user_var_values().all() {
self.clear_user_var_value(&var, arch_and_addr.address)?;
for user_value in self.user_var_values().iter() {
self.clear_user_var_value(&user_value.var, user_value.def_site.address)?;
}
Ok(())
}
Expand Down
87 changes: 55 additions & 32 deletions rust/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3355,55 +3355,78 @@ impl Drop for LookupTableEntryRaw {
/////////////////////////
// ArchAndAddr

#[repr(C)]
#[derive(Copy, Clone, Eq, Hash, PartialEq)]
pub struct ArchAndAddr {
pub arch: CoreArchitecture,
pub address: u64,
}

impl ArchAndAddr {
pub(crate) unsafe fn from_raw(value: BNArchitectureAndAddress) -> Self {
mem::transmute(value)
}
}

/////////////////////////
// UserVariableValues
// UserVariableValue

pub struct UserVariableValue {
pub var: Variable,
pub def_site: ArchAndAddr,
pub possible_value: PossibleValueSet,
}

impl UserVariableValue {
pub(crate) unsafe fn from_raw(value: BNUserVariableValue) -> Self {
let var = Variable::from_raw(value.var);
let def_site = ArchAndAddr::from_raw(value.defSite);
let value = PossibleValueSet::from_raw(value.value);
Self {
var,
def_site,
possible_value: value,
}
}
}

pub struct UserVariableValues {
pub(crate) vars: *const [BNUserVariableValue],
impl CoreArrayProvider for UserVariableValue {
type Raw = BNUserVariableValue;
type Context = ();
type Wrapped<'a> = Self;
}

impl UserVariableValues {
pub fn into_hashmap(self) -> HashMap<Variable, HashMap<ArchAndAddr, PossibleValueSet>> {
unsafe impl CoreArrayProviderInner for UserVariableValue {
unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
unsafe { BNFreeUserVariableValues(raw) };
}

unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> {
Self::from_raw(*raw)
}
}

impl Array<UserVariableValue> {
pub fn into_hashmap(&self) -> HashMap<Variable, HashMap<ArchAndAddr, PossibleValueSet>> {
let mut result: HashMap<Variable, HashMap<ArchAndAddr, PossibleValueSet>> = HashMap::new();
for (var, def_site, possible_val) in self.all() {
for value in self.iter()
{
result
.entry(var)
.entry(value.var)
.or_default()
.entry(def_site)
.or_insert(possible_val);
.entry(value.def_site)
.or_insert(value.possible_value);
}
result
}
pub fn all(&self) -> impl Iterator<Item = (Variable, ArchAndAddr, PossibleValueSet)> {
unsafe { &*self.vars }.iter().map(|var_val| {
let var = unsafe { Variable::from_raw(var_val.var) };
let def_site = ArchAndAddr {
arch: unsafe { CoreArchitecture::from_raw(var_val.defSite.arch) },
address: var_val.defSite.address,
};
let possible_val = unsafe { PossibleValueSet::from_raw(var_val.value) };
(var, def_site, possible_val)
})
}
pub fn values_from_variable(
&self,
var: Variable,
) -> impl Iterator<Item = (ArchAndAddr, PossibleValueSet)> {
self.all()
.filter(move |(t_var, _, _)| t_var == &var)
.map(|(_var, def_site, possible_val)| (def_site, possible_val))
}
}

impl Drop for UserVariableValues {
fn drop(&mut self) {
unsafe { BNFreeUserVariableValues(self.vars as *mut BNUserVariableValue) };
pub fn values_from_variable<'a>(
&'a self,
var: Variable,
) -> impl Iterator<Item = (ArchAndAddr, PossibleValueSet)> + 'a {
self.iter()
.filter(move |value| value.var == var)
.map(|value| (value.def_site, value.possible_value))
}
}

Expand Down

0 comments on commit 37f168e

Please sign in to comment.