Skip to content

Commit

Permalink
More rust cleanup
Browse files Browse the repository at this point in the history 
- Rename Label to LowLevelILLabel
- Update the functions label map automatically

This fixed a major api blunder where the label map is returned as a reference and originally resulted in UB prone lifetime semantics. It was temporarily "fixed" with explicit label updates in the architecture implementation code. But that was less than ideal and was easy to mess up. Now the label map will be updated automatically as the location of labels is now tracked.
  • Loading branch information
@emesare
emesare committed Jan 19, 2025
1 parent fc17319 commit b2e4cc5
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 60 deletions.
12 changes: 3 additions & 9 deletions arch/msp430/src/lift.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use crate::flag::{Flag, FlagWrite};
use crate::register::Register;
use crate::Msp430;

use binaryninja::{architecture::FlagCondition, low_level_il::lifting::Label};
use binaryninja::{architecture::FlagCondition, low_level_il::lifting::LowLevelILLabel};

use msp430_asm::emulate::Emulated;
use msp430_asm::instruction::Instruction;
Expand Down Expand Up @@ -142,12 +142,12 @@ macro_rules! conditional_jump {

let mut true_label = $il.label_for_address(true_addr).unwrap_or_else(|| {
new_true = true;
Label::new()
LowLevelILLabel::new()
});

let mut false_label = $il.label_for_address(false_addr).unwrap_or_else(|| {
new_false = true;
Label::new()
LowLevelILLabel::new()
});

$il.if_expr($cond, &mut true_label, &mut false_label)
Expand All @@ -156,14 +156,10 @@ macro_rules! conditional_jump {
if new_true {
$il.mark_label(&mut true_label);
$il.jump($il.const_ptr(true_addr)).append();
} else {
$il.update_label_for_address(true_addr, true_label);
}

if new_false {
$il.mark_label(&mut false_label);
} else {
$il.update_label_for_address(false_addr, false_label);
}
};
}
Expand Down Expand Up @@ -289,7 +285,6 @@ pub(crate) fn lift_instruction(
match label {
Some(mut label) => {
il.goto(&mut label).append();
il.update_label_for_address(fixed_addr, label);
}
None => {
il.jump(il.const_ptr(fixed_addr)).append();
Expand Down Expand Up @@ -424,7 +419,6 @@ pub(crate) fn lift_instruction(
let dest = if let Some(Operand::Immediate(dest)) = inst.destination() {
if let Some(mut label) = il.label_for_address(*dest as u64) {
il.goto(&mut label).append();
il.update_label_for_address(*dest as u64, label);
return;
} else {
il.const_ptr(*dest as u64)
Expand Down
24 changes: 8 additions & 16 deletions arch/riscv/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,9 @@ use binaryninja::confidence::{Conf, MAX_CONFIDENCE, MIN_CONFIDENCE};
use binaryninja::logger::Logger;
use binaryninja::low_level_il::expression::{LowLevelILExpressionKind, ValueExpr};
use binaryninja::low_level_il::instruction::LowLevelILInstructionKind;
use binaryninja::low_level_il::lifting::{Label, LiftableLowLevelIL, LiftableLowLevelILWithSize};
use binaryninja::low_level_il::lifting::{
LiftableLowLevelIL, LiftableLowLevelILWithSize, LowLevelILLabel,
};
use binaryninja::low_level_il::{
expression::ExpressionHandler, instruction::InstructionHandler, LowLevelILRegister,
MutableLiftedILExpr, MutableLiftedILFunction, RegularLowLevelILFunction,
Expand Down Expand Up @@ -1220,11 +1222,7 @@ impl<D: 'static + RiscVDisassembler + Send + Sync> architecture::Architecture fo
let target = addr.wrapping_add(j.imm() as i64 as u64);

match (j.rd().id(), il.label_for_address(target)) {
(0, Some(mut l)) => {
let jump_expr = il.goto(&mut l);
il.update_label_for_address(target, l);
jump_expr
}
(0, Some(mut l)) => il.goto(&mut l),
(0, None) => il.jump(il.const_ptr(target)),
(_, _) => il.call(il.const_ptr(target)),
}
Expand Down Expand Up @@ -1289,27 +1287,23 @@ impl<D: 'static + RiscVDisassembler + Send + Sync> architecture::Architecture fo

let mut f = il.label_for_address(ft).unwrap_or_else(|| {
new_false = true;
Label::new()
LowLevelILLabel::new()
});

let mut t = il.label_for_address(tt).unwrap_or_else(|| {
new_true = true;
Label::new()
LowLevelILLabel::new()
});

il.if_expr(cond_expr, &mut t, &mut f).append();

if new_true {
il.mark_label(&mut t);
il.jump(il.const_ptr(tt)).append();
} else {
il.update_label_for_address(tt, t);
}

if new_false {
il.mark_label(&mut f);
} else {
il.update_label_for_address(ft, f);
}
}

Expand Down Expand Up @@ -1461,14 +1455,14 @@ impl<D: 'static + RiscVDisassembler + Send + Sync> architecture::Architecture fo
il.set_reg(max_width, dest_reg, il.unimplemented()).append();

let mut new_false = false;
let mut t = Label::new();
let mut t = LowLevelILLabel::new();

let cond_expr = il.cmp_e(max_width, dest_reg, 0u64);

let ft = addr.wrapping_add(inst_len);
let mut f = il.label_for_address(ft).unwrap_or_else(|| {
new_false = true;
Label::new()
LowLevelILLabel::new()
});

il.if_expr(cond_expr, &mut t, &mut f).append();
Expand All @@ -1480,8 +1474,6 @@ impl<D: 'static + RiscVDisassembler + Send + Sync> architecture::Architecture fo

if new_false {
il.mark_label(&mut f);
} else {
il.update_label_for_address(ft, f);
}
}
Op::AmoSwap(a)
Expand Down
81 changes: 46 additions & 35 deletions rust/src/low_level_il/lifting.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1097,8 +1097,8 @@ where
pub fn if_expr<'a: 'b, 'b, C>(
&'a self,
cond: C,
true_label: &'b mut Label,
false_label: &'b mut Label,
true_label: &'b mut LowLevelILLabel,
false_label: &'b mut LowLevelILLabel,
) -> LowLevelILExpression<'a, A, Mutable, NonSSA<LiftedNonSSA>, VoidExpr>
where
C: LiftableLowLevelIL<'b, A, Result = ValueExpr>,
Expand All @@ -1119,24 +1119,27 @@ where
};

// Update the labels after they have been resolved.
*true_label = Label::from(raw_true_label);
*false_label = Label::from(raw_false_label);
*true_label = LowLevelILLabel::from(raw_true_label);
*false_label = LowLevelILLabel::from(raw_false_label);
self.update_label_map_for_label(true_label);
self.update_label_map_for_label(false_label);

LowLevelILExpression::new(self, LowLevelExpressionIndex(expr_idx))
}

// TODO: Wtf are these lifetimes??
pub fn goto<'a: 'b, 'b>(
&'a self,
label: &'b mut Label,
label: &'b mut LowLevelILLabel,
) -> LowLevelILExpression<'a, A, Mutable, NonSSA<LiftedNonSSA>, VoidExpr> {
use binaryninjacore_sys::BNLowLevelILGoto;

let mut raw_label = BNLowLevelILLabel::from(*label);
let expr_idx = unsafe { BNLowLevelILGoto(self.handle, &mut raw_label) };

// Update the labels after they have been resolved.
*label = Label::from(raw_label);
*label = LowLevelILLabel::from(raw_label);
self.update_label_map_for_label(label);

LowLevelILExpression::new(self, LowLevelExpressionIndex(expr_idx))
}
Expand Down Expand Up @@ -1549,51 +1552,58 @@ where
}
}

