Added flag and map comparison.

Author: DieracDelta

wp/plugin/lib/analysis.ml

@@ -259,6 +259,21 @@ let single (bap_ctx : ctxt) (z3_ctx : Z3.context) (var_gen : Env.var_gen)
     Printf.printf "\nSub:\n%s\n%!" (Sub.to_string main_sub);
   { pre = pre; orig = env, main_sub; modif = env, main_sub }
 
+let check_syntax_equality (bap_ctx : ctxt)
+    (p : Params.t) (file1 : string) (file2 : string) : bool =
+  (* TODO replace with param flag *)
+  if p.syntax_equality then
+    let prog1, _ = Utils.read_program bap_ctx
+        ~loader:Utils.loader ~filepath:file1 in
+    let prog2, _ = Utils.read_program bap_ctx
+        ~loader:Utils.loader ~filepath:file2 in
+    let subs1 = Term.enum sub_t prog1 in
+    let subs2 = Term.enum sub_t prog2 in
+    let main_sub1 = Utils.find_func_err subs1 p.func in
+    let main_sub2 = Utils.find_func_err subs2 p.func in
+    Eq.exist_isomorphism main_sub1 main_sub2
+  else false
+
 (* Runs a comparative analysis. *)
 let comparative (bap_ctx : ctxt) (z3_ctx : Z3.context) (var_gen : Env.var_gen)
     (p : Params.t) (file1 : string) (file2 : string) : combined_pre =
@@ -270,8 +285,6 @@ let comparative (bap_ctx : ctxt) (z3_ctx : Z3.context) (var_gen : Env.var_gen)
   let subs2 = Term.enum sub_t prog2 in
   let main_sub1 = Utils.find_func_err subs1 p.func in
   let main_sub2 = Utils.find_func_err subs2 p.func in
-  let syntactically_equal = Eq.cmp_overall main_sub1 main_sub2 in
-  Printf.printf "\nTHE SUBS ARE EQUAL: %b\n" syntactically_equal;
   let stack_range = Utils.update_stack ~base:p.stack_base ~size:p.stack_size in
   let env2, pointer_vars_2 =
     let to_inline2 = Utils.match_inline p.inline subs2 in
@@ -361,8 +374,12 @@ let run (p : Params.t) (files : string list) (bap_ctx : ctxt)
     single bap_ctx z3_ctx var_gen p file
     |> check_pre p z3_ctx
   | [file1; file2] ->
-    comparative bap_ctx z3_ctx var_gen p file1 file2
-    |> check_pre p z3_ctx
+    begin
+      match check_syntax_equality bap_ctx p file1 file2 with
+      | false ->
+        comparative bap_ctx z3_ctx var_gen p file1 file2 |> check_pre p z3_ctx
+      | true -> Ok ()
+    end
   | _ ->
     let err =
       Printf.sprintf "WP can only analyze one binary for a single analysis or \

wp/plugin/lib/equality.ml

@@ -13,6 +13,7 @@ module TidSet = Set.Make(Tid)
 
 module TidMap = Map.Make(Tid)
 module VarMap = Map.Make(Var)
+module VarSet = Set.Make(Var)
 module TidTupleMap = Map.Make(struct type t = Tid.t * Tid.t [@@deriving sexp, compare] end)
 
 type varToVarMap = Var.t VarMap.t
@@ -296,11 +297,33 @@ let check_isomorphism
         let blk2 = TidMap.find_exn tid_to_blk2 tid2 in
         compare_blk_tid_only graph blk1 blk2)
 
