staged concolic miniwasm

.github/workflows/scala.yml

@@ -79,3 +79,4 @@ jobs:
         sbt 'testOnly gensym.wasm.TestConcolic'
         sbt 'testOnly gensym.wasm.TestDriver'
         sbt 'testOnly gensym.wasm.TestStagedEval'
+        sbt 'testOnly gensym.wasm.TestStagedConcolicEval'

benchmarks/wasm/branch-strip-buggy.wat

@@ -29,6 +29,7 @@
       else
         i32.const 0
         call 2
+        i32.const 1 ;; to satisfy the type checker, this line will never be reached
       end
     end
   )

benchmarks/wasm/staged/brtable_concolic.wat

@@ -0,0 +1,22 @@
+(module $brtable
+  (global (;0;) (mut i32) (i32.const 1048576))
+  (type (;0;) (func (param i32)))
+  (func (;0;) (type 1) (result i32)
+    i32.const 2
+    (block
+      (block
+        (block
+          i32.const 0
+          i32.symbolic
+          br_table 0 1 2 0 ;; br_table will consume an element from the stack
+        )
+        i32.const 1
+        call 1
+        br 1
+      )
+      i32.const 0
+      call 1
+    )
+  )
+  (import "console" "assert" (func (type 0)))
+  (start 0))

headers/wasm.hpp

@@ -2,5 +2,7 @@
 #define WASM_HEADERS
 
 #include "wasm/concrete_rt.hpp"
-
+#include "wasm/symbolic_rt.hpp"
+#include "wasm/concolic_driver.hpp"
+#include "wasm/utils.hpp"
 #endif

headers/wasm/concolic_driver.hpp

@@ -0,0 +1,105 @@
+#ifndef CONCOLIC_DRIVER_HPP
+#define CONCOLIC_DRIVER_HPP
+
+#include "concrete_rt.hpp"
+#include "smt_solver.hpp"
+#include "symbolic_rt.hpp"
+#include <functional>
+#include <ostream>
+#include <string>
+#include <vector>
+
+class ConcolicDriver {
+  friend class ManagedConcolicCleanup;
+
+public:
+  ConcolicDriver(std::function<void()> entrypoint, std::string tree_file)
+      : entrypoint(entrypoint), tree_file(tree_file) {}
+  ConcolicDriver(std::function<void()> entrypoint)
+      : entrypoint(entrypoint), tree_file(std::nullopt) {}
+  void run();
+
+private:
+  Solver solver;
+  std::function<void()> entrypoint;
+  std::optional<std::string> tree_file;
+};
+
+class ManagedConcolicCleanup {
+  const ConcolicDriver &driver;
+
+public:
+  ManagedConcolicCleanup(const ConcolicDriver &driver) : driver(driver) {}
+  ~ManagedConcolicCleanup() {
+    if (driver.tree_file.has_value())
+      ExploreTree.dump_graphviz(driver.tree_file.value());
+  }
+};
+
+inline void ConcolicDriver::run() {
+  ManagedConcolicCleanup cleanup{*this};
+  while (true) {
+    ExploreTree.reset_cursor();
+
+    auto unexplored = ExploreTree.pick_unexplored();
+    if (!unexplored) {
+      std::cout << "No unexplored nodes found, exiting..." << std::endl;
+      return;
+    }
+    auto cond = unexplored->collect_path_conds();
+    auto result = solver.solve(cond);
+    if (!result.has_value()) {
+      // TODO: current implementation is buggy, there could be other reachable
+      // unexplored paths
+      std::cout << "Found an unreachable path, marking it as unreachable..."
+                << std::endl;
+      unexplored->fillUnreachableNode();
+      continue;
+    }
+    auto new_env = result.value();
+    SymEnv.update(std::move(new_env));
+    try {
+      entrypoint();
+      std::cout << "Execution finished successfully with symbolic environment:"
+                << std::endl;
+      std::cout << SymEnv.to_string() << std::endl;
+    } catch (...) {
+      ExploreTree.fillFailedNode();
+      std::cout << "Caught runtime error with symbolic environment:"
+                << std::endl;
+      std::cout << SymEnv.to_string() << std::endl;
+      return;
+    }
+  }
+}
+
+static std::monostate reset_stacks() {
+  Stack.reset();
+  Frames.reset();
+  SymStack.reset();
+  SymFrames.reset();
+  initRand();
+  Memory = Memory_t(1);
+  return std::monostate{};
+}
+
+static void start_concolic_execution_with(
+    std::function<std::monostate(std::monostate)> entrypoint,
+    std::string tree_file) {
+  ConcolicDriver driver([=]() { entrypoint(std::monostate{}); }, tree_file);
+  driver.run();
+}
+
+static void start_concolic_execution_with(
+    std::function<std::monostate(std::monostate)> entrypoint) {
+
+  const char *env_tree_file = std::getenv("TREE_FILE");
+
+  ConcolicDriver driver =
+      env_tree_file ? ConcolicDriver([=]() { entrypoint(std::monostate{}); },
+                                     env_tree_file)
+                    : ConcolicDriver([=]() { entrypoint(std::monostate{}); });
+  driver.run();
+}
+
+#endif // CONCOLIC_DRIVER_HPP

headers/wasm/concrete_rt.hpp

@@ -1,3 +1,6 @@
+#ifndef WASM_CONCRETE_RT_HPP
+#define WASM_CONCRETE_RT_HPP
+
 #include <cassert>
 #include <cstdint>
 #include <cstdio>
@@ -49,8 +52,6 @@ static Num I32V(int v) { return v; }
 
 static Num I64V(int64_t v) { return v; }
 
