Use BB cache prepopulation

(Ab)use the new API for basic block cache prepopulation to
pretranslate in the parent forkserver process basic blocks
executed in children. That is, implement an optimization
already used by the qemu_mode instrumentation.

CMakeLists.txt

@@ -22,6 +22,11 @@ if (NOT DynamoRIO_FOUND)
 endif(NOT DynamoRIO_FOUND)
 configure_DynamoRIO_client(afl-dr)
 
+option(BB_PREPOPULATE "Enable basic block prepopulation in the parent forkserver process" ON)
+if (BB_PREPOPULATE)
+    add_definitions(-DBB_PREPOPULATE)
+endif(BB_PREPOPULATE)
+
 # The droption extension seems not to be available through this command,
 # so using this kludge to add the include path
 use_DynamoRIO_extension(afl-dr drutil)

afl-dr.cpp

@@ -76,6 +76,10 @@ static dr_emit_flags_t event_basic_block(void *drcontext, void *tag, instrlist_t
                                          bool for_trace, bool translating) {
     app_pc pc = dr_fragment_app_pc(tag);
 
+    if (!translating) {
+        trace_bb_instrumentation(pc, for_trace);
+    }
+
     if (!opt_instrument_everything.get_value() && !dr_module_contains_addr(main_module, pc)) {
         return DR_EMIT_DEFAULT;
     }

afl-dr.h

@@ -10,5 +10,6 @@
     }
 
 void start_forkserver(void);
+void trace_bb_instrumentation(app_pc pc, bool for_trace);
 
 #endif

afl-forksrv.cpp

@@ -4,6 +4,7 @@
 #include <unistd.h>
 #include <sys/wait.h>
 #include <sys/types.h>
+#include <fcntl.h>
 
 #include <droption.h>
 
@@ -16,6 +17,14 @@ static droption_t<bool> opt_private_fork(DROPTION_SCOPE_CLIENT, "private-fork",
                                          "Use fork function from the private libc",
                                          "Use fork function from the private libc");
 
+#ifdef BB_PREPOPULATE
+static droption_t<bool> opt_disable_prepop(DROPTION_SCOPE_CLIENT, "no-prepop", false,
+                                           "Disable basic block cache prepopulation",
+                                           "Disable qemu_mode-like prepopulation of basic block cache in the parent forkserver process");
+
+static int prepop_fd = -1;
+#endif
+
 void start_forkserver() {
     // For references, see https://lcamtuf.blogspot.ru/2014/10/fuzzing-binaries-without-execve.html
     // and __afl_start_forkserver in llvm_mode/afl-llvm-rt.o.c from AFL sources
@@ -52,12 +61,44 @@ void start_forkserver() {
         fork_ptr = fork;
     }
 
+#ifdef BB_PREPOPULATE
+    int prepop_pipe[2];
+    if (!opt_disable_prepop.get_value()) {
+        EXIT_IF_FAILED(pipe2(prepop_pipe, O_NONBLOCK) == 0, "Cannot create the bb-prepop pipe.\n", 1);
+    }
+#endif
+
     while (true) {
+#ifdef BB_PREPOPULATE
+        // Pretranslate basic blocks inside the parent forkserver process,
+        // like qemu_mode already does.
+        if (!opt_disable_prepop.get_value()) {
+            app_pc pcs[500];
+            while (true) {
+                int res = read(prepop_pipe[0], pcs, sizeof(pcs));
+                if (res <= 0)
+                    break;
+
+                // At the time of writing, this is kind of API abuse, see:
+                // https://github.com/DynamoRIO/dynamorio/issues/2463
+                // https://github.com/DynamoRIO/dynamorio/pull/2505
+                void *drcontext = dr_get_current_drcontext(); // Save before dr_prepopulate_cache()
+                dr_prepopulate_cache(pcs, res / sizeof(pcs[0]));
+                dr_switch_to_dr_state_ex(drcontext, DR_STATE_GO_NATIVE);
+            }
+        }
+#endif
         EXIT_IF_FAILED(read(FROM_FUZZER_FD, &was_killed, 4) == 4, "Incorrect spawn command from fuzzer.\n", 1)
         int child_pid = fork_ptr();
         EXIT_IF_FAILED(child_pid >= 0, "Cannot fork.\n", 1)
 
         if (child_pid == 0) {
+#ifdef BB_PREPOPULATE
+            if (!opt_disable_prepop.get_value()) {
+                close(prepop_pipe[0]);
+                prepop_fd = prepop_pipe[1];
+            }
+#endif
             close(TO_FUZZER_FD);
             close(FROM_FUZZER_FD);
             return;
@@ -69,3 +110,11 @@ void start_forkserver() {
         }
     }
 }
+
+void trace_bb_instrumentation(app_pc pc, bool for_trace) {
+#ifdef BB_PREPOPULATE
+    if (!for_trace && prepop_fd > 0) {
+        EXIT_IF_FAILED(write(prepop_fd, &pc, sizeof(pc)) == sizeof(pc), "Cannot write pc for prepop.\n", 1);
+    }
+#endif
+}