Network

driver/aflpp/cov.h

@@ -10,13 +10,31 @@ namespace rgd {
 
 struct BranchContext {
   void *addr;
-  uint32_t id;
   bool direction;
+};
+
+struct HybridBranchContext : public BranchContext {
+  uint32_t id;
+};
+
+struct ContextAwareBranchContext : public BranchContext {
   uint32_t context;
+};
+
+struct LoopAwareBranchContext : public BranchContext {
   uint32_t loop_counter;
+};
+
+struct HistoryAwareBranchContext : public BranchContext {
   uint32_t history;
 };
 
+struct FullBranchContext : public HybridBranchContext,
+                          public ContextAwareBranchContext,
+                          public LoopAwareBranchContext,
+                          public HistoryAwareBranchContext {
+};
+
 class CovManager {
 public:
   virtual ~CovManager() {}
@@ -27,54 +45,33 @@ class CovManager {
 };
 
 class EdgeCovManager : public CovManager {
+private:
+  using BranchTargets = std::pair<bool, bool>;
+  std::unordered_map<void*, BranchTargets> branches;
+  std::shared_ptr<BranchContext> _ctx;
+
 public:
+  EdgeCovManager() { _ctx = std::make_shared<BranchContext>(); }
+
   const std::shared_ptr<BranchContext>
   add_branch(void *addr, uint32_t id, bool direction, uint32_t context, bool is_loop_header, bool is_loop_exit) override {
-    auto itr = branches.find(addr);
-    if (itr == branches.end()) {
-      auto ctx = std::make_shared<BranchContext>();
-      ctx->addr = addr;
-      ctx->direction = direction;
-      std::shared_ptr<BranchContext> tt = direction? ctx : nullptr;
-      std::shared_ptr<BranchContext> ft = direction? nullptr : ctx;
-      branches.insert({addr, {tt, ft}});
-      return ctx;
-    } else {
-      if (direction) {
-        auto ctx = itr->second.first;
-        if (!ctx) {
-          ctx = std::make_shared<BranchContext>();
-          ctx->addr = addr;
-          ctx->direction = direction;
-          itr->second.first = ctx;
-        }
-        return ctx;
-      } else {
-        auto ctx = itr->second.second;
-        if (!ctx) {
-          ctx = std::make_shared<BranchContext>();
-          ctx->addr = addr;
-          ctx->direction = direction;
-          itr->second.second = ctx;
-        }
-        return ctx;
-      }
-    }
+    auto &itr = branches[addr];
+    itr.first = direction? true : false;
+    itr.second = direction? false : true;
+    _ctx->addr = addr;
+    _ctx->direction = direction;
+    return _ctx;
   }
 
   bool is_branch_interesting(const std::shared_ptr<BranchContext> context) override {
     auto itr = branches.find(context->addr);
     assert(itr != branches.end());
     if (context->direction) {
-      return itr->second.first == nullptr;
+      return itr->second.first == false;
     } else {
-      return itr->second.second == nullptr;
+      return itr->second.second == false;
     }
   }
-
-private:
-  typedef std::pair<std::shared_ptr<BranchContext>, std::shared_ptr<BranchContext>> BranchTargets;
-  std::unordered_map<void*, BranchTargets> branches;
 };
 
 }; // namespace rgd

driver/aflpp/jit-solver.cpp

@@ -125,6 +125,36 @@ JITSolver::solve(std::shared_ptr<SearchTask> task,
       DEBUGF("generate_input offset:%zu => %u\n", offset, value);
       out_buf[offset] = value;
     }
+    // handle atoi bytes
+    if (unlikely(!task->atoi_info.empty())) {
+      // if there are atoi bytes, handle them
+      for (auto const &[offset, info] : task->atoi_info) {
+        uint64_t val = 0;
+        uint32_t length = std::get<0>(info);
+        memcpy(out_buf + offset, in_buf + offset, length); // restore??
+        for (auto i = length; i != 0; --i) {
+          DEBUGF("generate_input atoi offset:%d => %lu\n", offset + i - 1, val);
+          auto itr = task->solution.find(offset + i - 1);
+          assert(itr != task->solution.end());
+          val |= itr->second << (8 * (i - 1));
+        }
+        uint32_t base = std::get<1>(info);
+        uint32_t orig_len = std::get<2>(info);
+        DEBUGF("generate_input atoi offset:%d => %lu, base = %d, original len = %d\n",
+            offset, val, base, orig_len);
+        const char *format = nullptr;
+        switch (base) {
+          case 2: format = "%lb"; break;
+          case 8: format = "%lo"; break;
+          case 10: format = "%ld"; break;
+          case 16: format = "%lx"; break;
+          default: WARNF("unsupported base %d\n", base);
+        }
+        if (format) {
+          snprintf((char*)out_buf + offset, in_size - offset, format, val);
+        }
+      }
+    }
     num_solved++;
     return SOLVER_SAT;
   } else {

driver/aflpp/symsan.cpp

@@ -51,7 +51,7 @@ using namespace __dfsan;
 
 #define NEED_OFFLINE 0
 
-#define PRINT_STATS 1
+#define PRINT_STATS 0
 
 #define MAX_AST_SIZE 200
 
@@ -373,6 +373,51 @@ static bool do_uta_rel(dfsan_label label, rgd::AstNode *ret,
     ret->set_hash(hash);
 #if NEED_OFFLINE
     ret->set_name("memcmp");
+#endif
+    return true;
+  } else if (info->op == __dfsan::fatoi) {
+    assert(info->l1 == 0 && info->l2 >= CONST_OFFSET);
+    dfsan_label_info *src = get_label_info(info->l2);
+    assert(src->op == Load);
+    visited.insert(info->l2);
+    uint64_t offset = get_label_info(src->l1)->op1.i;
+    assert(offset < buf_size);
+    ret->set_bits(info->size);
+    ret->set_label(label);
+    ret->set_index(offset);
+    // special handling for atoi, we are introducing the result/output of
+    // atoi as fake inputs, and solve constraints over the output,
+    // once solved, we convert it back to string
+    // however, because the input is fake, we need to map it specially
+    ret->set_kind(rgd::Read);
+    auto itr = constraint->local_map.find(offset);
+    if (itr != constraint->local_map.end()) {
+      WARNF("atoi inputs should not be involved in other constraints\n");
+      return false;
+    }
+    uint32_t hash = 0;
+    uint32_t length = info->size / 8; // bits to bytes
+    // record the offset, base, and original length
+    constraint->atoi_info[offset] = std::make_tuple(length, (uint32_t)info->op1.i, (uint32_t)info->op2.i);
+    for (uint32_t i = 0; i < length; ++i, ++offset) {
+      u8 val = 0; // XXX: use 0 as initial value?
+      // because this is fake input, we always map it to a new index
+      uint32_t arg_index = (uint32_t)constraint->input_args.size();
+      constraint->inputs.insert({offset, val});
+      constraint->local_map[offset] = arg_index;
+      constraint->input_args.push_back(std::make_pair(true, 0)); // 0 is to be filled in the aggragation
+      if (i == 0) {
+        constraint->shapes[offset] = length;
+        // from solver's perspective, atoi and read are the same
+        // they both introduce a new symbolic input as arg_index
+        hash = rgd::xxhash(length * 8, rgd::Read, arg_index);
+      } else {
+        constraint->shapes[offset] = 0;
+      }
+    }
+    ret->set_hash(hash);
+#if NEED_OFFLINE
+    ret->set_name("atoi");
 #endif
     return true;
   }
@@ -1048,7 +1093,9 @@ static void scan_labels(dfsan_label label, size_t buf_size) {
       branch_to_inputs.emplace_back(input_dep_t(buf_size));
       auto &itr = branch_to_inputs[i];
       itr.set(info->op1.i); // input offset
+#if DEBUG
       assert(branch_to_inputs[i].find_first() == info->op1.i);
+#endif
       // nested cmp?
       nested_cmp_cache.push_back(0);
     } else if (info->op == __dfsan::Load) {
@@ -1062,7 +1109,10 @@ static void scan_labels(dfsan_label label, size_t buf_size) {
         // DEBUGF("adding input: %lu <- %lu\n", i, offset + n);
         itr.set(offset + n); // input offsets
       }
-      assert(branch_to_inputs[i].find_first() == offset);
+#if DEBUG
+      if (likely(info->l2 > 0))
+        assert(branch_to_inputs[i].find_first() == offset);
+#endif
       // nested cmp?
       nested_cmp_cache.push_back(0);
     } else {
@@ -1270,7 +1320,11 @@ static bool add_data_flow_constraints(bool direction, dfsan_label label,
         }
       }
       auto root = itr.find_first();
-      assert(root != input_dep_t::npos);
+      if (root == input_dep_t::npos) {
+        // not actual input dependency, skip
+        // this can happen for atoi
+        continue;
+      }
       // update uion find
       for (auto input = itr.find_next(root); input != input_dep_t::npos;
            input = itr.find_next(input)) {
@@ -1315,7 +1369,9 @@ static void handle_cond(pipe_msg &msg, const u8 *buf, size_t buf_size,
     // add the tasks to the task manager
     for (auto const& task : tasks) {
       my_mutator->task_mgr->add_task(neg_ctx, task);
+#if PRINT_STATS
       task_size_dist[task->constraints.size()] += 1;
+#endif
     }
 
     total_tasks += tasks.size();
@@ -1326,11 +1382,6 @@ static void handle_cond(pipe_msg &msg, const u8 *buf, size_t buf_size,
     // add the current branch direction as nested conditions
     add_data_flow_constraints(ctx->direction, msg.label, buf, buf_size);
   }
-
-  if (msg.flags & F_ADD_CONS) {
-    // add the current branch direction as nested conditions
-    add_data_flow_constraints(ctx->direction, msg.label, buf, buf_size);
-  }
 }
 
 static void handle_gep(gep_msg &gmsg, pipe_msg &msg) {
@@ -1493,6 +1544,7 @@ static int spawn_symsan_child(my_mutator_t *data, const u8 *buf, size_t buf_size
   if (pid == 0) {
     close(pipefds[0]); // close the read fd
     setenv("TAINT_OPTIONS", (char*)options, 1);
+    unsetenv("LD_PRELOAD"); // don't preload anything
     if (data->afl->fsrv.use_stdin) {
       close(0);
       lseek(data->out_fd, 0, SEEK_SET);

driver/aflpp/task.h

@@ -1,12 +1,14 @@
 #pragma once
 
 #include <stdint.h>
-#include <vector>
+
+#include <bitset>
 #include <map>
 #include <memory>
-#include <unordered_map>
-#include <bitset>
 #include <queue>
+#include <tuple>
+#include <unordered_map>
+#include <vector>
 
 #include "ast.h"
 #include "cov.h"
@@ -46,6 +48,8 @@ struct Constraint {
   std::unordered_map<uint32_t, uint8_t> inputs;
   // shape information about the input (e.g., 1, 2, 4, 8 bytes)
   std::unordered_map<uint32_t, uint32_t> shapes;
+  // special infomation for atoi: offset -> (base, length)
+  std::unordered_map<uint32_t, std::tuple<uint32_t, uint32_t, uint32_t>> atoi_info;
   // number of constant in the input array
   uint32_t const_num;
   // recorded comparison operands
@@ -80,6 +84,8 @@ struct SearchTask {
   std::vector<std::pair<uint32_t, uint8_t>> inputs;
   // shape information at each offset
   std::unordered_map<uint32_t, uint32_t> shapes;
+  // aggreated atoi info
+  std::unordered_map<uint32_t, std::tuple<uint32_t, uint32_t, uint32_t>> atoi_info;
   // max number of constants in the input array
   uint32_t max_const_num;
   // record constraints that use a certain input byte
@@ -158,6 +164,24 @@ struct SearchTask {
       // save the last set of consecutive input bytes
       cm->i2s_candidates.push_back({last_offset + 1 - size, size});
 
+      // process atoi
+      for (const auto& [offset, info] : constraints[i]->atoi_info) {
+        // check dependencies
+        uint32_t length = std::get<2>(info);
+        for (auto j = 0; j < length; ++j) {
+          auto ditr = cmap.find(offset + j);
+          if (ditr != cmap.end()) {
+            fprintf(stderr, "atoi bytes (%d) used in other constraints\n", offset + j);
+          }
+        }
+        auto itr = atoi_info.find(offset);
+        if (itr != atoi_info.end()) {
+          fprintf(stderr, "atoi bytes (%d) already exists\n", offset);
+          assert(info == itr->second);
+        }
+        atoi_info[offset] = info;
+      }
+
       // update the number of required constants in the input array
       if (max_const_num < constraints[i]->const_num)
         max_const_num = constraints[i]->const_num;

driver/aflpp/z3-solver.cpp

@@ -299,6 +299,37 @@ Z3Solver::solve(std::shared_ptr<SearchTask> task,
       out_size = in_size;
       z3::model m = solver_.get_model();
       extract_model(m, out_buf, out_size, task->solution);
+      if (unlikely(!task->atoi_info.empty())) {
+        // if there are atoi bytes, handle them
+        for (auto const &[offset, info] : task->atoi_info) {
+          uint64_t val = 0;
+          uint32_t length = std::get<0>(info);
+          memcpy(out_buf + offset, in_buf + offset, length); // restore?
+          for (auto i = length; i != 0; --i) {
+            DEBUGF("generate_input atoi offset:%d => %lu\n", offset + i - 1, val);
+            auto itr = task->solution.find(offset + i - 1);
+            if (itr != task->solution.end())
+              val |= itr->second << (8 * (i - 1));
+            else
+              val |= 0 << (8 * (i - 1));
+          }
+          uint32_t base = std::get<1>(info);
+          uint32_t orig_len = std::get<2>(info);
+          DEBUGF("generate_input atoi offset:%d => %lu, base = %d, original len = %d\n",
+              offset, val, base, orig_len);
+          const char *format = nullptr;
+          switch (base) {
+            case 2: format = "%lb"; break;
+            case 8: format = "%lo"; break;
+            case 10: format = "%ld"; break;
+            case 16: format = "%lx"; break;
+            default: WARNF("unsupported base %d\n", base);
+          }
+          if (format) {
+            snprintf((char*)out_buf + offset, in_size - offset, format, val);
+          }
+        }
+      }
       task->solved = true;
       return SOLVER_SAT;
     } else if (ret == z3::unsat) {