From aab1f8694a845d8c702d4053e454a9b4c2569e3c Mon Sep 17 00:00:00 2001
From: Jorge Gorbe <slack@codemaniacs.com>
Date: Tue, 25 Nov 2014 19:44:05 +0100
Subject: [PATCH] WIP: don't instrument basic block in other modules, to speed
 things up

---
 drafl.c | 67 ++++++++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 47 insertions(+), 20 deletions(-)

diff --git a/drafl.c b/drafl.c
index 8fa542e..c98912a 100644
--- a/drafl.c
+++ b/drafl.c
@@ -49,7 +49,7 @@
 #include "afl-0.68b/config.h"
 
 //#define USE_FIFO_HACK
-
+#define ENABLE_LOG
 
 typedef unsigned char u8;
 typedef unsigned int  u32;
@@ -67,19 +67,30 @@ static u8          afl_setup_failure = 0;
 static native_word afl_temp = 0;
 static native_word afl_fork_pid = 0;
 
-static FILE *logfile;
+static int ctl_fd;
+static int st_fd;
+
+static const module_data_t *moduleinfos[100];
+static int modulecount = 0;
 
+#ifdef ENABLE_LOG
+static FILE *logfile;
 #define LOG(s) do { fprintf(logfile, s); fflush(logfile); } while(0)
 #define LOG1(s, a1) do { fprintf(logfile, s, a1); fflush(logfile); } while(0)
+#else
+#define LOG(s) do {} while(0)
+#define LOG1(s, a1) do {} while(0)
+#endif
 
-int ctl_fd;
-int st_fd;
 
 static void
 event_exit(void)
 {
     LOG("event_exit()\n");
+
+#ifdef ENABLE_LOG
     fclose(logfile);
+#endif
 
     //drx_exit();
     //drwrap_exit();
@@ -89,7 +100,9 @@ event_exit(void)
 static void afl_die()
 {
     LOG("afl_die()\n");
+#ifdef ENABLE_LOG
     fclose(logfile);
+#endif
     exit(0);
 }
 
@@ -118,43 +131,43 @@ static void afl_store(native_word loc)
 static void
 afl_forkserver()
 {
-    LOG("writing hello message to st pipe\n");
+    //LOG("writing hello message to st pipe\n");
     int n = write(st_fd, &afl_temp, 4);
-    LOG1("write returned %d\n", n);
+    //LOG1("write returned %d\n", n);
     if (n != 4) {
         close(ctl_fd);
         close(st_fd);
-        LOG("returning\n");
+        //LOG("returning\n");
         return;
     }
 
     while (1)
     {
-        LOG("waiting for parent reading from ctl pipe\n");
+        //LOG("waiting for parent reading from ctl pipe\n");
         n = read(ctl_fd, &afl_temp, 4);
-        LOG1("read returned %d\n", n);
+        //LOG1("read returned %d\n", n);
         if (n != 4)
             afl_die();
 
-        LOG("forking\n");
+        //LOG("forking\n");
         afl_fork_pid = fork();
         if (afl_fork_pid < 0)
             afl_die();
 
         if (afl_fork_pid != 0) // parent
         {
-            LOG1("parent: writing PID %d to st pipe\n", (int)afl_fork_pid);
+            //LOG1("parent: writing PID %d to st pipe\n", (int)afl_fork_pid);
             // write PID
             n = write(st_fd, &afl_fork_pid, 4);
-            LOG1("write returned %d\n", n);
+            //LOG1("write returned %d\n", n);
 
             // wait for PID
-            LOG("parent: calling waitpid\n");
+            //LOG("parent: calling waitpid\n");
             int err = waitpid(afl_fork_pid, (int*)&afl_temp, WUNTRACED);
             if (err <= 0)
                 afl_die();
 
-            LOG("parent: writing wait status after waitpid\n");
+            //LOG("parent: writing wait status after waitpid\n");
             // write wait status
             write(st_fd, &afl_temp, 4);
 
@@ -162,7 +175,7 @@ afl_forkserver()
         }
         else  // child
         {
-            LOG("child: closing descriptors");
+            //LOG("child: closing descriptors\n");
             close(ctl_fd);
             close(st_fd);
             break;
@@ -221,12 +234,27 @@ static dr_emit_flags_t
 event_basic_block(void *drcontext, void *tag, instrlist_t *bb,
                   bool for_trace, bool translating)
 {
+    app_pc pc = instr_get_app_pc(instrlist_first(bb));
+
+    // skip instrumentation for basic blocks in other modules
+    for (int i=1; i<modulecount; ++i)
+    {
+        if (dr_module_contains_addr(moduleinfos[i], pc))
+            return DR_EMIT_DEFAULT;
+    }
+
     instr_t *instr = instrlist_first(bb);
     native_word random_id = random() % MAP_SIZE;
     dr_insert_clean_call(drcontext, bb, instr, afl_maybe_log, false, 1, OPND_CREATE_INTPTR(random_id));
     return DR_EMIT_DEFAULT;
 }
 
+void event_module_load(void *drcontext, const module_data_t *info, bool loaded)
+{
+    LOG1("event_module_load: %s\n", info->full_path);
+    moduleinfos[modulecount++] = dr_copy_module_data(info);
+}
+
 
 DR_EXPORT void
 dr_init(client_id_t id)
@@ -237,14 +265,13 @@ dr_init(client_id_t id)
     dr_register_exit_event(event_exit);
     dr_register_bb_event(event_basic_block);
 
+    dr_register_module_load_event(event_module_load);
 
-//#ifdef UNIX
-//    dr_register_fork_init_event(event_fork);
-//#endif
-//    drmgr_register_module_load_event(event_module_load);
-//    drmgr_register_module_unload_event(event_module_unload);
 
+
+#ifdef ENABLE_LOG
     logfile = fopen("/tmp/afl.txt", "w");
+#endif
     LOG("log file created in dr_init()\n");
 
 
-- 
2.34.1