Cwe476 (#47)

refactored cwe476 to add stack tracking

CHANGES.md

+dev
+====
+
+-   Added a lot more test cases to acceptance tests (PR #46)
+-   Reworked CWE-476 check to track stack variables (PR #47)
+
 0.3 (2019-12)
 ====
 

src/analysis/mem_region.ml

@@ -152,3 +152,25 @@ let rec equal (mem_region1:'a t) (mem_region2:'a t) ~data_equal : bool =
     else
       false
   | _ -> false
+
+let map_data (mem_region: 'a t) ~(f: 'a -> 'b) : 'b t =
+  List.map mem_region ~f:(fun mem_node ->
+    { pos = mem_node.pos;
+      size = mem_node.size;
+      data = Result.map mem_node.data ~f
+    }
+  )
+
+let list_data (mem_region: 'a t) : 'a List.t =
+  List.filter_map mem_region ~f:(fun mem_node ->
+    match mem_node.data with
+    | Ok(value) -> Some(value)
+    | Error(_) -> None
+  )
+
+let list_data_pos (mem_region: 'a t) : (Bitvector.t * 'a) List.t =
+  List.filter_map mem_region ~f:(fun mem_node ->
+    match mem_node.data with
+    | Ok(value) -> Some( mem_node.pos, value )
+    | Error(_) -> None
+  )

src/analysis/mem_region.mli

@@ -37,3 +37,12 @@ val equal: 'a t -> 'a t -> data_equal:('a -> 'a -> bool) -> bool
 
 (** Mark an area in the mem_region as containing errors. *)
 val mark_error: 'a t -> pos:Bitvector.t -> size:Bitvector.t -> 'a t
+
+(** Map the contained data to new values. *)
+val map_data: 'a t -> f:('a -> 'b) -> 'b t
+
+(** List the contained data (ignoring error values). *)
+val list_data: 'a t -> 'a List.t
+
+(** List the contained data (ignoring error values) together with their positions. *)
+val list_data_pos: 'a t -> (Bitvector.t * 'a) List.t

src/analysis/type_inference.ml

@@ -196,10 +196,7 @@ module TypeInfo = struct
   let remove_virtual_registers state =
     { state with reg = Map.filter_keys state.reg ~f:(fun var -> Var.is_physical var) }
 
-(** if the addr_exp is a (computable) stack offset, return the offset. In cases where addr_expr
-    may or may not be a stack offset (i.e. offset of a register which may point to the stack or
-    to some other memory region), it still returns an offset. *)
-  let compute_stack_offset state addr_exp ~sub_tid ~project : Bitvector.t  Option.t =
+  let compute_stack_offset (state: t) (addr_exp: Exp.t) ~(sub_tid: Tid.t) ~(project: Project.t) : Bitvector.t Option.t =
     let (register, offset) = match addr_exp with
       | Bil.Var(var) -> (Some(var), Bitvector.of_int 0 ~width:(Symbol_utils.arch_pointer_size_in_bytes project * 8))
       | Bil.BinOp(Bil.PLUS, Bil.Var(var), Bil.Int(num)) -> (Some(var), num)
@@ -280,7 +277,7 @@ let value_of_exp exp =
   | _ -> None
 
 
-let rec type_of_exp exp (state: TypeInfo.t) ~sub_tid ~project =
+let rec type_of_exp (exp: Exp.t) (state: TypeInfo.t) ~(sub_tid: Tid.t) ~(project: Project.t) : (Register.t, unit) Result.t Option.t =
   let open Register in
   match exp with
   | Bil.Load(_) -> (* TODO: Right now only the stack is tracked for type infos. *)
@@ -321,7 +318,8 @@ let rec type_of_exp exp (state: TypeInfo.t) ~sub_tid ~project =
   | Bil.Unknown(_) -> None
   | Bil.Ite(_if_, then_, else_) -> begin
       match (type_of_exp then_ state ~sub_tid ~project, type_of_exp else_ state ~sub_tid ~project) with
-      | (Some(value1), Some(value2)) -> if value1 = value2 then Some(value1) else None
+      | (Some(Ok(value1)), Some(Ok(value2))) -> if Register.equal value1 value2 then Some(Ok(value1)) else None
+      | (Some(Error ()), Some(Error ())) -> Some(Error ())
       | _ -> None
     end
   | Bil.Extract(_) -> Some(Ok(Data)) (* TODO: Similar to cast: Are there cases of 32bit-64bit-address-conversions here? *)
@@ -643,7 +641,26 @@ let print_type_info_tags ~project ~tid_map =
                                (Tid.name (Term.tid block)) in
              Log_utils.error error_str
         )
-    )
+  )
+
+let get_type_info_of_block ~(project: Project.t) (block: Blk.t) ~(sub_tid: Tid.t) : TypeInfo.t Tid.Map.t =
+  match Term.get_attr block type_info_tag with
+  | Some(start_state) ->
+      let elements = Blk.elts block in
+      let (type_info_map, _) = Seq.fold elements ~init:(Tid.Map.empty, start_state) ~f:(fun (type_info_map, state) element ->
+        match element with
+        | `Phi _ -> (type_info_map, state)
+        | `Def term ->
+            let new_type_info_map = Tid.Map.set type_info_map ~key:(Term.tid term) ~data:state in
+            let new_state = update_type_info element state ~sub_tid ~project in
+            (new_type_info_map, new_state)
+        | `Jmp term ->
+            let new_type_info_map = Tid.Map.set type_info_map ~key:(Term.tid term) ~data:state in
+            let new_state = update_type_info element state ~sub_tid ~project in
+            (new_type_info_map, new_state)
+      ) in
+      type_info_map
+  | None -> failwith "[cwe_checker] Error: Tag not found"
 
 (* Functions made available for unit tests *)
 module Private = struct

src/analysis/type_inference.mli

@@ -45,6 +45,12 @@ module TypeInfo : sig
 
   (* Pretty Printer. At the moment, the output is not pretty at all. *)
   val pp: Format.formatter -> t -> unit
+
+  (** if the addr_exp is a (computable) stack offset, return the offset. In cases where addr_expr
+      may or may not be a stack offset (i.e. offset of a register which may point to the stack or
+      to some other memory region), it still returns an offset. *)
+  val compute_stack_offset: t -> Exp.t -> sub_tid:Tid.t -> project:Project.t -> Bitvector.t Option.t
+
 end
 
 (** A tag for TypeInfo.t, so that we can annotate basic blocks with known type information
@@ -65,6 +71,12 @@ val print_type_info_tags: project:Project.t -> tid_map:word Tid.Map.t -> unit
     which is internally used to mark which pointers point to the current stack frame.*)
 val update_type_info: Blk.elt -> TypeInfo.t -> sub_tid:Tid.t -> project:Project.t -> TypeInfo.t
 
+(** Get the type info for each def term and jump term in a block.
+    The returned type info for the tid of a (def/jmp) term is the tid before that term.
+    The sub_tid is the Tid of the current function which is internally used to mark
+    which pointers point to the current stack frame. *)
+val get_type_info_of_block: project:Project.t -> Blk.t -> sub_tid:Tid.t -> TypeInfo.t Tid.Map.t
+
 (* functions made available for unit tests: *)
 module Private : sig
   val update_block_analysis: Blk.t -> TypeInfo.t -> sub_tid:Tid.t -> project:Project.t -> TypeInfo.t

src/checkers/cwe_476.mli

@@ -8,9 +8,8 @@
 
     {1 How the check works}
 
-    We search for an execution path where a memory access using the return value of
-    a symbol happens before the return value is checked through a conditional
-    jump instruction.
+    Using dataflow analysis we search for an execution path where a memory access using the return value of
+    a symbol happens before the return value is checked through a conditional jump instruction.
 
     Note that the check relies on Bap-generated stubs to identify return registers of the
     checked functions. Therefore it only works for functions for which Bap generates
@@ -20,8 +19,14 @@
 
     - strict_call_policy=\{true, false\}: Determines behaviour on call and return instructions.
       If false, we assume that the callee, resp. the caller on a return instruction,
-      checks all unchecked values still contained in the registers. If true, every
+      checks all unchecked values still contained in parameter registers. If true, every
       unchecked value on a call or return instruction gets reported.
+    - strict_mem_policy=|{true, false|}:
+      Determines behaviour on writing an unchecked return value to a memory region other than the stack.
+      If true, these instances get reported.
+      Depending on the coding style, this can lead to a lot false positives if return values are
+      only checked after writing them to their target destination.
+      If false, these instances do not get reported, which in turn can lead to false negatives.
     - max_steps=<num>: Max number of steps for the dataflow fixpoint algorithm.
 
     {2 Symbols configurable in config.json}
@@ -31,9 +36,10 @@
 
     {1 False Positives}
 
-    - The check does not yet track values on the stack. Thus instances, where the
-    return value gets written onto the stack before the check happens get incorrectly
-    flagged. This happens a lot on unoptimized binaries but rarely on optimized ones.
+    - If strict_mem_policy is set to true, writing a return value to memory other than the stack
+    gets reported even if a NULL pointer check happens right afterwards.
+    - The check has no knowledge about the actual number of parameters that an extern function call takes.
+      This can lead to false positives if strict_call_policy is set to true.
 
     {1 False Negatives}
 
@@ -43,9 +49,41 @@
     for the return value being NULL or something else
     - For functions with more than one return value we do not distinguish between
     the return values.
+    - If strict_mem_policy is set to false, unchecked return values that are
+    saved somewhere other than the stack may be missed.
+    - The check has no knowledge about the actual number of parameters that an extern function call takes.
+      This can lead to false negatives, especially if function parameters are passed on the stack.
 *)
 
+open Bap.Std
+open Core_kernel
+
 val name : string
 val version : string
 
 val check_cwe : Bap.Std.program Bap.Std.term -> Bap.Std.project -> Bap.Std.word Bap.Std.Tid.Map.t -> string list list -> string list -> unit
+
+
+(**/**)
+(* Functions made public for unit tests *)
+module Private : sig
+
+  module Taint : module type of Tid.Set
+  module State : sig
+    type t
+    val empty: t
+    val set_register: t -> Var.t -> Taint.t -> t
+    val find_register: t -> Var.t -> Taint.t Option.t
+    val union: t -> t -> t
+  end
+
+  module StackInfo : sig
+    type t
+    val assemble_mock_info: Tid.t -> Project.t -> t
+  end
+
+  val flag_unchecked_return_values: State.t -> cwe_hits: Taint.t ref -> project: Project.t -> State.t
+  val flag_register_taints: State.t -> cwe_hits: Taint.t ref -> State.t
+  val flag_parameter_register: State.t -> cwe_hits: Taint.t ref -> project: Project.t -> State.t
+  val untaint_non_callee_saved_register: State.t -> project: Project.t -> State.t
+end

src/config.json

@@ -98,6 +98,7 @@
     "_comment1": "included functions of the following libs: stdlib.h, locale.h, stdio.h, cstring.h, wchar.h",
     "parameters": [
       "strict_call_policy=true",
+      "strict_memory_policy=false",
       "max_steps=100"
     ],
     "symbols": [

src/utils/cconv.ml

@@ -41,6 +41,58 @@ let is_callee_saved var project =
     callee_saved_registers := Some(String.Set.of_list (callee_saved_register_list project));
     String.Set.mem (Option.value_exn !callee_saved_registers) (Var.name var)
 
+(** Return a list of all registers that may hold function arguments. *)
+let get_parameter_register_list (project: Project.t) : String.t List.t =
+  let architecture = Project.arch project in
+  match architecture with
+  | `x86 ->
+      [] (* TODO: This is the value for the standard C calling convention. But it is incorrect for some of the other calling conventions for x86! *)
+  | `x86_64 -> (* System V ABI. TODO: Floationg Point registers are mising! TODO: Microsoft calling convention uses different register. *)
+      "RDI" :: "RSI" :: "RDX" :: "RCX" :: "R8" :: "R9" :: []
+  | `armv4 | `armv5 | `armv6 | `armv7
+  | `armv4eb | `armv5eb | `armv6eb | `armv7eb
+  | `thumbv4 | `thumbv5 | `thumbv6 | `thumbv7
+  | `thumbv4eb | `thumbv5eb | `thumbv6eb | `thumbv7eb ->
+      "R0" :: "R1" :: "R2" :: "R3" :: []
+  | `aarch64 | `aarch64_be -> (* ARM 64bit *)
+      "X0" :: "X1" :: "X2" :: "X3" :: "X4" :: "X5" :: "X6" :: "X7" :: []
+  | `ppc (* 32bit PowerPC *) (* TODO: add floating point register! *)
+  | `ppc64 | `ppc64le -> (* 64bit PowerPC *)
+      "R3" :: "R4" :: "R5" :: "R6" :: "R7" :: "R8" :: "R9" :: "R10" :: []
+  | `mips | `mips64 | `mips64el | `mipsel -> (* TODO: MIPS has also a calling convention with less arguments. TODO: check whether BAP actually uses A4-A7 as register names or gives them different names *)
+      "A0" :: "A1" :: "A2" :: "A3" :: "A4" :: "A5" :: "A6" :: "A7" :: []
+  | _ -> failwith "No calling convention implemented for the given architecture"
+
+let is_parameter_register (var: Var.t) (project: Project.t) : Bool.t =
+  let param_register = get_parameter_register_list project in
+  Option.is_some (List.find param_register ~f:(String.equal (Var.name var)))
+
+(** Return all registers that may contain return values of function calls *) (* TODO: Add Floating Point register! *)
+let get_return_register_list (project: Project.t) : String.t List.t =
+  let architecture = Project.arch project in
+  match architecture with
+  | `x86 ->
+      "EAX" :: []
+  | `x86_64 -> (* System V ABI *)
+      "RAX" :: "RDX" :: []
+  | `armv4 | `armv5 | `armv6 | `armv7
+  | `armv4eb | `armv5eb | `armv6eb | `armv7eb
+  | `thumbv4 | `thumbv5 | `thumbv6 | `thumbv7
+  | `thumbv4eb | `thumbv5eb | `thumbv6eb | `thumbv7eb ->
+      "R0" :: "R1" :: "R2" :: "R3" :: []
+  | `aarch64 | `aarch64_be -> (* ARM 64bit *)
+      "X0" :: "X1" :: "X2" :: "X3" :: "X4" :: "X5" :: "X6" :: "X7" :: []
+  | `ppc (* 32bit PowerPC *) (* TODO: add floating point register! *)
+  | `ppc64 | `ppc64le -> (* 64bit PowerPC *)
+      "R3" :: "R4" :: []
+  | `mips | `mips64 | `mips64el | `mipsel ->
+      "V0" :: "V1" :: []
+  | _ -> failwith "No calling convention implemented for the given architecture"
+
+let is_return_register (var: Var.t) (project: Project.t) : Bool.t =
+  let ret_register = get_return_register_list project in
+  Option.is_some (List.find ret_register ~f:(String.equal (Var.name var)))
+
 (** Parse a line from the dyn-syms output table of readelf. Return the name of a symbol if the symbol is an extern function name. *)
 let parse_dyn_sym_line line =
   let line = ref (String.strip line) in

src/utils/cconv.mli

 open Bap.Std
 open Core_kernel
 
+
 (** Returns whether a variable is callee saved according to the calling convention
     of the target architecture. Should only used for calls to functions outside
     of the program, not for calls between functions inside the program. *)
-val is_callee_saved: Var.t -> Project.t -> bool
+val is_callee_saved: Var.t -> Project.t -> Bool.t
+
+
+(** Returns whether a variable may be used to pass parameters to a function.
+    This depends on the calling convention of the target architecture and should only be used for extern function calls. *)
+val is_parameter_register: Var.t -> Project.t -> Bool.t
+
+
+(** Returns whether a variable may be used for return values of function calls.
+    This depends on the calling convention of the target architecture and should only be used for extern function calls. *)
+val is_return_register: Var.t -> Project.t -> Bool.t
 
 
 (** Returns a list of those function names that are extern symbols.

test/acceptance/test_check_path.py

@@ -16,4 +16,4 @@ class TestCheckPath(unittest.TestCase):
         output = subprocess.check_output(self.cmd.split())
         j = json.loads(output)
         self.assertTrue('check_path' in j)
-        self.assertEqual(len(j['check_path']), 7)
+        self.assertEqual(len(j['check_path']), 5)

test/acceptance/test_cwe476.py

@@ -14,34 +14,31 @@ class TestCwe476(unittest.TestCase):
             self.target, self.target, 'x64', 'gcc', self.string)
         self.assertEqual(res, expect_res)
 
-    @unittest.skip('FIXME!')
     def test_cwe476_01_x64_clang(self):
         expect_res = 1
         res = cwe_checker_testlib.execute_and_check_occurence(
             self.target, self.target, 'x64', 'clang', self.string)
         self.assertEqual(res, expect_res)
 
+    @unittest.skip('FIXME!')
     def test_cwe476_01_x86_gcc(self):
         expect_res = 1
         res = cwe_checker_testlib.execute_and_check_occurence(
             self.target, self.target, 'x86', 'gcc', self.string)
         self.assertEqual(res, expect_res)
 
-    @unittest.skip('FIXME!')
     def test_cwe476_01_x86_clang(self):
         expect_res = 1
         res = cwe_checker_testlib.execute_and_check_occurence(
             self.target, self.target, 'x86', 'clang', self.string)
         self.assertEqual(res, expect_res)
 
-    @unittest.skip('FIXME!')
     def test_cwe476_01_arm_gcc(self):
         expect_res = 1
         res = cwe_checker_testlib.execute_and_check_occurence(
             self.target, self.target, 'arm', 'gcc', self.string)
         self.assertEqual(res, expect_res)
 
-    @unittest.skip('FIXME!')
     def test_cwe476_01_arm_clang(self):
         expect_res = 1
         res = cwe_checker_testlib.execute_and_check_occurence(
@@ -118,7 +115,6 @@ class TestCwe476(unittest.TestCase):
             self.target, self.target, 'mips64el', 'clang', self.string)
         self.assertEqual(res, expect_res)
 
-    @unittest.skip("Fix issue in CWE476 implementation to support PPC")
     def test_cwe476_01_ppc_gcc(self):
         expect_res = 1
         res = cwe_checker_testlib.execute_and_check_occurence(

test/artificial_samples/SConstruct

@@ -78,7 +78,7 @@ def which(pgm):
 
 
 def optimize(filename):
-    optimize_me = ['cwe_476.c']
+    optimize_me = []
     if filename in optimize_me:
         return ' -O3'
     else:
@@ -140,7 +140,7 @@ def build_cpp(arch, compiler):
             env['CXXFLAGS'] = cpp_flags[arch] + optimize(str(prog))
             if arch == 'x86':
                 env['LINKFLAGS'] = '-m32'
- 
+
             compiler_abrev = get_compiler_abrev(compiler)
             if compiler_abrev == 'mingw32-g++' and str(prog) in skip_for_pe:
                     continue

test/unit/analysis/mem_region_test.ml

@@ -87,6 +87,22 @@ let test_around_zero () =
   let x = Mem_region.add x "Two" ~pos:(bv (-5)) ~size:(bv 20) in
   check "around_zero2" (Some(Error()) = Mem_region.get x (bv 0))
 
+let test_list_data () =
+  let bv num = Bitvector.of_int num ~width:32 in
+  let x = Mem_region.empty () in
+  let x = Mem_region.add x "One" ~pos:(bv (15)) ~size:(bv 10) in
+  let x = Mem_region.add x "Two" ~pos:(bv 0) ~size:(bv 10) in
+  let x = Mem_region.add x "Three" ~pos:(bv (-5)) ~size:(bv 10) in
+  let x = Mem_region.add x "Four" ~pos:(bv (-15)) ~size:(bv 10) in
+  let data_list = "Four" :: "Three" :: "One" :: [] in
+  check "list_data" (Mem_region.list_data x = data_list);
+  let data_pos_list = (bv (-15), "Four") :: (bv (-5), "Three") :: (bv 15, "One") ::[] in
+  check "list_data_pos" (Mem_region.list_data_pos x = data_pos_list);
+  let pos_minus = Bitvector.to_int_exn (Bitvector.signed (bv (-5))) in
+  let pos_plus = Bitvector.to_int_exn (Bitvector.signed (bv 5)) in
+  check "expected_sign_minus" (pos_minus < 0);
+  check "expected_sign_plus" (pos_plus > 0)
+
 let tests = [
   "Add", `Quick, test_add;
   "Negative Indices", `Quick, test_minus;
@@ -95,4 +111,5 @@ let tests = [
   "Merge", `Quick, test_merge;
   "Equal", `Quick, test_equal;
   "Around Zero", `Quick, test_around_zero;
+  "List data", `Quick, test_list_data;
 ]

test/unit/checkers/cwe_476_test.ml

+open Bap.Std
+open Core_kernel
+open Cwe_checker_core
+open Cwe_checker_core.Cwe_476.Private
+
+let check msg x = Alcotest.(check bool) msg true x
+
+let example_project : Project.t Option.t ref = ref None
+
+let call_handling_test () =
+  let project = Option.value_exn !example_project in
+  let state = State.empty in
+  let mock_tid = Tid.create () in
+  let mock_taint = Taint.add Taint.empty mock_tid in
+  let mock_hits = ref Taint.empty in
+  let rax_register = Var.create "RAX" (Bil.Imm (Symbol_utils.arch_pointer_size_in_bytes project * 8)) in
+  let rbx_register = Var.create "RBX" (Bil.Imm (Symbol_utils.arch_pointer_size_in_bytes project * 8)) in
+  let rdx_register = Var.create "RDX" (Bil.Imm (Symbol_utils.arch_pointer_size_in_bytes project * 8)) in
+
+  let state = State.set_register state rax_register mock_taint in
+  let _state = flag_unchecked_return_values state ~cwe_hits:mock_hits ~project in
+  check "flag_RAX_return" (false = Taint.is_empty !mock_hits);
+  let state = State.empty in
+  let state = State.set_register state rbx_register mock_taint in
+  mock_hits := Taint.empty;
+  let _state = flag_unchecked_return_values state ~cwe_hits:mock_hits ~project in
+  check "dont_flag_RBX_return" (Taint.is_empty !mock_hits);
+
+  let state = State.empty in
+  mock_hits := Taint.empty;
+  let state = State.set_register state rbx_register mock_taint in
+  let _state = flag_register_taints state ~cwe_hits:mock_hits in
+  check "flag_all_registers" (false = Taint.is_empty !mock_hits);
+
+  let state = State.empty in
+  mock_hits := Taint.empty;
+  let other_mock_taint = Taint.add Taint.empty (Tid.create ()) in
+  let state = State.set_register state rdx_register mock_taint in
+  let state = State.set_register state rbx_register other_mock_taint in
+  let state = flag_parameter_register state ~cwe_hits:mock_hits ~project in
+  check "flag_RDX_parameter" (false = Taint.is_empty !mock_hits && Option.is_none (State.find_register state rdx_register));
+  check "dont_flag_RBX_parameter" (Option.is_some (State.find_register state rbx_register));
+
+  let state = State.empty in
+  mock_hits := Taint.empty;
+  let state = State.set_register state rax_register mock_taint in
+  let state = State.set_register state rbx_register other_mock_taint in
+  let state = untaint_non_callee_saved_register state ~project in
+  check "RAX_non_callee_saved" (Option.is_none (State.find_register state rax_register));
+  check "RBX_callee_saved" (Option.is_some (State.find_register state rbx_register));
+  ()
+
+let state_test () =
+  let project = Option.value_exn !example_project in
+  let state = State.empty in
+  let mock_tid = Tid.create () in
+  let mock_taint = Taint.add Taint.empty mock_tid in
+  let _mock_hits = ref Taint.empty in
+  let rax_register = Var.create "RAX" (Bil.Imm (Symbol_utils.arch_pointer_size_in_bytes project * 8)) in
+  let rbx_register = Var.create "RBX" (Bil.Imm (Symbol_utils.arch_pointer_size_in_bytes project * 8)) in
+  let rdx_register = Var.create "RDX" (Bil.Imm (Symbol_utils.arch_pointer_size_in_bytes project * 8)) in
+  let state1 = State.set_register state rax_register mock_taint in
+  let state2 = State.set_register state rbx_register mock_taint in
+  let union_state = State.union state1 state2 in
+  check "state_union_RAX" (Option.is_some (State.find_register union_state rax_register));
+  check "state_union_RBX" (Option.is_some (State.find_register union_state rbx_register));
+  check "state_union_not_RDX" (Option.is_none (State.find_register union_state rdx_register));
+  ()
+
+(* TODO: write checks for expression handling!! *)
+
+let tests = [
+  "Call Handling", `Quick, call_handling_test;
+  "State Operations", `Quick, state_test;
+]

test/unit/checkers/cwe_476_test.mli

+open Bap.Std
+open Core_kernel
+
+
+val example_project: Project.t option ref
+
+val tests: unit Alcotest.test_case list

test/unit/cwe_checker_unit_tests.ml

@@ -5,10 +5,12 @@ open Core_kernel
 let run_tests project =
   Type_inference_test.example_project := Some(project);
   Cconv_test.example_project := Some(project);
+  Cwe_476_test.example_project := Some(project);
   Alcotest.run "Unit tests" ~argv:[|"DoNotComplainWhenRunAsABapPlugin";"--color=always";|] [
     "Mem_region_tests", Mem_region_test.tests;
     "Type_inference_tests", Type_inference_test.tests;
     "Cconv_tests", Cconv_test.tests;
+    "CWE_476_tests", Cwe_476_test.tests;
     "CWE_560_tests", Cwe_560_test.tests;
      ]
 

test/unit/utils/cconv_test.ml

@@ -17,6 +17,24 @@ let test_callee_saved () =
   let () = check "caller_saved_register" (is_callee_saved register project = false) in
   ()
 
+let test_parameter_register () =
+  (* this test assumes, that the example project is a x64 binary *)
+  let project = Option.value_exn !example_project in
+  let register = Var.create "RDX" (Bil.Imm (Symbol_utils.arch_pointer_size_in_bytes project * 8)) in
+  let () = check "return_register" (is_parameter_register register project) in
+  let register = Var.create "RAX" (Bil.Imm (Symbol_utils.arch_pointer_size_in_bytes project * 8)) in
+  let () = check "no_return_register" (is_parameter_register register project = false) in
+  ()
+
+let test_return_register () =
+  (* this test assumes, that the example project is a x64 binary *)
+  let project = Option.value_exn !example_project in
+  let register = Var.create "RAX" (Bil.Imm (Symbol_utils.arch_pointer_size_in_bytes project * 8)) in
+  let () = check "return_register" (is_return_register register project) in
+  let register = Var.create "R12" (Bil.Imm (Symbol_utils.arch_pointer_size_in_bytes project * 8)) in
+  let () = check "no_return_register" (is_return_register register project = false) in
+  ()
+
 let test_parse_dyn_syms () =
 (* this test assumes, that the example project is the arrays_x64.out binary from the artificial samples. *)
   let project = Option.value_exn !example_project in
@@ -29,5 +47,7 @@ let test_parse_dyn_syms () =
 
 let tests = [
   "Callee saved register", `Quick, test_callee_saved;
+  "Parameter register", `Quick, test_parameter_register;
+  "Return register", `Quick, test_return_register;
   "Parse dynamic symbols", `Quick, test_parse_dyn_syms;
 ]