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.ml

@@ -3,166 +3,306 @@ open Bap.Std
 open Log_utils
 
 let name = "CWE476"
-let version = "0.2"
+let version = "0.3"
 
-(* Access type denotes whether a variable var gets accessed or the memory stored at
-   address var gets accessed *)
-type access_type = | Access of Bil.var | MemAccess of Bil.var | NoAccess
+(* TODO: This check is based on Mem_region, which does not support partial access yet.
+   Thus partially written tainted values may be marked as error and thus the taint is falsely forgotten. *)
 
-(* The union of two accesses is the higher access with MemAcces > Access > NoAccess *)
-let union_access access1 access2 : access_type =
-  match (access1, access2) with
-  | (MemAccess(_), _) -> access1
-  | (_, MemAccess(_))
-  | (_, Access(_))    -> access2
-  | _                 -> access1
 
+(** Each taint is denoted by the Tid of the basic block where it originated from.
+    Each value can be tainted by different sources at the same time. *)
+module Taint = Tid.Set
 
-(* union of three accesses for convenience *)
-let union_access_triple access1 access2 access3 =
-  union_access access1 access2
-  |> union_access access3
-
-(* the state contains a list of pairs of register names containing an unchecked
-   return value and the term identifiers of the block where the unchecked
-   return value was generated. *)
+(** The state contains taint information for all registers and stack variables. *)
 module State = struct
-  type t = (Var.t * Tid.t) list
+  type t = {
+    register: Taint.t Var.Map.t;
+    stack: Taint.t Mem_region.t;
+  } [@@deriving bin_io, compare, sexp]
+
+  (** Get an empty state without tainted values. *)
+  let empty : t =
+    { register = Var.Map.empty;
+      stack = Mem_region.empty () }
 
-  (** adds var as a tainted register (with the taint source given by tid) *)
-  let add state var tid =
-    let state = List.Assoc.remove state ~equal:Var.(=) var in
-    List.Assoc.add state ~equal:Var.(=) var tid
+  (** equality function for states *)
+  let equal (state1: t) (state2: t) : Bool.t =
+    let reg_equal = Var.Map.equal Taint.equal state1.register state2.register in
+    let stack_equal = Mem_region.equal state1.stack state2.stack ~data_equal:Taint.equal in
+    reg_equal && stack_equal
 
-  (** returns Some(tid) if var is a tainted register, None otherwise *)
-  let find state var =
-    List.Assoc.find state ~equal:Var.(=) var
+  (** set the taint of a register *)
+  let set_register (state: t) (register: Var.t) (taint: Taint.t) : t =
+    { state with register = Var.Map.set state.register ~key:register ~data: taint}
 
-  (** returns the tid associated with a tainted register *)
-  let find_exn state var =
-    Option.value_exn (find state var)
+  (** return the taint of a register *)
+  let find_register (state: t) (register: Var.t) : Taint.t Option.t =
+    Var.Map.find state.register register
 
   (** only remove the register var from the list of tainted registers *)