pub fn label_for_address<L: Into<Location>>(&self, loc: L) -> Option<Label> {
pub fn label_for_address<L: Into<Location>>(&self, loc: L) -> Option<LowLevelILLabel> {
use binaryninjacore_sys::BNGetLowLevelILLabelForAddress;

let loc: Location = loc.into();
let arch = loc.arch.unwrap_or_else(|| *self.arch().as_ref());
let raw_label =
unsafe { BNGetLowLevelILLabelForAddress(self.handle, arch.handle, loc.addr) };
match raw_label.is_null() {
false => Some(unsafe { Label::from(*raw_label) }),
false => {
let mut label = unsafe { LowLevelILLabel::from(*raw_label) };
// Set the location so that calls to [Self::update_label_map_for_label] will update the label map.
label.location = Some(loc);
Some(label)
}
true => None,
}
}

// TODO: Make this private and then force the updates in the expressions.
/// Call this after updating the label through an il operation or via [`Self::mark_label`].
///
/// If you retrieved a label via [`Self::label_for_address`] than you very likely want to use this.
pub fn update_label_for_address<L: Into<Location>>(&self, loc: L, label: Label) {
fn update_label_map_for_label(&self, label: &LowLevelILLabel) {
use binaryninjacore_sys::BNGetLowLevelILLabelForAddress;

let loc: Location = loc.into();
let arch = loc.arch.unwrap_or_else(|| *self.arch().as_ref());
// Add the label into the label map
unsafe { BNAddLowLevelILLabelForAddress(self.handle, arch.handle, loc.addr) };
// Retrieve a pointer to the label in the map
let raw_label =
unsafe { BNGetLowLevelILLabelForAddress(self.handle, arch.handle, loc.addr) };
// We should always have a valid label here
assert!(!raw_label.is_null(), "Failed to add label for address!");
// Update the label in the map with `label`
unsafe { *raw_label = label.into() };
// Only need to update the label if there is an associated address.
if let Some(loc) = label.location {
let loc: Location = loc.into();
let arch = loc.arch.unwrap_or_else(|| *self.arch().as_ref());
// Add the label into the label map
unsafe { BNAddLowLevelILLabelForAddress(self.handle, arch.handle, loc.addr) };
// Retrieve a pointer to the label in the map
let raw_label =
unsafe { BNGetLowLevelILLabelForAddress(self.handle, arch.handle, loc.addr) };
// We should always have a valid label here
assert!(!raw_label.is_null(), "Failed to add label for address!");
// Update the label in the map with `label`
unsafe { *raw_label = label.into() };
}
}

pub fn mark_label(&self, label: &mut Label) {
pub fn mark_label(&self, label: &mut LowLevelILLabel) {
use binaryninjacore_sys::BNLowLevelILMarkLabel;

let mut raw_label = BNLowLevelILLabel::from(*label);
unsafe { BNLowLevelILMarkLabel(self.handle, &mut raw_label) };
*label = Label::from(raw_label);
*label = LowLevelILLabel::from(raw_label);
self.update_label_map_for_label(label);
}
}

// TODO: Rename to LowLevelILLabel
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Label {
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct LowLevelILLabel {
/// Used to update the label map if the label is associated with a location.
pub location: Option<Location>,
pub resolved: bool,
// TODO: This expr_ref is not actually a valid one sometimes...
// TODO: We should make these non public and only accessible if resolved is true.
Expand All @@ -1602,7 +1612,7 @@ pub struct Label {
pub operand: usize,
}

impl Label {
impl LowLevelILLabel {
pub fn new() -> Self {
use binaryninjacore_sys::BNLowLevelILInitLabel;

Expand All @@ -1612,18 +1622,19 @@ impl Label {
}
}

impl From<BNLowLevelILLabel> for Label {
impl From<BNLowLevelILLabel> for LowLevelILLabel {
fn from(value: BNLowLevelILLabel) -> Self {
Self {
location: None,
resolved: value.resolved,
expr_ref: LowLevelExpressionIndex(value.ref_),
operand: value.operand,
}
}
}

impl From<Label> for BNLowLevelILLabel {
fn from(value: Label) -> Self {
impl From<LowLevelILLabel> for BNLowLevelILLabel {
fn from(value: LowLevelILLabel) -> Self {
Self {
resolved: value.resolved,
ref_: value.expr_ref.0,
Expand All @@ -1632,13 +1643,13 @@ impl From<Label> for BNLowLevelILLabel {
}
}

impl From<&Label> for BNLowLevelILLabel {
fn from(value: &Label) -> Self {
impl From<&LowLevelILLabel> for BNLowLevelILLabel {
fn from(value: &LowLevelILLabel) -> Self {
Self::from(*value)
}
}

impl Default for Label {
impl Default for LowLevelILLabel {
fn default() -> Self {
Self::new()
}
Expand Down

0 comments on commit b2e4cc5

Please sign in to comment.