Add inline remapping for function lifter

Author: zznop

rust/src/architecture.rs

@@ -600,9 +600,9 @@ pub struct FunctionLifterContext {
     pub blocks: Vec<Ref<BasicBlock<NativeBlock>>>,
     pub no_return_calls: HashSet<Location>,
     pub contextual_returns: HashMap<Location, bool>,
-    //pub inline_remapping: HashMap<ArchAndAddr, ArchAndAddr>>,
-    //pub user_indirect_branches: HashMap<ArchAndAddr, HashSet<ArchAndAddr>>>,
-    //pub auto_indirect_branches: HashMap<ArchAndAddr, HashSet<ArchAndAddr>>>,
+    pub inlined_remapping: HashMap<Location, Location>,
+    pub user_indirect_branches: HashMap<Location, HashSet<Location>>,
+    pub auto_indirect_branches: HashMap<Location, HashSet<Location>>,
     //pub inlined_calls: HashSet<u64>,
 }
 
@@ -651,6 +651,43 @@ impl FunctionLifterContext {
             .zip(raw_contextual_return_vals.iter().copied())
             .collect();
 
+        let inlined_remapping: HashMap<Location, Location> = {
+            let raw_inline_remap_locs: &[BNArchitectureAndAddress] = std::slice::from_raw_parts(
+                flc_ref.inlinedRemappingKeys,
+                flc_ref.inlinedRemappingEntryCount,
+            );
+
+            let raw_inline_remap_dests: &[BNArchitectureAndAddress] = std::slice::from_raw_parts(
+                flc_ref.inlinedRemappingValues,
+                flc_ref.inlinedRemappingEntryCount,
+            );
+
+            raw_inline_remap_locs
+                .iter()
+                .map(Location::from)
+                .zip(raw_inline_remap_dests.iter().map(Location::from))
+                .collect()
+        };
+
+        let mut user_indirect_branches: HashMap<Location, HashSet<Location>> = HashMap::new();
+        let mut auto_indirect_branches: HashMap<Location, HashSet<Location>> = HashMap::new();
+        for i in 0..flc_ref.indirectBranchesCount {
+            let entry = unsafe { *flc_ref.indirectBranches.add(i) };
+            let src = Location::new(Some(CoreArchitecture::from_raw(entry.sourceArch)), entry.sourceAddr);
+            let dest = Location::new(Some(CoreArchitecture::from_raw(entry.destArch)), entry.destAddr);
+            if entry.autoDefined {
+                auto_indirect_branches
+                    .entry(src)
+                    .or_insert_with(HashSet::new)
+                    .insert(dest);
+            } else {
+                user_indirect_branches
+                    .entry(src)
+                    .or_insert_with(HashSet::new)
+                    .insert(dest);
+            }
+        }
+
         FunctionLifterContext {
             handle,
             function: BNNewLowLevelILFunctionReference(function),
@@ -659,6 +696,9 @@ impl FunctionLifterContext {
             blocks,
             no_return_calls,
             contextual_returns,
+            inlined_remapping,
+            user_indirect_branches,
+            auto_indirect_branches,
         }
     }
 

rust/src/low_level_il/function.rs

@@ -209,6 +209,20 @@ where
             Some(unsafe { BasicBlock::ref_from_raw(block, LowLevelILBlock { function: self }) })
         }
     }
+
+    pub fn set_indirect_branches(&self, branches: &Vec<Location>) {
+        let mut bn_branches: Box<[BNArchitectureAndAddress]> = branches
+            .iter()
+            .map(|loc| BNArchitectureAndAddress {
+                address: loc.addr,
+                arch: self.arch().handle,
+            })
+            .collect();
+
+        unsafe {
+            BNLowLevelILSetIndirectBranches(self.handle, bn_branches.as_mut_ptr(), branches.len());
+        }
+    }
 }
 
 impl<M: FunctionMutability> LowLevelILFunction<M, NonSSA> {