-let rec get_isomorphism (candidate_map : TidSet.t TidMap.t) (used_set : TidSet.t) (node_stack : (blk term) list)
-    (graph : Tid.t TidMap.t) (tid_to_blk1 : (blk term) TidMap.t) (tid_to_blk2 : (blk term) TidMap.t): (Tid.t TidMap.t) option  =
+let merge_maps (merged_map : varToVarMap) (new_map : varToVarMap) : varToVarMap option =
+  let new_merged_map = VarMap.merge merged_map new_map
+      ~f:(fun ~key:_key v ->
+          match v with
+          | `Left v1 -> Some v1
+          | `Right v2 -> Some v2
+          | `Both (v1, v2) -> if Var.equal v1 v2 then Some v1 else None)
+  in
+  let get_keys = (fun m -> VarMap.keys m |> VarSet.of_list) in
+  let expected_len = VarSet.union (get_keys merged_map) (get_keys new_map) |> VarSet.length in
+  let actual_len = VarMap.keys new_merged_map |> List.length in
+  match actual_len = expected_len with
+    false -> None
+  | true -> Some new_merged_map
+
+let rec get_isomorphism
+    (candidate_map : TidSet.t TidMap.t) (used_set : TidSet.t)
+    (node_stack : (blk term) list) (graph : Tid.t TidMap.t)
+    (tid_to_blk1 : (blk term) TidMap.t) (tid_to_blk2 : (blk term) TidMap.t)
+    (var_maps : varToVarMap TidTupleMap.t) (merged_map : varToVarMap) :
+  (Tid.t TidMap.t) option  =
   let node = List.hd node_stack in
   match node with
-  | None -> if check_isomorphism graph tid_to_blk1 tid_to_blk2 then Some graph else None
+  | None ->
+    if check_isomorphism graph tid_to_blk1 tid_to_blk2
+    then Some graph
+    else None
   | Some node ->
     let tid = Term.tid node in
     match TidMap.find candidate_map tid with
@@ -312,58 +335,41 @@ let rec get_isomorphism (candidate_map : TidSet.t TidMap.t) (used_set : TidSet.t
             (* cannot be empty *)
             let node_stack = List.tl_exn node_stack in
             let graph = TidMap.update graph tid ~f:(fun _ -> tid_mapped_to) in
-            get_isomorphism candidate_map used_set node_stack graph tid_to_blk1 tid_to_blk2)
+            (* TODO change this to a map lookup *)
+            TidTupleMap.find var_maps (tid, tid_mapped_to) >>=
+            fun new_map ->
+            merge_maps merged_map new_map >>=
+            fun merged_map ->
+            get_isomorphism candidate_map used_set node_stack
+              graph tid_to_blk1 tid_to_blk2 var_maps merged_map)
     | None -> None
 
+let get_length cfg : int =
+  let f = ( fun node acc ->
+      let lbl = get_label node in
+      let def_len = Term.enum def_t lbl |> Sequence.to_list |> List.length in
+      let jmp_len = Term.enum jmp_t lbl |> Sequence.to_list |> List.length in
+      acc + def_len + jmp_len
+    ) in
+  BFS.fold f 0 cfg
+
 let exist_isomorphism (sub1: Sub.t) (sub2 : Sub.t) : bool =
   let cfg1, cfg2 = Sub.to_cfg sub1, Sub.to_cfg sub2 in
+  printf "\nLENGTH IS: %d\n" (get_length cfg1);
   let tid_to_blk1, tid_to_blk2 = get_node_maps cfg1 cfg2 in
   (* the variable mappings on a per-block basis are not actually used *)
-  let tid_map, _var_maps = compare_blocks sub1 sub2 in
+  let tid_map, var_maps = compare_blocks sub1 sub2 in
+  let _ = printf "PERCENTAGE_IDENTICAL_BLOCKS_IS: %.4f\n" (((List.length (TidMap.keys tid_map)) |> float_of_int) /. ((List.length (TidMap.keys tid_to_blk1)) |> float_of_int)) in
   (* List.iter (TidMap.keys tid_map) ~f:(fun key -> printf "TID MAP: %s\n" (Tid.to_string key)) ; *)
   let used_set = TidSet.empty in
   let node_stack = TidMap.data tid_to_blk1 in
   let node2_stack = TidMap.data tid_to_blk2 in
   let graph = TidMap.empty in
-  printf "BLOCK LENGTH: %d" (List.length node_stack);
+  printf "BLOCK LENGTH: %d\n" (List.length node_stack);
   (* short circuit if we already know that the length does not match *)
   if (List.length node_stack) = (List.length node2_stack) then
-    match get_isomorphism tid_map used_set node_stack graph tid_to_blk1 tid_to_blk2 with
+    match get_isomorphism tid_map used_set node_stack graph tid_to_blk1 tid_to_blk2 var_maps VarMap.empty with
     | None -> false
-    | Some _x -> true
+    | Some _x ->
+      printf "Blocks are syntactically equal! Not performing WP analysis.\nUNSAT!\n"; true
   else false
-
-(* let evaluator_overall node1 cur_res =
- *   match cur_res with
- *   | Some (seen_set, v_map, cfg2) ->
- *     let inner_evaluator =
- *       (fun (node2 : Bap.Std.Graphs.Ir.Node.t) (found_match, seen_set, v_map) ->
- *          let blk2 = get_label node2 in
- *          let tid2 = Term.tid blk2 in
- *          match found_match with
- *          | true -> true, seen_set, v_map
- *          | false ->
- *            (\* if already matched, then don't consider again *\)
- *            if TidSet.exists seen_set ~f:(fun a -> a = tid2) then false, seen_set, v_map
- *            (\* otherwise, perform comparison *\)
- *            else
- *              match compare_block (get_label node1) blk2 v_map with
- *              | Some v_map_updated ->
- *                true, TidSet.union seen_set (TidSet.singleton tid2), v_map_updated
- *              | None -> false, seen_set, v_map) in
- *     let r, s, m = BFS.fold inner_evaluator (false, seen_set, v_map) cfg2 in
- *     if r then
- *       Some (s, m, cfg2)
- *     else None
- *   | None -> None
- *
- * (\* performs the overarching comparison for exact syntactic equality between subs*\)
- * let cmp_overall (sub1 : Sub.t) (sub2 : Sub.t) : bool =
- *   let cfg1, cfg2 = Sub.to_cfg sub1, Sub.to_cfg sub2 in
- *   let seen_set = TidSet.empty in
- *   let v_map = VarMap.empty in
- *   let r = BFS.fold evaluator_overall (Some (seen_set, v_map, cfg2)) cfg1 in
- *   match r with
- *   | Some _ -> true
- *   | None -> false *)
-

wp/plugin/lib/parameters.ml

@@ -47,7 +47,8 @@ type t = {
   debug : string list;
   stack_base : int option;
   stack_size : int option;
-  show : string list
+  syntax_equality : bool;
+  show : string list;
 }
 
 (* Ensures the user inputted a function for analysis. *)

wp/plugin/lib/parameters.mli

@@ -50,7 +50,8 @@ type t = {
   debug : string list;
   stack_base : int option;
   stack_size : int option;
-  show : string list
+  syntax_equality : bool;
+  show : string list;
 }
 
 (** [validate flags files] ensures the user inputted the appropriate flags for

wp/plugin/wp.ml

@@ -180,6 +180,13 @@ let stack_size = Cmd.parameter Typ.(some int) "stack-size"
     ~doc:{|Sets the size of the stack, which should be denoted in bytes. By
            default, the size of the stack is 0x800000 which is 8MB.|}
 
+let syntax_equality = Cmd.flag "syntax-equality"
+    ~doc:{| Short-circuits WP and returns UNSAT if the functions
+           are syntactically equal. Use with caution; the comparison must
+           be testing for function equality in the same sense as syntactic
+           equality. This is currently found with the mem_offset flag and
+           pointer flag enabled. Only applicable in the comparative case. |}
+
 let grammar = Cmd.(
     args
     $ func
@@ -200,6 +207,7 @@ let grammar = Cmd.(
     $ show
     $ stack_base
     $ stack_size
+    $ syntax_equality
     $ files)
 
 (* The callback run when the command is invoked from the command line. *)
@@ -222,6 +230,7 @@ let callback
     (show : string list)
     (stack_base : int option)
     (stack_size : int option)
+    (syntax_equality : bool)
     (files : string list)
     (ctxt : ctxt) =
   let params = Parameters.({
@@ -242,7 +251,8 @@ let callback
       debug = debug;
       show = show;
       stack_base = stack_base;
-      stack_size = stack_size
+      stack_size = stack_size;
+      syntax_equality = syntax_equality;
     })
   in
   Parameters.validate params files >>= fun () ->

wp/resources/sample_binaries/syntax_equality/run_wp.sh

@@ -12,6 +12,7 @@ run_unsat () {
     --num-unroll=0 \
     --no-byteweight \
     --mem-offset \
+    --syntax-equality \
     -- main_1 main_2
 }