-using Slice = std::vector<Num>;
-
 const int STACK_SIZE = 1024 * 64;
 
 class Stack_t {
@@ -115,9 +116,12 @@ class Stack_t {
   }
 
   void initialize() {
-    // do nothing for now
+    // todo: remove this method
+   reset();
   }
 
+  void reset() { count = 0; }
+
 private:
   int32_t count;
   Num *stack_ptr;
@@ -148,6 +152,8 @@ class Frames_t {
     count += size;
   }
 
+  void reset() { count = 0; }
+
 private:
   int32_t count;
   Num *stack_ptr;
@@ -200,4 +206,6 @@ struct Memory_t {
   }
 };
 
-static Memory_t Memory(1); // 1 page memory
+static Memory_t Memory(1); // 1 page memory
+
+#endif // WASM_CONCRETE_RT_HPP

headers/wasm/smt_solver.hpp

@@ -0,0 +1,126 @@
+#ifndef SMT_SOLVER_HPP
+#define SMT_SOLVER_HPP
+
+#include "concrete_rt.hpp"
+#include "symbolic_rt.hpp"
+#include "z3++.h"
+#include <array>
+#include <set>
+#include <string>
+#include <tuple>
+#include <vector>
+
+class Solver {
+public:
+  Solver() {}
+  std::optional<std::vector<Num>> solve(const std::vector<SymVal> &conditions) {
+    // make an conjunction of all conditions
+    z3::expr conjunction = z3_ctx.bool_val(true);
+    for (const auto &cond : conditions) {
+      auto z3_cond = build_z3_expr(cond);
+      conjunction = conjunction && z3_cond != z3_ctx.bv_val(0, 32);
+    }
+#ifdef DEBUG
+    std::cout << "Symbolic conditions size: " << conditions.size() << std::endl;
+    std::cout << "Solving conditions: " << conjunction << std::endl;
+#endif
+    // call z3 to solve the condition
+    z3::solver z3_solver(z3_ctx);
+    z3_solver.add(conjunction);
+    switch (z3_solver.check()) {
+    case z3::unsat:
+      return std::nullopt; // No solution found
+    case z3::sat: {
+      z3::model model = z3_solver.get_model();
+      std::vector<Num> result;
+      // Reference:
+      // https://github.com/Z3Prover/z3/blob/master/examples/c%2B%2B/example.cpp#L59
+
+      std::cout << "Solved Z3 model" << std::endl << model << std::endl;
+      for (unsigned i = 0; i < model.size(); ++i) {
+        z3::func_decl var = model[i];
+        z3::expr value = model.get_const_interp(var);
+        std::string name = var.name().str();
+        if (name.starts_with("s_")) {
+          int id = std::stoi(name.substr(2));
+          if (id >= result.size()) {
+            result.resize(id + 1);
+          }
+          result[id] = Num(value.get_numeral_int64());
+        } else {
+          std::cout << "Find a variable that is not created by GenSym: " << name
+                    << std::endl;
+        }
+      }
+      return result;
+    }
+    case z3::unknown:
+      throw std::runtime_error("Z3 solver returned unknown status");
+    }
+    return std::nullopt; // Should not reach here
+  }
+
+private:
+  z3::context z3_ctx;
+  z3::expr build_z3_expr(const SymVal &sym_val);
+};
+
+inline z3::expr Solver::build_z3_expr(const SymVal &sym_val) {
+  if (auto sym = std::dynamic_pointer_cast<Symbol>(sym_val.symptr)) {
+    return z3_ctx.bv_const(("s_" + std::to_string(sym->get_id())).c_str(), 32);
+  } else if (auto concrete =
+                 std::dynamic_pointer_cast<SymConcrete>(sym_val.symptr)) {
+    return z3_ctx.bv_val(concrete->value.value, 32);
+  } else if (auto binary =
+                 std::dynamic_pointer_cast<SymBinary>(sym_val.symptr)) {
+    auto bit_width = 32;
+    z3::expr zero_bv =
+        z3_ctx.bv_val(0, bit_width); // Represents 0 as a 32-bit bitvector
+    z3::expr one_bv =
+        z3_ctx.bv_val(1, bit_width); // Represents 1 as a 32-bit bitvector
+
+    z3::expr left = build_z3_expr(binary->lhs);
+    z3::expr right = build_z3_expr(binary->rhs);
+    // TODO: make sure the semantics of these operations are aligned with wasm
+    switch (binary->op) {
+    case EQ: {
+      auto temp_bool = left == right;
+      return z3::ite(temp_bool, one_bv, zero_bv);
+    }
+    case NEQ: {
+      auto temp_bool = left != right;
+      return z3::ite(temp_bool, one_bv, zero_bv);
+    }
+    case LT: {
+      auto temp_bool = left < right;
+      return z3::ite(temp_bool, one_bv, zero_bv);
+    }
+    case LEQ: {
+      auto temp_bool = left <= right;
+      return z3::ite(temp_bool, one_bv, zero_bv);
+    }
+    case GT: {
+      auto temp_bool = left > right;
+      return z3::ite(temp_bool, one_bv, zero_bv);
+    }
+    case GEQ: {
+      auto temp_bool = left >= right;
+      return z3::ite(temp_bool, one_bv, zero_bv);
+    }
+    case ADD: {
+      return left + right;
+    }
+    case SUB: {
+      return left - right;
+    }
+    case MUL: {
+      return left * right;
+    }
+    case DIV: {
+      return left / right;
+    }
+    }
+  }
+  throw std::runtime_error("Unsupported symbolic value type");
+}
+#endif // SMT_SOLVER_HPP