Initial commit of precondition smtlib code

minor fixes to formatting and 1 silly bug

wp/lib/bap_wp/src/compare.ml

@@ -116,19 +116,43 @@ let compare_subs_empty_post
   in
   compare_subs ~postcond:postcond ~hyps:hyps ~original:original ~modified:modified
 
+(** [mk_smtlib2_compare] builds a constraint out of an smtlib2 string that can be used 
+    as a comparison predicate between an original and modified binary. *)
+let mk_smtlib2_compare (env1 : Env.t) (env2 : Env.t) (smtlib_str : string) : Constr.t =
+  let var_map1 = (Env.get_var_map env1) in
+  let var_map2 = (Env.get_var_map env2) in
+  let smtlib_str = Env.EnvMap.fold var_map1 ~init:smtlib_str 
+      ~f:(fun ~key:var ~data:z3_var smtlib_str ->
+          String.substr_replace_all smtlib_str ~pattern:((Var.name var) ^ "_orig") ~with_:(Z3.Expr.to_string z3_var)
+        ) in
+  let smtlib_str = Env.EnvMap.fold var_map2 ~init:smtlib_str 
+      ~f:(fun ~key:var ~data:z3_var smtlib_str ->
+          String.substr_replace_all smtlib_str ~pattern:((Var.name var) ^ "_mod") ~with_:(Z3.Expr.to_string z3_var)
+        ) in
+  info "New smtlib string: %s \n" smtlib_str;
+  let declsym1 = Env.get_decls_and_symbols env1 in
+  let declsym2 = Env.get_decls_and_symbols env2 in
+  let declsym = declsym1 @ declsym2 in
+  let ctx = Env.get_context env1 in
+  Constr.mk_smtlib2 ctx smtlib_str declsym
+
 let compare_subs_eq
     ~input:(input : Var.Set.t)
     ~output:(output : Var.Set.t)
     ~original:(original : Sub.t * Env.t)
     ~modified:(modified : Sub.t * Env.t)
+    ~smtlib_post:(smtlib_post : string)
+    ~smtlib_pre:(smtlib_pre : string)
   : Constr.t * Env.t * Env.t =
   let postcond ~original:(_, env1) ~modified:(_, env2) ~rename_set:_ =
     let post_eqs, env1, env2 = set_to_eqs env1 env2 output in