-  let remove_var state var =
-    List.Assoc.remove state ~equal:Var.(=) var
-
-  (** filters out all registers from the state with the same tid *)
-  let remove_tid state var =
-    let tid = find_exn state var in
-    List.filter state ~f:(fun (_, state_elem_tid) -> not (state_elem_tid = tid))
-
-  (** two states are equal if they contain the same set of tainted registers*)
-  let equal state1 state2 =
-    (List.length state1) = (List.length state2) &&
-    not (List.exists state1 ~f:(fun (var, _tid) -> Option.is_none (find state2 var) ))
-
-  (** The union of two states is the union of the tainted registers*)
-  let union state1 state2 =
-    List.fold state2 ~init:state1 ~f:(fun state (var, tid) ->
-        if Option.is_some (find state var) then
-          state
-        else
-          (var, tid) :: state
-      )
-
-  (** remove virtual registers from the state (useful at the end of a block) *)
-  let remove_virtual_registers state =
-    List.filter state ~f:(fun (var, _tid) -> Var.is_physical var)
+  let remove_register (state: t) (register: Var.t) : t =
+    { state with register = Var.Map.remove state.register register }
+
+  (** set the taint of a stack element *)
+  let set_stack (state: t) ~(pos: Bitvector.t) ~(size: Bitvector.t) (taint: Taint.t) : t =
+    { state with stack = Mem_region.add state.stack taint ~pos ~size }
+
+  (** get the taint from the stack
+      TODO: Mem_region is currently unsound for only partially loaded values, which might lead to errors here. *)
+  let find_stack (state: t) ~(pos: Bitvector.t) : Taint.t Option.t =
+    match Mem_region.get state.stack pos with
+    | Some(Ok(taint, _size)) -> Some(taint)
+    | _ -> None
+
+  (** remove a stack element *)
+  let remove_stack (state: t) ~(pos: Bitvector.t) ~(size: Bitvector.t) : t =
+    { state with stack = Mem_region.remove state.stack ~pos ~size}
+
+  (** remove all Tids contained in the taint from all taints in the state *)
+  let remove_taint (state: t) (taint_to_remove: Taint.t) : t =
+    let register_list = Var.Map.to_alist state.register in
+    let cleaned_register = List.fold register_list ~init:Var.Map.empty ~f:(fun cleaned_register (register, taint) ->
+      let cleaned_taint = Tid.Set.diff taint taint_to_remove in
+      if Tid.Set.is_empty cleaned_taint then
+        cleaned_register
+      else
+        Var.Map.set cleaned_register ~key:register ~data:cleaned_taint
+    ) in
+    let cleaned_stack = Mem_region.map_data state.stack ~f:(fun taint ->
+      Tid.Set.diff taint taint_to_remove
+    ) in
+    { register = cleaned_register;
+      stack = cleaned_stack; }
+
+  (** The union of two states is the union of all taints *)
+  let union (state1: t) (state2: t) : t =
+    let register = Var.Map.merge state1.register state2.register ~f:(fun ~key:_ values->
+      match values with
+      | `Both (taint1, taint2) -> Some (Taint.union taint1 taint2)
+      | `Left taint | `Right taint -> Some taint
+    ) in
+    let stack = Mem_region.merge state1.stack state2.stack ~data_merge:(fun taint1 taint2 ->
+      Some( Ok(Taint.union taint1 taint2) )
+    ) in
+    { register = register;
+      stack = stack; }
+
+  (** remove virtual register from the state (useful at the end of a block) *)
+  let remove_virtual_register (state: t) : t =
+    { state with register = Var.Map.filter_keys state.register ~f:(fun var -> Var.is_physical var) }
+
+end
+
 
+(** The stack info contains all necessary information to access stack variables. *)
+module StackInfo = struct
+  type t = {
+    type_info: Type_inference.TypeInfo.t;
+    sub_tid:  Tid.t;
+    project: Project.t;
+    strict_mem_policy: Bool.t;
+  }
+
+  (** If the expression denotes an address on the stack, return the address. *)
+  let get_address (stack_info: t) (expression: Exp.t) : Bitvector.t Option.t =
+    Type_inference.TypeInfo.compute_stack_offset stack_info.type_info expression ~sub_tid:stack_info.sub_tid ~project:stack_info.project
+
+  (** Assemble a StackInfo.t object. *)
+  let assemble (pointer_info_map: Type_inference.TypeInfo.t Tid.Map.t) (term_tid: Tid.t) ~(sub_tid: Tid.t) ~(project: Project.t) ~(strict_mem_policy: Bool.t) : t =
+    { type_info = Tid.Map.find_exn pointer_info_map term_tid;
+      sub_tid = sub_tid;
+      project = project;
+      strict_mem_policy = strict_mem_policy; }
+
+  (**/**)
+  (* assemble a mock StackInfo for unit tests *)
+  let assemble_mock_info (mock_tid: Tid.t) (project: Project.t) : t =
+    { type_info = { Type_inference.TypeInfo.stack = Mem_region.empty (); Type_inference.TypeInfo.reg = Var.Map.empty};
+      sub_tid = mock_tid;
+      project = project;
+      strict_mem_policy = false; }
+  (**/**)
 end
 
-(* check whether an expression contains an unchecked value. *)
-let rec contains_unchecked exp state : access_type =
+
+(** append taint to the list of already found cwe_hits *)
+let append_to_hits (cwe_hits:Taint.t ref) (taint: Taint.t) : unit =
+  cwe_hits := Taint.union !cwe_hits taint
+
+
+(** Check whether an expression contains a tainted value.
+    Memory accesses through tainted values are added to cwe_hits, but the Tids are not removed from the state. *)
+let rec contains_taint (exp: Exp.t) (state: State.t) ~(cwe_hits: Taint.t ref) ~(stack: StackInfo.t) : Taint.t =
   match exp with
-  | Bil.Load(_mem, addr, _, _)->
+  | Bil.Load(_mem, addr, _endian, _size)->
     begin
-      let acc = contains_unchecked addr state in
-      match acc with
-      | MemAccess(_) -> acc
-      | Access(var) -> MemAccess(var)
-      | NoAccess -> NoAccess
+      let access_taint = contains_taint addr state ~cwe_hits ~stack in
+      let () = if Taint.is_empty access_taint = false then append_to_hits cwe_hits access_taint in
+      match StackInfo.get_address stack addr with
+      | Some(stack_offset) -> Option.value (State.find_stack state ~pos:stack_offset) ~default:Taint.empty
+      | None -> Taint.empty
     end
   | Bil.Store(_mem, addr, val_expression, _,_) ->
     begin
-      let acc = union_access (contains_unchecked addr state) (contains_unchecked val_expression state) in
-      match acc with
-      | MemAccess(_) -> acc
-      | Access(var) -> MemAccess(var)
-      | NoAccess -> NoAccess
-    end
-  | Bil.BinOp(_, exp1, exp2) -> union_access (contains_unchecked exp1 state) (contains_unchecked exp2 state)
-  | Bil.UnOp(_, exp) -> contains_unchecked exp state
-  | Bil.Var(var) ->
-    begin
-      match State.find state var with
-      | Some(_) -> Access(var)
-      | None -> NoAccess
+      let access_taint = contains_taint addr state ~cwe_hits ~stack in
+      let value_taint = contains_taint val_expression state ~cwe_hits ~stack in
+      let () = if Taint.is_empty access_taint = false then append_to_hits cwe_hits access_taint in
+      match StackInfo.get_address stack addr with
+      | Some(_) -> Taint.empty
+      | None ->
+          let () = if stack.strict_mem_policy && (Taint.is_empty value_taint = false) then append_to_hits cwe_hits value_taint in
+          Taint.empty
     end
-  | Bil.Int(_) -> NoAccess
-  | Bil.Cast(_, _, exp) -> contains_unchecked exp state
+  | Bil.BinOp(Bil.XOR, Bil.Var(var1), Bil.Var(var2)) when var1 = var2 -> Taint.empty (* standard assembly shortcut for setting a register to NULL *)
+  | Bil.BinOp(_, exp1, exp2) -> Taint.union (contains_taint exp1 state ~cwe_hits ~stack) (contains_taint exp2 state ~cwe_hits ~stack)
+  | Bil.UnOp(_, exp) -> contains_taint exp state ~cwe_hits ~stack
+  | Bil.Var(var) -> Option.value (State.find_register state var) ~default:Taint.empty
+  | Bil.Int(_) -> Taint.empty
+  | Bil.Cast(_, _, exp) -> contains_taint exp state ~cwe_hits ~stack
   | Bil.Let(var, exp1, exp2) ->
-    union_access_triple (contains_unchecked exp1 state) (contains_unchecked exp2 state) (contains_unchecked (Bil.var var) state)
-  | Bil.Unknown(_) -> NoAccess
+      Taint.union_list (
+        (contains_taint exp1 state ~cwe_hits ~stack)
+        :: (contains_taint exp2 state ~cwe_hits ~stack)
+        :: (contains_taint (Bil.var var) state ~cwe_hits ~stack) :: [])
+  | Bil.Unknown(_) -> Taint.empty
   | Bil.Ite(if_, then_, else_) ->
-    union_access_triple (contains_unchecked if_ state) (contains_unchecked then_ state) (contains_unchecked else_ state)
-  | Bil.Extract(_,_, exp) -> contains_unchecked exp state
-  | Bil.Concat(exp1, exp2) -> union_access (contains_unchecked exp1 state) (contains_unchecked exp2 state)
+      Taint.union_list (
+        (contains_taint if_ state ~cwe_hits ~stack)
+        :: (contains_taint then_ state ~cwe_hits ~stack)
+        :: (contains_taint else_ state ~cwe_hits ~stack) :: [])
+  | Bil.Extract(_,_, exp) -> contains_taint exp state ~cwe_hits ~stack
+  | Bil.Concat(exp1, exp2) -> Taint.union (contains_taint exp1 state ~cwe_hits ~stack) (contains_taint exp2 state ~cwe_hits ~stack)
 
-(* If an formerly unchecked return value was checked then remove all registers pointing
-   to the source of this return value from state. *)
-let checks_value exp state : State.t =
+
+(** Parse an expression for memory accesses through tainted values and taint contained in the value itself.
+    All memory accesses except for loading/storing values from/to the stack get flagged as cwe_hits.
+    Returns the taint of the expression and the new state, with the Tids of new cwe_hits removed from both. *)
+let parse_taint_of_exp (exp: Exp.t) (state: State.t) ~(cwe_hits: Taint.t ref) ~(stack: StackInfo.t) : Taint.t * State.t =
+  let hits_to_clean : Taint.t ref = ref Taint.empty in
+  let unchecked_taint = contains_taint exp state ~cwe_hits:hits_to_clean ~stack in
+  let () = append_to_hits cwe_hits !hits_to_clean in
+  let state = State.remove_taint state !hits_to_clean in
+  let unchecked_taint = Taint.diff unchecked_taint !hits_to_clean in
+  (unchecked_taint, state)
+
+
+(** If an formerly unchecked return value was checked then remove all registers pointing
+    to the source of this return value from state. *)
+let checks_value (exp: Exp.t) (state: State.t) ~(cwe_hits: Taint.t ref) ~(stack: StackInfo.t) : State.t =
   match exp with
   | Bil.Ite(if_, _then_, _else_) -> begin
-      match contains_unchecked if_ state with
-      | Access(var) ->
-        (* We filter out all registers with the same generating tid, since we have checked
-           the return value of this source *)
-        State.remove_tid state var
-      | MemAccess(_) (* This is a memory access before checking the return value, so do nothing here. *)
-      | NoAccess -> state
+      let (taint_to_remove, state) = parse_taint_of_exp if_ state ~cwe_hits ~stack in
+      if Taint.is_empty taint_to_remove = false then
+        State.remove_taint state taint_to_remove
+      else
+        state
     end
   | _ -> state
 
-let append_to_hits (cwe_hits:Tid.t list ref) (tid:Tid.t) =
-  match List.find cwe_hits.contents ~f:(fun elem -> elem = tid) with
-  | Some(_) -> ()
-  | None -> (cwe_hits := (tid :: cwe_hits.contents))
 
 (** flags any access (not just memory access) from an unchecked source as a cwe_hit. *)
-let flag_any_access exp state ~cwe_hits =
-  match contains_unchecked exp state with
-  | MemAccess(var) | Access(var) ->
-    let tid = State.find_exn state var in
-    append_to_hits cwe_hits tid;
-    State.remove_tid state var
-  | NoAccess -> state
-
-(** flag all unchecked registers as cwe_hits, return empty state *)
-let flag_all_unchecked_registers state ~cwe_hits =
-  let () = List.iter state ~f:(fun (_var, tid) ->
-      append_to_hits cwe_hits tid) in
-  []
+let flag_any_access (exp: Exp.t) (state: State.t) ~(cwe_hits: Taint.t ref) ~(stack: StackInfo.t) : State.t=
+  let (taint_to_flag, state) = parse_taint_of_exp exp state ~cwe_hits ~stack in
+  let () = append_to_hits cwe_hits taint_to_flag in
+  State.remove_taint state taint_to_flag
+
+
+(** flag all unchecked registers and stack variables that may be used as return values.
+    That means stack variables above the return pointer get flagged,
+    but variables below the return pointer are treated as local variables and do not get flagged.
+    Return empty state *)
+let flag_unchecked_return_values (state: State.t) ~(cwe_hits: Taint.t ref) ~(project: Project.t) : State.t =
+  let taint_to_flag = Var.Map.fold state.register ~init:Taint.empty ~f:(fun ~key ~data taint_accum ->
+    if Cconv.is_return_register key project then
+      Taint.union taint_accum data
+    else
+      taint_accum
+  ) in
+  let taint_to_flag = List.fold (Mem_region.list_data_pos state.stack) ~init:taint_to_flag ~f:(fun taint_accum (position_unsigned, taint_value) ->
+    let position = Bitvector.to_int_exn (Bitvector.signed position_unsigned) in
+    if position >= 0 then
+      Taint.union taint_accum taint_value
+    else
+      taint_accum
+  ) in
+  let () = append_to_hits cwe_hits taint_to_flag in
+  State.empty
+
+
+(** flag all register taints as cwe_hits, but not taints that are only contained in stack variables *)
+let flag_register_taints (state: State.t) ~(cwe_hits: Taint.t ref) : State.t =
+  let taint_to_flag = List.fold (Var.Map.data state.register) ~init: Taint.empty ~f:(fun taint_accum register_taint ->
+    Taint.union taint_accum register_taint
+  ) in
+  let () = append_to_hits cwe_hits taint_to_flag in
+  State.remove_taint state taint_to_flag
+
+
+(** Flag all possible parameter register as cwe_hits. These registers may be input values to an extern function call.
+    This can lead to false positives if a function does not use all of these registers for argument passing. *)
+let flag_parameter_register (state: State.t) ~(cwe_hits: Taint.t ref) ~(project: Project.t) : State.t =
+  let taint_to_flag = Var.Map.fold state.register ~init:Taint.empty ~f:(fun ~key ~data taint_accum ->
+    if Cconv.is_parameter_register key project then
+      Taint.union taint_accum data
+    else
+      taint_accum
+  ) in
+  let () = append_to_hits cwe_hits taint_to_flag in
+  State.remove_taint state taint_to_flag
+
+
+(** Remove the taint of non-callee-saved register (without flagging them).
+    For taints in parameter register we assume that they are checked by the callee, thus we also remove the corresponding Tids from the state. *)
+let untaint_non_callee_saved_register (state: State.t) ~(project: Project.t) : State.t =
+  let taint_to_remove = Var.Map.fold state.register ~init:Taint.empty ~f:(fun ~key ~data taint_accum ->
+    if Cconv.is_callee_saved key project then
+      taint_accum
+    else
+      Taint.union taint_accum data
+  ) in
+  let state = State.remove_taint state taint_to_remove in
+  Var.Map.fold state.register ~init:state ~f:(fun ~key ~data:_ state ->
+    if Cconv.is_callee_saved key project then
+      state
+    else
+      State.remove_register state key
+  )
+
+
+(** If the expression is a store onto a stack variable, write the corresponding taint to the stack. *)
+let update_stack_on_stores (exp: Exp.t) (state: State.t) ~(stack: StackInfo.t) : State.t =
+  let pointer_size = Symbol_utils.arch_pointer_size_in_bytes stack.project in
+  match exp with
+  | Bil.Store(_mem, address_exp, value, _endian, size) -> begin
+      let value_taint = contains_taint value state ~cwe_hits:(ref Taint.empty) ~stack in
+      match StackInfo.get_address stack address_exp with
+      | Some(address) ->
+          if Taint.is_empty value_taint then
+            State.remove_stack state ~pos:address ~size:(Bitvector.of_int (Size.in_bytes size) ~width:pointer_size)
+          else
+            State.set_stack state value_taint ~pos:address ~size:(Bitvector.of_int (Size.in_bytes size) ~width:pointer_size)
+      | None -> state
+    end
+  | _ -> state
+
 
 (** Updates the state depending on the def. If memory is accessed using an unchecked return value,
     then the access is added to the list of cwe_hits. *)
-let update_state_def def state ~cwe_hits =
+let update_state_def (def: Def.t) (state: State.t) ~(cwe_hits: Taint.t ref) ~(stack: StackInfo.t) : State.t =
   let (lhs, rhs) = (Def.lhs def, Def.rhs def) in
-  let state = checks_value rhs state in
-  match contains_unchecked rhs state with
-  | MemAccess(var) -> begin (* we found a case of unchecked return value *)
-      let tid = State.find_exn state var in
-      append_to_hits cwe_hits tid;
-      State.remove_tid state var
-    end
-  | Access(var) -> (* taint the lhs as an unchecked return value *)
-    let tid = State.find_exn state var in
-    State.add state lhs tid
-  | NoAccess -> (* no access to an unchecked return value in rhs. Since lhs is overwritten, it cannot be an unchecked return value anymore. *)
-    State.remove_var state lhs
+  let state = checks_value rhs state ~cwe_hits ~stack in
+  let (rhs_taint, state) = parse_taint_of_exp rhs state ~cwe_hits ~stack in
+  let state =
+    if Taint.is_empty rhs_taint then
+      State.remove_register state lhs
+    else
+      State.set_register state lhs rhs_taint in
+  update_stack_on_stores rhs state ~stack
+
 
 (** Taint the return registers of a function as unchecked return values. *)
-let taint_return_registers func_tid state ~program ~block =
-  let func = Term.find_exn sub_t program func_tid in
+let taint_return_registers (func_tid: Tid.t) (state: State.t) ~(project: Project.t) ~(block: Blk.t) : State.t =
+  let func = Term.find_exn sub_t (Project.program project) func_tid in
   let arguments = Term.enum arg_t func in
   (* Every return register is tainted as unchecked return value. *)
   Seq.fold arguments ~init:state ~f:(fun state arg ->
@@ -172,126 +312,186 @@ let taint_return_registers func_tid state ~program ~block =
         let variable = match Bap.Std.Arg.rhs arg with
           | Bil.Var(var) -> var
           | _ -> failwith "[CWE476] Return register wasn't a register." in
-        State.add state variable (Term.tid block)
+        State.set_register state variable (Taint.add Taint.empty (Term.tid block))
     )
 
 (** Updates the state depending on the jump. On a jump to a function from the function list
     taint all return registers as unchecked return values. *)
-let update_state_jmp jmp state ~cwe_hits ~function_names ~program ~block ~strict_call_policy =
+let update_state_jmp
+      (jmp: Jmp.t)
+      (state: State.t)
+      ~(cwe_hits: Taint.t ref)
+      ~(malloc_like_functions: String.t List.t)
+      ~(extern_functions: String.Set.t)
+      ~(stack: StackInfo.t)
+      ~(block: Blk.t)
+      ~(strict_call_policy: Bool.t) : State.t =
   (* first check the guard condition for unchecked access. Any normal access clears the access from being unchecked *)
   let condition_exp = Jmp.cond jmp in
   let state = begin
-    match contains_unchecked condition_exp state with
-    | Access(var) ->
-      State.remove_tid state var
-    | MemAccess(var) -> (* a memory access using an unchecked value is still an error *)
-      let tid = State.find_exn state var in
-      let () = append_to_hits cwe_hits tid in
-      State.remove_tid state var
-    | NoAccess -> state
+    let (condition_taint, state) = parse_taint_of_exp condition_exp state ~cwe_hits ~stack in
+    if Taint.is_empty condition_taint then
+      state
+    else
+      State.remove_taint state condition_taint
   end in
   match Jmp.kind jmp with
-  | Goto(Indirect(exp)) -> flag_any_access exp state ~cwe_hits
+  | Goto(Indirect(exp)) -> flag_any_access exp state ~cwe_hits ~stack
   | Goto(Direct(_)) -> state
   | Ret(_) -> if strict_call_policy then
-      flag_all_unchecked_registers state ~cwe_hits
+      flag_unchecked_return_values state ~cwe_hits ~project:stack.project
     else
       state
-  | Int(_, _) -> flag_all_unchecked_registers state ~cwe_hits
+  | Int(_, _) -> flag_register_taints state ~cwe_hits
   | Call(call) ->
-    let state = match Call.return call with
-      | Some(Indirect(exp)) -> flag_any_access exp state ~cwe_hits
-      | _ -> state in
-    let state = match Call.target call with
-      | Indirect(exp) -> flag_any_access exp state ~cwe_hits
-      | _ -> state in
-    let state = match strict_call_policy with
-      | true -> (* all unchecked registers get flagged as hits *)
-        flag_all_unchecked_registers state ~cwe_hits
-      | false -> (* we assume that the callee will check all remaining unchecked values *)
-        [] in
-    match Call.target call with
-    | Indirect(_) -> state (* already handled above *)
-    | Direct(tid) ->
-      if List.exists function_names ~f:(fun elem -> String.(=) elem (Tid.name tid)) then
-        taint_return_registers tid state ~program ~block
-      else
-        state
+      (* flag tainted values in the call and return expressions of indirect calls *)
+      let state = match Call.return call with
+        | Some(Indirect(exp)) -> flag_any_access exp state ~cwe_hits ~stack
+        | _ -> state in
+      let state =  begin match Call.target call with
+        | Indirect(exp) -> flag_any_access exp state ~cwe_hits ~stack
+        | _ -> state end in
+      (* flag tainted values in the parameter registers (if strict_call_policy is set to true)*)
+      let state = match (Call.target call, strict_call_policy) with
+        | (Indirect(_), false)
+        | (Direct(_), false) -> state
+        | (Indirect(_), true) -> flag_parameter_register state ~cwe_hits ~project:stack.project (* TODO: indirect calls are handled as extern calls right now. Change that *)
+        | (Direct(tid), true) ->
+            let sub = Term.find_exn sub_t (Project.program stack.project) tid in
+            if Set.mem extern_functions (Sub.name sub) then
+              flag_parameter_register state ~cwe_hits ~project:stack.project
+            else (* flag all registers for intern calls, as these do not necessarily adhere to any calling convention *)
+              flag_register_taints state ~cwe_hits
+      in
+      (* remove the taint of non-callee-saved registers *)
+      let state = match Call.target call with
+        | Direct(tid) ->
+            let sub = Term.find_exn sub_t (Project.program stack.project) tid in
+            if Set.mem extern_functions (Sub.name sub) then
+              untaint_non_callee_saved_register state ~project:stack.project
+            else (* we untaint all registers for internal function calls, as these do not necessarily adhere to any calling convention *)
+              { state with register = Var.Map.empty }
+        | Indirect(_) -> (* we treat all indirect calls as extern function calls, since we cannot handle indirect calls properly yet *)
+            untaint_non_callee_saved_register state ~project:stack.project
+      in
+      (* introduce new taint for the return values of malloc_like_functions *)
+      match Call.target call with
+      | Indirect(_) -> state
+      | Direct(tid) ->
+          if List.exists malloc_like_functions ~f:(fun elem -> String.(=) elem (Tid.name tid)) then
+            taint_return_registers tid state ~project:stack.project ~block
+          else
+            state
+
 
 (** updates a block analysis.
     The strict call policy decides the behaviour on call and return instructions:
-    strict: all unchecked values get flagged as cwe-hits
-    non-strict: the state gets cleared, it is assumed that the target of the call/return
-    instruction checks all remaining unchecked values. *)
-let update_block_analysis block register_state ~cwe_hits ~function_names ~program ~strict_call_policy =
+    strict: unchecked values in registers get flagged as cwe_hits
+    non-strict: unchecked values in registers get marked as checked. It is assumed that the callee checks these values. *)
+let update_block_analysis
+      (block: Blk.t)
+      (state: State.t)
+      ~(cwe_hits: Taint.t ref)
+      ~(malloc_like_functions: String.t List.t)
+      ~(extern_functions: String.Set.t)
+      ~(sub_tid: Tid.t)
+      ~(project: Project.t)
+      ~(strict_call_policy: Bool.t)
+      ~(strict_mem_policy: Bool.t)  : State.t =
   let elements = Blk.elts block in
-  let register_state = Seq.fold elements ~init:register_state ~f:(fun state element ->
+  let type_info_map = Type_inference.get_type_info_of_block ~project block ~sub_tid in
+  let state = Seq.fold elements ~init:state ~f:(fun state element ->
       match element with
-      | `Def def -> update_state_def def state ~cwe_hits
+      | `Def def ->
+          let stack = StackInfo.assemble type_info_map (Term.tid def) ~sub_tid ~project ~strict_mem_policy in
+          update_state_def def state ~cwe_hits ~stack
       | `Phi _phi -> state (* We ignore phi terms for this analysis. *)
-      | `Jmp jmp -> update_state_jmp jmp state ~cwe_hits ~function_names ~program ~block ~strict_call_policy
+      | `Jmp jmp ->
+          let stack = StackInfo.assemble type_info_map (Term.tid jmp) ~sub_tid ~project ~strict_mem_policy in
+          update_state_jmp jmp state ~cwe_hits ~malloc_like_functions ~extern_functions ~stack ~block ~strict_call_policy
     ) in
