246 Add Output Parameters to The Set of Pointer Variables (#247)

If the slicer enters a function due to following an output parameter, the corresponding parameter is added to the set of known pointers in that function.

mole/common/helper/function.py

@@ -22,7 +22,7 @@ class FunctionHelper:
         FunctionHelper.get_mlil_direct_call_insts.cache_clear()
         FunctionHelper.get_mlil_param_insts.cache_clear()
         FunctionHelper.get_var_map.cache_clear()
-        FunctionHelper.get_ptr_map.cache_clear()
+        FunctionHelper._get_ptr_map.cache_clear()
         FunctionHelper.get_ssa_memory_definitions.cache_clear()
         FunctionHelper.get_il_code.cache_clear()
         FunctionHelper.get_pseudo_c_code.cache_clear()
@@ -265,17 +265,16 @@ class FunctionHelper:
 
     @staticmethod
     @lru_cache(maxsize=maxsize)
-    def get_ptr_map(
-        func: bn.MediumLevelILFunction,
+    def _get_ptr_map(
+        func: bn.MediumLevelILFunction, param_idxs: frozenset[int]
     ) -> Dict[
         bn.SSAVariable,
         Tuple[Optional[bn.HighLevelILVar | bn.HighLevelILAddressOf], int],
     ]:
-        """
-        This method returns a dictionary, where keys are MLIL SSA variables corresponding to
-        pointers. The values are the corresponding HLIL_VAR/HLIL_ADDRESS_OF together with an offset.
-        """
-        func = func.ssa_form if func is not None else None
+        ptr_map: Dict[
+            bn.SSAVariable,
+            Tuple[Optional[bn.HighLevelILVar | bn.HighLevelILAddressOf], int],
+        ] = {}
 
         # Find MLIL SSA variables corresponding to pointers and their corresponding HLIL_VAR /
         # HLIL_ADDRESS_OF instructions
@@ -397,46 +396,91 @@ class FunctionHelper:
                                 )
             return None, (None, 0)
 
-        # Map MLIL SSA variables (corresponding to pointers) to their HLIL_VAR / HLIL_ADDRESS_OF
-        # instructions
-        ptr_map: Dict[
-            bn.SSAVariable,
-            Tuple[Optional[bn.HighLevelILVar | bn.HighLevelILAddressOf], int],
-        ] = {}
-        if func is not None:
-            var_map = FunctionHelper.get_var_map(func)
-            # Find pointers
-            for mlil_ptr_ssa_var, hlil_ptr_inst in func.traverse(find_mlil_ptrs):
-                mlil_ptr_ssa_var = mlil_ptr_ssa_var  # type: Optional[bn.SSAVariable]
-                hlil_ptr_inst = hlil_ptr_inst  # type: Optional[bn.HighLevelILInstruction]
-                if mlil_ptr_ssa_var is None or hlil_ptr_inst is None:
-                    continue
-                if mlil_ptr_ssa_var not in ptr_map:
-                    ptr_map[mlil_ptr_ssa_var] = (hlil_ptr_inst, 0)
-                # Find pointer aliases
-                for mlil_ptr_ssa_var_alias in var_map.get(mlil_ptr_ssa_var, set()):
-                    if mlil_ptr_ssa_var_alias not in ptr_map:
-                        ptr_map[mlil_ptr_ssa_var_alias] = (hlil_ptr_inst, 0)
-            # Find pointers with offsets
-            for mlil_ptr_ssa_var, (hlil_ptr_inst, hlil_ptr_offset) in func.traverse(
-                find_mlil_ptrs_offsets
-            ):
-                mlil_ptr_ssa_var = mlil_ptr_ssa_var  # type: Optional[bn.SSAVariable]
-                hlil_ptr_inst = hlil_ptr_inst  # type: Optional[bn.HighLevelILInstruction]
-                hlil_ptr_offset = hlil_ptr_offset  # type: int
-                if mlil_ptr_ssa_var is None or hlil_ptr_inst is None:
-                    continue
-                if mlil_ptr_ssa_var not in ptr_map:
-                    ptr_map[mlil_ptr_ssa_var] = (hlil_ptr_inst, hlil_ptr_offset)
-                # Find pointer aliases
-                for mlil_ptr_ssa_var_alias in var_map.get(mlil_ptr_ssa_var, set()):
-                    if mlil_ptr_ssa_var_alias not in ptr_map:
-                        ptr_map[mlil_ptr_ssa_var_alias] = (
-                            hlil_ptr_inst,
-                            hlil_ptr_offset,
-                        )
+        # Get variable map of the function
+        var_map = FunctionHelper.get_var_map(func)
+        # Get MLIL instructions of function parameters
+        mlil_param_insts = FunctionHelper.get_mlil_param_insts(func)
+        # Add known pointer parameters to the pointer map
+        for param_idx in param_idxs:
+            # Get the parameter's MLIL SSA variable and HLIL instruction
+            if not (0 < param_idx <= len(mlil_param_insts)):
+                continue
+            if (mlil_param_inst := mlil_param_insts[param_idx - 1]) is None:
+                continue
+            if (mlil_param_ssa_var := mlil_param_inst.var) is None:
+                continue
+            if (
+                hlil_param_inst := (
+                    mlil_param_inst.hlil
+                    if isinstance(mlil_param_inst.hlil, bn.HighLevelILVar)
+                    else None
+                )
+            ) is None:
+                continue
+            # Add parameter to the pointer map
+            if mlil_param_ssa_var not in ptr_map:
+                ptr_map[mlil_param_ssa_var] = (hlil_param_inst, 0)
+            # Add aliases of the parameter to the pointer map
+            for mlil_param_ssa_var_alias in var_map.get(mlil_param_ssa_var, set()):
+                if mlil_param_ssa_var_alias not in ptr_map:
+                    ptr_map[mlil_param_ssa_var_alias] = (hlil_param_inst, 0)
+        # Find pointers (without offsets) in the function and add them to the pointer map
+        for mlil_ptr_ssa_var, hlil_ptr_inst in func.traverse(find_mlil_ptrs):
+            # Get the pointer's MLIL SSA variable and HLIL instruction
+            mlil_ptr_ssa_var = mlil_ptr_ssa_var  # type: Optional[bn.SSAVariable]
+            hlil_ptr_inst = hlil_ptr_inst  # type: Optional[bn.HighLevelILInstruction]
+            if mlil_ptr_ssa_var is None or hlil_ptr_inst is None:
+                continue
+            # Add pointer to the pointer map
+            if mlil_ptr_ssa_var not in ptr_map:
+                ptr_map[mlil_ptr_ssa_var] = (hlil_ptr_inst, 0)
+            # Add aliases of the pointer to the pointer map
+            for mlil_ptr_ssa_var_alias in var_map.get(mlil_ptr_ssa_var, set()):
+                if mlil_ptr_ssa_var_alias not in ptr_map:
+                    ptr_map[mlil_ptr_ssa_var_alias] = (hlil_ptr_inst, 0)
+        # Find pointers with offsets in the function and add them to the pointer map
+        for mlil_ptr_ssa_var, (hlil_ptr_inst, hlil_ptr_offset) in func.traverse(
+            find_mlil_ptrs_offsets
+        ):
+            # Get the pointer's MLIL SSA variable, HLIL instruction and offset
+            mlil_ptr_ssa_var = mlil_ptr_ssa_var  # type: Optional[bn.SSAVariable]
+            hlil_ptr_inst = hlil_ptr_inst  # type: Optional[bn.HighLevelILInstruction]
+            hlil_ptr_offset = hlil_ptr_offset  # type: int
+            if mlil_ptr_ssa_var is None or hlil_ptr_inst is None:
+                continue
+            # Add pointer to the pointer map
+            if mlil_ptr_ssa_var not in ptr_map:
+                ptr_map[mlil_ptr_ssa_var] = (hlil_ptr_inst, hlil_ptr_offset)
+            # Add aliases of the pointer to the pointer map
+            for mlil_ptr_ssa_var_alias in var_map.get(mlil_ptr_ssa_var, set()):
+                if mlil_ptr_ssa_var_alias not in ptr_map:
+                    ptr_map[mlil_ptr_ssa_var_alias] = (hlil_ptr_inst, hlil_ptr_offset)
+        # Return pointer map
         return ptr_map
 
+    @staticmethod
+    def get_ptr_map(
+        func: bn.MediumLevelILFunction,
+        param_idxs: Set[int] = set(),
+    ) -> Dict[
+        bn.SSAVariable,
+        Tuple[Optional[bn.HighLevelILVar | bn.HighLevelILAddressOf], int],
+    ]:
+        """
+        This method finds MLIL SSA variables in function `func` that correspond to pointers.
+        Function parameters known to be pointers can be specified in `param_idxs` (set of parameter
+        indices). A dictionary is returned, where keys are MLIL SSA variables corresponding to
+        pointers. The values are the corresponding HLIL_VAR/HLIL_ADDRESS_OF together with an offset.
+        """
+        # Get the function in SSA form
+        if func is None or func.ssa_form is None:
+            return {}
+        func = func.ssa_form
+        # Make parameter indices immutable
+        param_idxs = frozenset(param_idxs) if param_idxs else frozenset({0})
+        # Find pointers in the function and add them to the pointer map
+        return FunctionHelper._get_ptr_map(func, param_idxs)
+
     @staticmethod
     @lru_cache(maxsize=maxsize)
     def get_ssa_memory_definitions(

mole/core/call.py

@@ -14,7 +14,10 @@ class MediumLevelILCallFrame:
 
     def __init__(self, func: bn.MediumLevelILFunction) -> None:
         self.func = func
-        self.func_params: Set[int] = set()
+        # Set of parameter indices that should be sliced when going back to the caller
+        self.hit_param_idxs: Set[int] = set()
+        # Set of parameter indices known to be output parameters
+        self.out_param_idxs: Set[int] = set()
         self.inst_stack: List[bn.MediumLevelILInstruction] = []
         self.last_inst: bn.MediumLevelILInstruction = None
         self.inst_graph: MediumLevelILInstructionGraph = MediumLevelILInstructionGraph()
@@ -46,6 +49,14 @@ class MediumLevelILCallTracker:
         )
         return
 
+    def get_out_params(self) -> Set[int]:
+        """
+        This method returns the set of output parameter indices of the current call frame.
+        """
+        if self._call_stack:
+            return self._call_stack[-1].out_param_idxs
+        return set()
+
     def get_call_level(self) -> int:
         """
         This method returns the current call level.
@@ -93,18 +104,18 @@ class MediumLevelILCallTracker:
     def push_func(
         self,
         to_inst: bn.MediumLevelILInstruction,
-        reverse: bool = False,
+        downwards: bool = True,
         param_idx: int = 0,
     ) -> bool:
         """
         This method creates a new call frame with the function `to_inst.function` and pushes it to
-        the top of the call stack. Further, it updates the call graph. If `reverse` is `False`,
-        `to_inst.function` is treated as the callee; if `reverse` is `True`, it is treated as the
-        caller. When traversing down the call graph (`reverse==False`), `param_idx` indicates the
+        the top of the call stack. Further, it updates the call graph. If `downwards` is `True`,
+        `to_inst.function` is treated as the callee; if `downwards` is `False`, it is treated as the
+        caller. When traversing down the call graph (`downwards==True`), `param_idx` indicates the
         callee's output parameter that was followed, or `0` if the traversal followed a return
-        instruction. When traversing up the call graph (`reverse==True`), `param_idx` indicates the
-        caller's relevant parameter. The function returns `True` in case of recursion, `False`
-        otherwise.
+        instruction. When traversing up the call graph (`downwards==False`), `param_idx` indicates
+        which parameter of the caller is relevant (marked in the graph). The function returns `True`
+        in case of recursion, `False` otherwise.
         """
         # Get the `to_inst`'s function
         func = to_inst.function
@@ -122,15 +133,23 @@ class MediumLevelILCallTracker:
         if not recursion:
             # Update call graph
             if len(self._call_stack) >= 2:
-                if not reverse:
+                if downwards:
                     caller_frame = self._call_stack[-2]
                     self._call_graph.add_edge(
-                        caller_frame.func, func, downwards=True, param_idx=param_idx
+                        caller_frame.func,
+                        func,
+                        downwards=downwards,
+                        param_idx=param_idx,
                     )
+                    if param_idx > 0:
+                        new_call_frame.out_param_idxs.add(param_idx)
                 else:
                     callee_frame = self._call_stack[-2]
                     self._call_graph.add_edge(
-                        func, callee_frame.func, downwards=False, param_idx=param_idx
+                        func,
+                        callee_frame.func,
+                        downwards=downwards,
+                        param_idx=param_idx,
                     )
             else:
                 self._call_graph.add_node(func)
@@ -139,7 +158,7 @@ class MediumLevelILCallTracker:
     def pop_func(self) -> Set[int]:
         """
         This method pops the top call frame from the call stack and returns a set of parameter
-        indices (`func_params`) that should be sliced further.
+        indices (`hit_param_idxs`) that should be sliced further.
         """
         if self._call_stack:
             # Pop old call frame and get its last instruction
@@ -168,7 +187,7 @@ class MediumLevelILCallTracker:
             old_inst_graph = nx.relabel_nodes(old_call_frame.inst_graph, mapping)
             self._inst_graph = nx.compose(self._inst_graph, old_inst_graph)
             # Return indices of parameters to be sliced further
-            return old_call_frame.func_params
+            return old_call_frame.hit_param_idxs
         return set()
 
     def push_inst(
@@ -216,14 +235,14 @@ class MediumLevelILCallTracker:
             self._call_stack[-1].mem_def_insts.append(inst)
         return
 
-    def add_func_param(self, param_idx: int) -> None:
+    def add_hit_param(self, param_idx: int) -> None:
         """
-        This method adds the given parameter to the `func_params` set of the current call frame.
-        `func_params` is the set of parameters that should be sliced when returning back to the
-        caller of the current function.
+        This method adds the given parameter to the `hit_param_idxs` set of the current call frame.
+        `hit_param_idxs` is the set of parameter indices that should be sliced when returning back
+        to the caller of the current function.
         """
         if self._call_stack:
-            self._call_stack[-1].func_params.add(param_idx)
+            self._call_stack[-1].hit_param_idxs.add(param_idx)
         return
 
     def print_call_stack(self) -> None:
@@ -263,7 +282,12 @@ class MediumLevelILCallTracker:
         This method prints the instruction slice (for debugging).
         """
         for call_level, call_frame in enumerate(self._call_stack):
+            hit_params = ", ".join(str(i) for i in sorted(call_frame.hit_param_idxs))
+            out_params = ", ".join(str(i) for i in sorted(call_frame.out_param_idxs))
             print(f"[{call_level:d}] {str(call_frame):s}")
+            print(f"Hit params: {hit_params:s}")
+            print(f"Out params: {out_params:s}")
+            print("Inst stack:")
             for inst in call_frame.inst_stack:
                 inst_info = InstructionHelper.get_inst_info(inst, False)
                 print(f"  - {inst_info:s}")

mole/core/slice.py

@@ -141,7 +141,10 @@ class MediumLevelILBackwardSlicer:
                         # Parameter is a variable
                         case bn.MediumLevelILVarSsa(var=param_ssa_var):
                             # Get pointers in the current function
-                            ptr_map = FunctionHelper.get_ptr_map(param_inst.function)
+                            out_params = self._call_tracker.get_out_params()
+                            ptr_map = FunctionHelper.get_ptr_map(
+                                param_inst.function, out_params
+                            )
                             # Get pointer instructions for `ssa_var` and `param_ssa_var`
                             ptr_inst_ssa_var, ptr_offset_ssa_var = ptr_map.get(
                                 ssa_var, (None, 0)
@@ -267,7 +270,7 @@ class MediumLevelILBackwardSlicer:
                     to_inst = call_inst.params[param_idx - 1]
                     recursion = self._call_tracker.push_func(
                         to_inst=to_inst,
-                        reverse=True,
+                        downwards=False,
                         param_idx=param_idx,
                     )
                     from_inst_info = InstructionHelper.get_inst_info(from_inst, False)
@@ -289,7 +292,7 @@ class MediumLevelILBackwardSlicer:
                     self._tag,
                     f"Follow parameter {param_idx:d} '{var_info:s}' when going back to specific caller",
                 )
-                self._call_tracker.add_func_param(param_idx)
+                self._call_tracker.add_hit_param(param_idx)
         return
 
     def _slice_backwards(
@@ -719,7 +722,7 @@ class MediumLevelILBackwardSlicer:
                                 # Proceed slicing the relevant output parameters, if we followed the
                                 # call due to reaching them
                                 if call_params:
-                                    # Get instructions defining the memory version of `inst`
+                                    # Get instructions defining the memory version of `ret_inst`
                                     mem_def_insts = (
                                         FunctionHelper.get_ssa_memory_definitions(
                                             dest_func,
@@ -793,7 +796,7 @@ class MediumLevelILBackwardSlicer:
                                                         # defining instruction
                                                         recursion = self._call_tracker.push_func(
                                                             to_inst=mem_def_inst,
-                                                            reverse=False,
+                                                            downwards=True,
                                                             param_idx=param_idx,
                                                         )
                                                         # Follow the output parameter defining
@@ -857,7 +860,7 @@ class MediumLevelILBackwardSlicer:
                                                     recursion = (
                                                         self._call_tracker.push_func(
                                                             to_inst=mem_def_inst,
-                                                            reverse=False,
+                                                            downwards=True,
                                                             param_idx=param_idx,
                                                         )
                                                     )
@@ -886,7 +889,7 @@ class MediumLevelILBackwardSlicer:
                                     # Push callee and proceed slicing its return instruction (if no recursion)
                                     recursion = self._call_tracker.push_func(
                                         to_inst=ret_inst,
-                                        reverse=False,
+                                        downwards=True,
                                         param_idx=0,
                                     )
                                     if recursion:
@@ -961,7 +964,7 @@ class MediumLevelILBackwardSlicer:
         This method backward slices the instruction `inst`.
         """
         self._call_tracker = MediumLevelILCallTracker()
-        self._call_tracker.push_func(to_inst=inst, reverse=True, param_idx=0)
+        self._call_tracker.push_func(to_inst=inst, downwards=False, param_idx=0)
         deque(
             inst.ssa_form.traverse(lambda inst: self._slice_backwards(inst)),
             maxlen=0,

plugin.json

@@ -24,7 +24,7 @@
     "Linux": "",
     "Windows": ""
   },
-  "version": "0.5.1",
+  "version": "0.5.2",
   "author": "Damian Pfammatter and Sergio Paganoni",
   "minimumbinaryninjaversion": 7290
 }

pyproject.toml

@@ -11,7 +11,7 @@ include = ["mole*"]
 
 [project]
 name = "mole"
-version = "0.5.1"
+version = "0.5.2"
 description = "A Binary Ninja plugin to identify interesting paths using static backward slicing"
 authors = [
     {name = "Damian Pfammatter"},

tests/data/src/function_out_params-09.c

@@ -0,0 +1,25 @@
+#include <stdio.h>
+#include <string.h>
+
+#define BUF_SIZE 8
+
+/*
+Testcase Description:
+- Output parameter 1 (size_t*): written and influence on the sink
+- Function output parameters are pointers
+*/
+
+__attribute__ ((noinline, optimize("O0")))
+void read_size(size_t* size) {
+    fscanf(stdin, "%zu", size);
+}
+
+int main() {
+    char dest[BUF_SIZE];
+    char src[] = "Hello, World!";
+    size_t size = 0;
+    read_size(&size);
+    memcpy(dest, src, size);
+    printf("%s\n", dest);
+    return 0;
+}

tests/slicing/test_function_out_params.py

@@ -83,3 +83,14 @@ class TestFunctionOutParams(TestSlicing):
             bv_callback=manually_set_types,
         )
         return
+
+    def test_function_out_params_09(
+        self, filenames: List[str] = ["function_out_params-09"]
+    ) -> None:
+        self.assert_paths(
+            srcs=[("fscanf", None), ("__isoc99_fscanf", None)],
+            snks=[("memcpy", 3)],
+            call_chains=[["main", "read_size"]],
+            filenames=filenames,
+        )
+        return