From b2e4cc5828309fdddc5db00a3e2ded8045165f7e Mon Sep 17 00:00:00 2001
From: Mason Reed <mason@vector35.com>
Date: Sun, 19 Jan 2025 18:02:58 -0500
Subject: [PATCH] More rust cleanup

- 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.
---
 arch/msp430/src/lift.rs          | 12 ++---
 arch/riscv/src/lib.rs            | 24 ++++------
 rust/src/low_level_il/lifting.rs | 81 ++++++++++++++++++--------------
 3 files changed, 57 insertions(+), 60 deletions(-)

diff --git a/arch/msp430/src/lift.rs b/arch/msp430/src/lift.rs
index 4fcc146a1..6ad7b67bd 100644
--- a/arch/msp430/src/lift.rs
+++ b/arch/msp430/src/lift.rs
@@ -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;
@@ -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)
@@ -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);
         }
     };
 }
@@ -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();
@@ -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)
diff --git a/arch/riscv/src/lib.rs b/arch/riscv/src/lib.rs
index b470c36e2..6d8a31ca8 100644
--- a/arch/riscv/src/lib.rs
+++ b/arch/riscv/src/lib.rs
@@ -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,
@@ -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)),
                 }
@@ -1289,12 +1287,12 @@ 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();
@@ -1302,14 +1300,10 @@ impl<D: 'static + RiscVDisassembler + Send + Sync> architecture::Architecture fo
                 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);
                 }
             }
 
@@ -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();
@@ -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)
diff --git a/rust/src/low_level_il/lifting.rs b/rust/src/low_level_il/lifting.rs
index 0803d41d9..61e514c15 100644
--- a/rust/src/low_level_il/lifting.rs
+++ b/rust/src/low_level_il/lifting.rs
@@ -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>,
@@ -1119,8 +1119,10 @@ 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))
     }
@@ -1128,7 +1130,7 @@ where
     // 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;
 
@@ -1136,7 +1138,8 @@ where
         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))
     }
@@ -1549,7 +1552,7 @@ 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();
@@ -1557,43 +1560,50 @@ where
         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.
@@ -1602,7 +1612,7 @@ pub struct Label {
     pub operand: usize,
 }
 
-impl Label {
+impl LowLevelILLabel {
     pub fn new() -> Self {
         use binaryninjacore_sys::BNLowLevelILInitLabel;
 
@@ -1612,9 +1622,10 @@ 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,
@@ -1622,8 +1633,8 @@ impl From<BNLowLevelILLabel> for Label {
     }
 }
 
-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,
@@ -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()
     }