-  State.remove_virtual_registers register_state (* virtual registers should not be accessed outside of the block where they are defined. *)
+  State.remove_virtual_register state (* virtual registers should not be accessed outside of the block where they are defined. *)
+
 
-let print_hit tid ~sub ~function_names ~tid_map =
+(** print a cwe_hit to the log *)
+let print_hit (tid: Tid.t) ~(sub: Sub.t) ~(malloc_like_functions: String.t List.t) ~(tid_map: Word.t Tid.Map.t) : unit =
   let block = Option.value_exn (Term.find blk_t sub tid) in
   let jmps = Term.enum jmp_t block in
   let _ = Seq.find_exn jmps ~f:(fun jmp ->
-      match Jmp.kind jmp with
-      | Call(call) -> begin
-          match Call.target call with
-          | Direct(call_tid) -> Option.is_some (List.find function_names ~f:(fun fn_name ->
-                                                    if fn_name = (Tid.name call_tid) then
-                                                      begin
-                                                      let address = Address_translation.translate_tid_to_assembler_address_string tid tid_map in
-                                                      let tids = [Address_translation.tid_to_string tid] in
-                                                      let description = sprintf
-                                                                          "(NULL Pointer Dereference) There is no check if the return value is NULL at %s (%s)."
-                                                                          address
-                                                                          fn_name in
-                                                      let cwe_warning = cwe_warning_factory
-                                                                          name
-                                                                          version
-                                                                          ~addresses:[address]
-                                                                          ~tids:tids
-                                                                          ~symbols:[fn_name]
-                                                                          description in
-                                                      collect_cwe_warning cwe_warning;
-                                                      true
-              end else
-                false
-            ))
-          | _ -> false
-        end
-      | _ -> false
-    ) in ()
-
-let check_cwe prog _proj tid_map symbol_names parameters =
+    match Jmp.kind jmp with
+    | Call(call) -> begin
+        match Call.target call with
+        | Direct(call_tid) -> Option.is_some (List.find malloc_like_functions ~f:(fun fn_name ->
+          if fn_name = (Tid.name call_tid) then
+            begin
+              let address = Address_translation.translate_tid_to_assembler_address_string tid tid_map in
+              let tids = [Address_translation.tid_to_string tid] in
+              let description = sprintf
+                                  "(NULL Pointer Dereference) There is no check if the return value is NULL at %s (%s)."
+                                  address
+                                  fn_name in
+              let cwe_warning = cwe_warning_factory
+                                  name
+                                  version
+                                  ~addresses:[address]
+                                  ~tids:tids
+                                  ~symbols:[fn_name]
+                                  description in
+              collect_cwe_warning cwe_warning;
+              true
+            end else
+            false
+        ))
+        | _ -> false
+      end
+    | _ -> false
+  ) in ()
+
+
+let check_cwe (_prog: Program.t) (project: Project.t) (tid_map: Word.t Tid.Map.t) (symbol_names: String.t List.t List.t) (parameters: String.t List.t) =
   let symbols = match symbol_names with
     | hd :: _ -> hd
     | _ -> failwith "[CWE476] symbol_names not as expected" in