-    Constr.mk_clause [] post_eqs, env1, env2
+    let post' = mk_smtlib2_compare env1 env2 smtlib_post in 
+    Constr.mk_clause [] (post' :: post_eqs), env1, env2
   in
   let hyps ~original:(_, env1) ~modified:(_, env2) ~rename_set:_ =
     let pre_eqs, env1, env2 = set_to_eqs env1 env2 input in
-    Constr.mk_clause [] pre_eqs, env1, env2
+    let pre' = mk_smtlib2_compare env1 env2 smtlib_pre in 
+    Constr.mk_clause [] (pre' :: pre_eqs), env1, env2
   in
   compare_subs ~postcond:postcond ~hyps:hyps ~original:original ~modified:modified
 

wp/lib/bap_wp/src/compare.mli

@@ -53,6 +53,8 @@ val compare_subs_eq
   -> output:Bap.Std.Var.Set.t
   -> original:(Bap.Std.Sub.t * Env.t)
   -> modified:(Bap.Std.Sub.t * Env.t)
+  -> smtlib_post:string
+  -> smtlib_pre:string
   -> Constr.t * Env.t * Env.t
 
 (** Compare two subroutines by composition for an empty postcondition:

wp/lib/bap_wp/src/constraint.ml

@@ -269,3 +269,24 @@ let get_refuted_goals ?filter_out:(filter_out = []) (constr : t)
       worker e current_path current_registers (olds @ o) (news @ n')
   in
   worker constr Jmp.Map.empty Jmp.Map.empty [] []
+
+
+let mk_smtlib2 (ctx : Z3.context) (smtlib_str : string) (decl_syms : (Z3.FuncDecl.func_decl * Z3.Symbol.symbol) list) : t =
+  let fun_decls, fun_symbols = List.unzip decl_syms in
+  let sort_symbols = [] in
+  let sorts = [] in
+  let asts = Z3.SMT.parse_smtlib2_string ctx smtlib_str
+      sort_symbols
+      sorts
+      fun_symbols
+      fun_decls
+  in
+  let goals = List.map (Z3.AST.ASTVector.to_expr_list asts)
+      ~f:(fun e ->
+          e
+          |> mk_goal (Expr.to_string e)
+          |> mk_constr)
+  in
+  mk_clause [] goals
+
+

wp/lib/bap_wp/src/constraint.mli

@@ -128,3 +128,8 @@ val eval_model_exn : Z3.Model.model -> z3_expr -> z3_expr
     if retrieving the model fails in the case that [check] was not yet called
     to the solver or the result of the last [check] did not result in SAT. *)
 val get_model_exn : Z3.Solver.solver -> Z3.Model.model
+
+(** [mk_smtlib2] parses a string in the context that already has a mapping of func_decl to symbols 
+    and returns a constraint [Constr.t] *)
+
+val mk_smtlib2 : Z3.context -> string -> ((Z3.FuncDecl.func_decl * Z3.Symbol.symbol) list)  -> t

wp/lib/bap_wp/src/environment.ml

@@ -255,6 +255,9 @@ let set_freshen (env : t) (freshen : bool) = { env with freshen = freshen }
 let add_var (env : t) (v : Var.t) (x : Constr.z3_expr) : t =
   { env with var_map = EnvMap.set env.var_map ~key:v ~data:x }
 
+let remove_var (env :t) (v : Var.t) : t =
+  { env with var_map = EnvMap.remove env.var_map v }
+
 let mk_exp_conds (env : t) (e : exp) : Constr.goal list * Constr.goal list =
   let { exp_conds; _ } = env in
   let conds = List.map exp_conds ~f:(fun gen -> gen env e) in
@@ -328,3 +331,27 @@ let is_x86 (a : Arch.t) : bool =
 
 let use_input_regs (env : t) : bool =
   env.fun_input_regs
+
+let get_decls_and_symbols (env : t) : ((Z3.FuncDecl.func_decl * Z3.Symbol.symbol) list) = 
+  let var_map = get_var_map env in
+  let ctx = get_context env in
+  EnvMap.fold var_map ~init:[]
+    ~f:(fun ~key:_ ~data:z3_var decls ->
+        assert (Z3.Expr.is_const z3_var);
+        let decl = Z3.FuncDecl.mk_const_decl_s ctx
+            (Z3.Expr.to_string z3_var)
+            (Z3.Expr.get_sort z3_var) in
+        let sym =  Z3.Symbol.mk_string ctx (Z3.Expr.to_string z3_var) in
+        (decl,sym)::decls
+      )
+
+let mk_smtlib2_single (env : t) (smt_post : string) : Constr.t =
+  let var_map = get_var_map env in
+  let smt_post = EnvMap.fold var_map ~init:smt_post 
+      ~f:(fun ~key:var ~data:z3_var smt_post ->
+          String.substr_replace_all smt_post ~pattern:((Var.name var) ^ " ") ~with_:((Z3.Expr.to_string z3_var) ^ " ")
+        ) in
+  info "New smt-lib string : %s\n" smt_post;
+  let decl_syms = get_decls_and_symbols env in
+  let ctx = get_context env in
+  Constr.mk_smtlib2 ctx smt_post decl_syms 

wp/lib/bap_wp/src/environment.mli

@@ -127,6 +127,9 @@ val wp_rec_call :
     typically a constant. *)
 val add_var : t -> Bap.Std.Var.t -> Constr.z3_expr -> t
 
+(** Remove a binding in the environment for a bap variable. *)
+val remove_var : t -> Bap.Std.Var.t -> t
+
 (** Add a precondition to be associated to a block b to the environment. *)
 val add_precond : t -> Bap.Std.Tid.t -> Constr.t -> t
 
@@ -209,4 +212,20 @@ val mk_z3_expr : Z3.context -> name:string -> typ:Bap.Std.Type.t -> Constr.z3_ex
     constant is "fresh" with the {!Environment.var_gen}. *)
 val new_z3_expr : ?name:string -> t -> Bap.Std.Type.t -> Constr.z3_expr
 
+
+(** [get_decls_and_symbols] builds from a the var_map in an environment 
+    a mapping of all Z3 func_decl to their symbol. This is a helper function for
+    [mk_smtlib2] *)
+
+val get_decls_and_symbols : t -> ((Z3.FuncDecl.func_decl * Z3.Symbol.symbol) list) 
+
+
+(** [mk_smtlib2_single env smtlib_str] takes in a string representing a
+    valid SMT-Lib-2 statement. 
+    The variables in the SMT-Lib statements need to appear in the
+    environment. The intended purpose of this function is generating hypothesis
+     and postconditions for single binary analysis *)
+val mk_smtlib2_single : t -> string -> Constr.t
+
+
 (*---------------------------------------------------*)

wp/lib/bap_wp/src/precondition.ml

@@ -217,11 +217,10 @@ let exp_to_z3 (exp : Exp.t) (env : Env.t) : Constr.z3_expr * Constr.t list * Con
       debug "Visiting let %s = %s in %s%!"
         (Var.to_string v) (Exp.to_string exp) (Exp.to_string body);
       let exp_val, env = exp_to_z3_body exp env in
-      (* FIXME: we're handling this incorrectly! The variable should
-         be removed from the context after leaving the scope of the
-         Let! *)
       let env' = Env.add_var env v exp_val in
-      exp_to_z3_body body env'
+      let z3_expr, env = exp_to_z3_body body env' in
+      let env = Env.remove_var env v in
+      (z3_expr, env)
     | Unknown (str, typ) ->
       debug "Visiting unknown: %s  Type:%s%!" str (Type.to_string typ);
       Env.new_z3_expr env ~name:("unknown_" ^ str) typ, env
@@ -847,42 +846,6 @@ let non_null_assert : Env.exp_cond = fun env exp ->
   else
     Some (Assume (Constr.mk_goal "assume" (Bool.mk_and ctx conds)))
 
-let mk_smtlib2_post (env : Env.t) (smt_post : string) : Constr.t =
-  let ctx = Env.get_context env in
-  let sort_symbols = [] in
-  let sorts = [] in
-  let fun_decls =
-    Env.EnvMap.fold (Env.get_var_map env) ~init:[]
-      ~f:(fun ~key:_ ~data:z3_var decls ->
-          assert (Z3.Expr.is_const z3_var);
-          Z3.FuncDecl.mk_const_decl_s ctx
-            (Z3.Expr.to_string z3_var)
-            (Z3.Expr.get_sort z3_var)
-          ::decls
-        )
-  in
-  let fun_symbols =
-    Env.EnvMap.fold (Env.get_var_map env) ~init:[]
-      ~f:(fun ~key:_ ~data:z3_var decls ->
-          assert (Z3.Expr.is_const z3_var);
-          Z3.Symbol.mk_string ctx
-            (Z3.Expr.to_string z3_var)
-          ::decls
-        )
-  in
-  let asts = Z3.SMT.parse_smtlib2_string ctx smt_post
-      sort_symbols
-      sorts
-      fun_symbols
-      fun_decls
-  in
-  let goals = List.map (Z3.AST.ASTVector.to_expr_list asts)
-      ~f:(fun e ->
-          e
-          |> Constr.mk_goal (Expr.to_string e)
-          |> Constr.mk_constr)
-  in
-  Constr.mk_clause [] goals
 
 let check ?refute:(refute = true) (solver : Solver.solver) (ctx : Z3.context)
     (pre : Constr.t) : Solver.status =
@@ -895,6 +858,7 @@ let check ?refute:(refute = true) (solver : Solver.solver) (ctx : Z3.context)
     else
       pre'
   in
+  Printf.printf "Z3 Query:\n%s\n" (Z3.Expr.to_string (Z3.Expr.simplify is_correct None));
   let () = Z3.Solver.add solver [is_correct] in
   Z3.Solver.check solver []
 

wp/lib/bap_wp/src/precondition.mli

@@ -196,14 +196,6 @@ val mk_env
   -> Env.var_gen
   -> Env.t
 
-(** [mk_smtlib2_post env sub post] takes in a string representing a
-    valid SMT-Lib-2 statement, and turns it into a valid postcondition
-    to be passed to {!visit_sub}.
-
-    The variables in the SMT-Lib statements need to appear in the
-    environment. *)
-val mk_smtlib2_post : Env.t -> string -> Constr.t
-
 (** Create a precondition for a given jump expression, depending on the postcondition
     and (potentially) the preconditions for the jump targets or the loop invariants.
 

wp/lib/bap_wp/tests/unit/test_compare.ml

@@ -123,7 +123,7 @@ let test_sub_pair_1 (test_ctx : test_ctxt) : unit =
   let output_vars = Var.Set.singleton z in
   let compare_prop, _, _ = Comp.compare_subs_eq
       ~input:input_vars ~output:output_vars
-      ~original:(sub1,env1) ~modified:(sub2,env2) in
+      ~original:(sub1,env1) ~modified:(sub2,env2) ~smtlib_post:"" ~smtlib_pre:"" in
   assert_z3_compare test_ctx ~orig:env1 ~modif:env2
     (Sub.to_string sub1) (Sub.to_string sub2)
     compare_prop Z3.Solver.UNSATISFIABLE
@@ -163,7 +163,7 @@ let test_sub_pair_2 (test_ctx : test_ctxt) : unit =
   let output_vars = Var.Set.singleton z in
   let compare_prop, _, _ = Comp.compare_subs_eq
       ~input:input_vars ~output:output_vars
-      ~original:(sub1,env1) ~modified:(sub2,env2) in
+      ~original:(sub1,env1) ~modified:(sub2,env2) ~smtlib_post:"" ~smtlib_pre:"" in
   assert_z3_compare test_ctx ~orig:env1 ~modif:env2
     (Sub.to_string sub1) (Sub.to_string sub2)
     compare_prop Z3.Solver.SATISFIABLE
@@ -197,7 +197,7 @@ let test_sub_pair_3 (test_ctx : test_ctxt) : unit =
   let output_vars = Var.Set.singleton z in
   let compare_prop, _, _ = Comp.compare_subs_eq
       ~input:input_vars ~output:output_vars
-      ~original:(sub1,env1) ~modified:(sub2,env2) in
+      ~original:(sub1,env1) ~modified:(sub2,env2) ~smtlib_post:"" ~smtlib_pre:"" in
   assert_z3_compare test_ctx ~orig:env1 ~modif:env2
     (Sub.to_string sub1) (Sub.to_string sub2)
     compare_prop Z3.Solver.UNSATISFIABLE
@@ -234,7 +234,7 @@ let test_sub_pair_4 (test_ctx : test_ctxt) : unit =
   let output_vars = Var.Set.singleton y in
   let compare_prop, _, _ = Comp.compare_subs_eq
       ~input:input_vars ~output:output_vars
-      ~original:(sub1,env1) ~modified:(sub2,env2) in
+      ~original:(sub1,env1) ~modified:(sub2,env2) ~smtlib_post:"" ~smtlib_pre:"" in
   assert_z3_compare test_ctx ~orig:env1 ~modif:env2
     (Sub.to_string sub1) (Sub.to_string sub2)
     compare_prop Z3.Solver.UNSATISFIABLE
@@ -271,7 +271,7 @@ let test_sub_pair_5 (test_ctx : test_ctxt) : unit =
   let output_vars = Var.Set.singleton y in
   let compare_prop, _, _ = Comp.compare_subs_eq
       ~input:input_vars ~output:output_vars
-      ~original:(sub1,env1) ~modified:(sub2,env2) in
+      ~original:(sub1,env1) ~modified:(sub2,env2) ~smtlib_post:"" ~smtlib_pre:"" in
   assert_z3_compare test_ctx ~orig:env1 ~modif:env2
     (Sub.to_string sub1) (Sub.to_string sub2)
     compare_prop Z3.Solver.SATISFIABLE
@@ -462,7 +462,7 @@ let test_fun_outputs_1 (test_ctx : test_ctxt) : unit =
   let output_vars = Var.Set.singleton ret_var in
   let compare_prop, _, _ = Comp.compare_subs_eq
       ~input:input_vars ~output:output_vars
-      ~original:(main_sub1, env1) ~modified:(main_sub2, env2) in
+      ~original:(main_sub1, env1) ~modified:(main_sub2, env2) ~smtlib_post:"" ~smtlib_pre:"" in
   assert_z3_compare test_ctx ~orig:env1 ~modif:env2
     (Sub.to_string main_sub1)
     (Sub.to_string main_sub2)
@@ -495,7 +495,7 @@ let test_fun_outputs_2 (test_ctx : test_ctxt) : unit =
   let output_vars = Var.Set.singleton ret_var in
   let compare_prop, _, _ = Comp.compare_subs_eq
       ~input:input_vars ~output:output_vars
-      ~original:(main_sub1, env1) ~modified:(main_sub2, env2) in
+      ~original:(main_sub1, env1) ~modified:(main_sub2, env2) ~smtlib_post:"" ~smtlib_pre:"" in
   assert_z3_compare test_ctx ~orig:env1 ~modif:env2
     (Sub.to_string main_sub1)
     (Sub.to_string main_sub2)
@@ -599,7 +599,7 @@ let test_sub_pair_mem_1 (test_ctx : test_ctxt) : unit =
   let output_vars = Var.Set.singleton mem in
   let compare_prop, _, _ = Comp.compare_subs_eq
       ~input:input_vars ~output:output_vars
-      ~original:(sub1,env1) ~modified:(sub2,env2) in
+      ~original:(sub1,env1) ~modified:(sub2,env2) ~smtlib_post:"" ~smtlib_pre:"" in
   assert_z3_compare test_ctx ~orig:env1 ~modif:env2
     (Sub.to_string sub1) (Sub.to_string sub2)
     compare_prop Z3.Solver.SATISFIABLE

wp/lib/bap_wp/tests/unit/test_precondition.ml

@@ -374,7 +374,7 @@ let test_subroutine_1_2 (test_ctx : test_ctxt) : unit =
   let env = Env.add_var env z (mk_z3_var env z) in
   (* The names x0, y0 and z0 are magical, as they are generated by BAP
      (with the "base names" x, y, z)*)
-  let post = Pre.mk_smtlib2_post env
+  let post = Env.mk_smtlib2_single env
       "(assert (and (bvsle (bvsub z0 x0) #x00000001) (bvsle #xffffffff (bvsub z0 x0))))"
   in
   let pre, _ = Pre.visit_sub env post sub in

wp/plugin/README.md

@@ -83,6 +83,9 @@ UNSAT!
 
 Meaning there is no way to reach the `assert(0)` statement.
 
+Alternatively one may assert custom assumptions specified as smt-lib expressions
+via the command line option `--wp-precond="(assert (= RDI #x0000000000000002))"`
+
 -------
 
 A more sophisticated example involves comparing two different
@@ -323,11 +326,15 @@ The various options are:
   function call sites. For example, If you want to inline everything, set to
   --wp-inline=.\* or --wp-inline=foo|bar will inline the functions foo and bar.
 
+- `--wp-precond=smt-lib-string`. If present, allows the introduction of assertions
+  to the beginning of a query. This allows pruning of possible models. For 
+  comparative predicates, one may refer to variables in the original and modified program by appending the suffix "_orig" and "_mod" to variable names in the smt-lib expression.
+  For examples `--wp-precond="(assert (= RDI_mod #x0000000000000003))  (assert (= RDI_orig #x0000000000000003))"`
+
 - `--wp-postcond=smt-lib-string`. If present, replaces the
   default post-condition by the user-specified one, using the
-  [smt-lib2] format. At the moment, the names of variables
-  representing memory and registers are a bit magical, so consider
-  this to be an experimental feature.
+  [smt-lib2] format. Similar to `--wp-precond`, one may create comparative 
+  postconditions on variables by appending "_orig" and "_mod" to register names.
 
 - `--wp-num-unroll=num`. If present, replaces the default number of
   times to unroll each loop. The number of loop unrollings is 5 by default.

wp/plugin/tests/test_wp.ml

@@ -119,7 +119,7 @@ let suite = [
   "Function Spec: no inlining"     >:: test_single_elf "function_spec" "main" "SAT!" ~inline:"NONEXISTENTGARBAGE";
   "Nested Function Calls"          >:: test_single_elf "nested_function_calls" "main" "SAT!" ~inline:"foo|bar";
   "Nested Calls: inline all"       >:: test_single_elf "nested_function_calls" "main" "SAT!" ~inline:".*";
-  "User Defined Postcondition"     >:: test_single_elf "return_argc" "main" "SAT!" ~post:"(assert (= RAX0 #x0000000000000000))";
+  "User Defined Postcondition"     >:: test_single_elf "return_argc" "main" "SAT!" ~post:"(assert (= RAX #x0000000000000000))";
 
   "Update Number of Unrolls"       >:: test_update_num_unroll (Some 3);
   "Original Number of Unrolls"     >:: test_update_num_unroll None;