Fix compatibility with WP init-mem parameter (#209)

* Collect code addrs for the verifier if `init-mem` is `true`

* Fix failing test

* Update

* Fixing tests for latest BAP

Co-authored-by: bmourad01 <[email protected]>

Author: bmourad01

bap-vibes/src/patch_c.ml

@@ -1633,7 +1633,7 @@ let data_of_tgt (target : Theory.target) : Data_model.t KB.t =
     Err.fail @@ Patch_c_error (
       Format.asprintf "Unsupported target %a"
         Theory.Target.pp target) in
-  if Theory.Target.matches target "arm" then
+  if Theory.Target.matches target "arm-family" then
     if Theory.Target.bits target = 32
     then KB.return Data_model.{sizes = `ILP32; schar = false}
     else fail ()

bap-vibes/src/pipeline.ml

@@ -114,10 +114,10 @@ let run_KB_computation (f : Data.cls KB.obj KB.t) (state : KB.state)
 
 (* This is the main CEGIS loop. It constructs a patched exe and verifies it.
    If its correct, it returns the filepath. If not, it runs again. *)
-let rec cegis ?count:(count=1) ?max_tries:(max_tries=None)
-    ~seed:(seed : Seeder.t) (config : Config.t) (orig_proj : project)
-    (orig_prog : Program.t) (state : KB.state)
-    : (string, Toplevel_error.t) result =
+let rec cegis ?(count : int = 1) ?(max_tries : int option = None)
+    ~(seed : Seeder.t) (config : Config.t) (orig_proj : project)
+    (orig_prog : Program.t) (orig_code : Bap_wp.Utils.Code_addrs.t)
+    (state : KB.state) : (string, Toplevel_error.t) result =
   Events.(send @@ Header "Starting CEGIS iteration");
   Events.(send @@ Info (Printf.sprintf "Iteration: %d" count));
 
@@ -132,27 +132,29 @@ let rec cegis ?count:(count=1) ?max_tries:(max_tries=None)
 
   (* Temporarily use the new KB state when loading the patched binary. *)
   Toplevel.set new_state;
-  let+ _, patch_prog = Utils.load_exe tmp_patched_filepath in
+  let+ patch_proj, patch_prog = Utils.load_exe tmp_patched_filepath in
+  let patch_code =
+    let open Bap_wp.Utils.Code_addrs in
+    if (Config.wp_params config).init_mem
+    then collect patch_proj
+    else empty in
   Toplevel.set state;
 
   let wp_params = Config.wp_params config in
   let target = Project.target orig_proj in
   if Config.perform_verification config then
-    let verif_res =
-      Verifier.verify target wp_params
-        ~orig_prog:(orig_prog, Config.exe config)
-        ~patch_prog:(patch_prog, tmp_patched_filepath)
-    in
+    (* let orig_code, patch_code = Addr.Set.empty, Addr.Set.empty in *)
+    let verif_res = Verifier.verify target wp_params
+        ~orig_prog:(orig_prog, Config.exe config, orig_code)
+        ~patch_prog:(patch_prog, tmp_patched_filepath, patch_code) in
     match verif_res with
+    | Error _ as err -> err
     | Ok Verifier.Done -> finalize_patched_exe value
     | Ok Verifier.Again ->
-      begin
-        let new_count = count + 1 in
-        let new_seed = Seeder.extract_seed value new_state in
-        cegis config orig_proj orig_prog state
-          ~count:new_count ~max_tries ~seed:new_seed
-      end
-    | Error e -> Error e
+      let new_count = count + 1 in
+      let new_seed = Seeder.extract_seed value new_state in
+      cegis config orig_proj orig_prog orig_code state
+        ~count:new_count ~max_tries ~seed:new_seed
   else
     let warning =
       "WARNING: No verification performed. Patched binary may be incorrect."
@@ -173,11 +175,16 @@ let run (config : Config.t) : (string, Toplevel_error.t) result =
   let filepath = Config.exe config in
   Events.(send @@ Info (Format.sprintf "Loading into BAP: %s..." filepath));
   let+ orig_proj, orig_prog = Utils.load_exe filepath in
+  let orig_code =
+    let open Bap_wp.Utils.Code_addrs in
+    if (Config.wp_params config).init_mem
+    then collect orig_proj
+    else empty in
 
   let state = Toplevel.current () in
   let computation = init config orig_proj in
   let+ obj, new_state = run_KB_computation computation state in
   let seed = Seeder.extract_seed obj new_state in
 
   let max_tries = Config.max_tries config in
-  cegis config orig_proj orig_prog new_state ~max_tries ~seed
+  cegis config orig_proj orig_prog orig_code new_state ~max_tries ~seed

bap-vibes/src/verifier.ml

@@ -52,18 +52,28 @@ let wp_verifier (p : Params.t) (inputs : Runner.input list) =
 *)
 let verify
     ?(verifier = wp_verifier)
-    ~(orig_prog : Program.t * string)
-    ~(patch_prog : Program.t * string)
+    ~(orig_prog : Program.t * string * Bap_wp.Utils.Code_addrs.t)
+    ~(patch_prog : Program.t * string * Bap_wp.Utils.Code_addrs.t)
     (tgt : Theory.target)
     (params : Params.t) : (next_step, Toplevel_error.t) result =
-  let prog1, name1 = orig_prog in
-  let prog2, name2 = patch_prog in
+  let prog1, name1, code1 = orig_prog in
+  let prog2, name2, code2 = patch_prog in
   Events.(send @@ Header "Starting Verifier");
   Events.(send @@ Info "Beginning weakest-precondition analysis...");
   Events.(send @@ Info (sprintf "Original program: %s" name1));
   Events.(send @@ Info (sprintf "Patched program: %s" name2));
-  let input1 = Runner.{program = prog1; target = tgt; filename = name1} in
-  let input2 = Runner.{program = prog2; target = tgt; filename = name2} in
+  let input1 = Runner.{
+      program = prog1;
+      target = tgt;
+      filename = name1;
+      code_addrs = code1;
+    } in
+  let input2 = Runner.{
+      program = prog2;
+      target = tgt;
+      filename = name2;
+      code_addrs = code2;
+    } in
   let* status =
     verifier params [input1; input2] |>
     Result.map_error ~f:(fun e -> Toplevel_error.WP_failure e) in

bap-vibes/src/verifier.mli

@@ -45,8 +45,8 @@ type next_step = Done | Again
     verifier's results. *)
 val verify :
   ?verifier:(verifier) ->
-  orig_prog:(Program.t * string) ->
-  patch_prog:(Program.t * string) ->
+  orig_prog:(Program.t * string * Utils.Code_addrs.t) ->
+  patch_prog:(Program.t * string * Utils.Code_addrs.t) ->
   Theory.target ->
   Run_parameters.t ->
   (next_step, Toplevel_error.t) result

bap-vibes/tests/integration/test_minizinc.ml

@@ -110,7 +110,7 @@ let ex1 : Ir.t = {
   blks = [blk1]
 }
 
-let arm_tgt = Arm_target.LE.v7
+let arm_tgt = Arm_target.LE.v7a
 let arm_lang = Arm_target.llvm_a32
 
 (* Ensure minizinc produces the expected output on our sample IR block. *)

bap-vibes/tests/unit/helpers.ml

@@ -59,9 +59,9 @@ let dummy_code : Cabs.definition =
 
 (* Get an empty program that can be used in tests. *)
 let prog_exn (proj : (Project.t * string, Error.t) result)
-    : Program.t * string =
+    : Program.t * string * Bap_wp.Utils.Code_addrs.t =
   let p, s = proj_exn proj in
-  (Project.program p, s)
+  (Project.program p, s, Bap_wp.Utils.Code_addrs.empty)
 
 (* Some dummy values that can be used in tests. *)
 let the_target () = Theory.Target.read ~package:"bap" "armv7+le"

bap-vibes/tests/unit/test_parse_c.ml

@@ -132,7 +132,7 @@ let test_call_hex _ =
     "int temp;
      if (temp == 0) { ((void (*)())0x3ec)(); }"
     "{
-       if (~temp) {
+       if (temp = 0) {
          jmp 0x3EC
        }
      }"
@@ -427,7 +427,7 @@ let test_bool_not _ =
        x = y;
      }"
     "{
-       if (~x) {
+       if (x = 0) {
          x := y
        }
      }"
@@ -437,7 +437,7 @@ let test_mixed_sorts _ =
     "int x, y;
      x += (y == 0);"
     "{
-       x := x + pad:32[~y]
+       x := x + pad:32[y = 0]
      }"
 
 let suite = [

bap-vibes/tests/unit/test_substituter.ml

@@ -35,7 +35,7 @@ module Test_result = struct
   let result = KB.Class.property ~package cls "test-result" domain
 end
 
-let x86_tgt = Theory.Target.get ~package:"bap" "amd64"
+let x86_tgt = Theory.Target.get ~package:"bap" "x86_64"
 
 (* Very lax equality to avoid tid comparison failures *)
 let eq_elt e1 e2 =

resources/exes/arm-isel-jmp2/config.json

@@ -1,5 +1,5 @@
 {
-  "max-tries": 1,
+  "max-tries": 10,
   "wp-params": {
     "func": "main",
     "postcond": "(assert (= R0_mod #x00000006))"
@@ -21,7 +21,6 @@
      "patch-size" : 128,
      "patch-vars": [
        {"name": "retvar",
-	    "at-entry": "R0",
         "at-exit": "R0"
        }
      ]

resources/exes/arm-simple-sum/config.json

@@ -1,5 +1,5 @@
 {
-  "max-tries": 1,
+  "max-tries": 10,
   "wp-params": {
     "func": "main",
     "postcond": "(assert (= R0_mod #x00000009))"
@@ -16,7 +16,6 @@
      "patch-size" : 128,
      "patch-vars": [
        {"name": "retvar",
-	    "at-entry": "R0",
         "at-exit": "R0"
        }
      ]