-  let (strict_call_policy_string, max_steps_string) = match parameters with
-    | par1 :: par2 :: _ -> (par1, par2)
+  let (strict_call_policy_string, strict_mem_policy_string, max_steps_string) = match parameters with
+    | par1 :: par2 :: par3 :: _ -> (par1, par2, par3)
     | _ -> failwith "[CWE476] parameters not as expected" in
   let strict_call_policy = match String.split strict_call_policy_string ~on:'=' with
     | "strict_call_policy" :: policy :: [] -> bool_of_string policy
     | _ -> failwith "[CWE476] parameters not as expected" in
+  let strict_mem_policy = match String.split strict_mem_policy_string ~on:'=' with
+    | "strict_memory_policy" :: policy :: [] -> bool_of_string policy
+    | _ -> failwith "[CWE476] parameters not as expected" in
   let max_steps = match String.split max_steps_string ~on:'=' with
     | "max_steps" :: num :: [] -> int_of_string num
     | _ -> failwith "[CWE476] parameters not as expected" in
-  let function_names = List.map symbols ~f:(fun symb -> "@" ^ symb)  in
-  let subfunctions = Term.enum sub_t prog in
+  let malloc_like_functions = List.map symbols ~f:(fun symb -> "@" ^ symb)  in
+  let extern_functions = Cconv.parse_dyn_syms project in
+  (* run the pointer inference analysis. TODO: This should be done somewhere else as this analysis will be needed in more than one check! *)
+  let project = Type_inference.compute_pointer_register project in
+  let subfunctions = Term.enum sub_t (Project.program project) in
   Seq.iter subfunctions ~f:(fun subfn ->
       let cfg = Sub.to_cfg subfn in
-      let cwe_hits = ref [] in
+      let cwe_hits = ref Taint.empty in
       let empty = Map.empty (module Graphs.Ir.Node) in
-      let init = Graphlib.Std.Solution.create empty [] in
+      let init = Graphlib.Std.Solution.create empty State.empty in
       let equal = State.equal in
       let merge = State.union in
       let f = (fun node state ->
           let block = Graphs.Ir.Node.label node in
-          update_block_analysis block state ~cwe_hits ~function_names ~program:prog ~strict_call_policy
+          update_block_analysis block state ~cwe_hits ~malloc_like_functions ~extern_functions ~sub_tid:(Term.tid subfn) ~project ~strict_call_policy ~strict_mem_policy
         ) in
       let _ = Graphlib.Std.Graphlib.fixpoint (module Graphs.Ir) cfg ~steps:max_steps ~rev:false ~init:init ~equal:equal ~merge:merge ~f:f in
-      List.iter (!cwe_hits) ~f:(fun hit -> print_hit hit ~sub:subfn ~function_names ~tid_map)
-    )
+      Tid.Set.iter (!cwe_hits) ~f:(fun hit -> print_hit hit ~sub:subfn ~malloc_like_functions ~tid_map)
+  )
+
+(**/**)
+(* Functions made public for unit tests *)
+module Private = struct
+  module StackInfo = StackInfo
+  module Taint = Taint
+  module State = State
+  let flag_unchecked_return_values = flag_unchecked_return_values
+  let flag_register_taints = flag_register_taints
+  let flag_parameter_register = flag_parameter_register
+  let untaint_non_callee_saved_register = untaint_non_callee_saved_register
+end

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;
 ]