added my changes

1 files changed, 67 insertions, 52 deletions

diff --git a/libAflDyninst.cpp b/libAflDyninst.cpp
index 514144f..ef9d046 100644
--- a/libAflDyninst.cpp
+++ b/libAflDyninst.cpp
@@ -14,70 +14,85 @@
 
 using namespace std;
 
-
-static u8* trace_bits;
-static s32 shm_id;                    /* ID of the SHM region             */
+static u8 *trace_bits;
+static s32 shm_id;
 static int __afl_temp_data;
 static pid_t __afl_fork_pid;
 static unsigned short prev_id;
+static long saved_di;
+register long rdi asm("di");    // the warning is fine - we need the warning because of a bug in dyninst
 
+void initAflForkServer() {
+  char *shm_env_var = getenv(SHM_ENV_VAR);
 
-void initAflForkServer()
-{
-    char *shm_env_var = getenv(SHM_ENV_VAR);
-    if(!shm_env_var) {
-        printf("Error getting shm\n");
-        return;
-    }
-    shm_id = atoi(shm_env_var);
-    trace_bits = (u8*)shmat(shm_id, NULL, 0);
-    if(trace_bits == (u8*)-1) {
-        perror("shmat");
-        return;
-    }
+  if (!shm_env_var) {
+    printf("Error getting shm\n");
+    return;
+  }
+  shm_id = atoi(shm_env_var);
+  trace_bits = (u8 *) shmat(shm_id, NULL, 0);
+  if (trace_bits == (u8 *) - 1) {
+    perror("shmat");
+    return;
+  }
+  // enter fork() server thyme!
+  int n = write(FORKSRV_FD + 1, &__afl_temp_data, 4);
 
-    // enter fork() server thyme!
-    int n = write(FORKSRV_FD+1, &__afl_temp_data,4);
-    if( n!=4 ) {
-        printf("Error writting fork server\n");
-        return;
+  if (n != 4) {
+    printf("Error writting fork server\n");
+    return;
+  }
+  while (1) {
+    n = read(FORKSRV_FD, &__afl_temp_data, 4);
+    if (n != 4) {
+      printf("Error reading fork server %x\n", __afl_temp_data);
+      return;
     }
-    while(1) {
-        n = read(FORKSRV_FD,&__afl_temp_data,4);
-        if(n != 4) {
-            printf("Error reading fork server %x\n",__afl_temp_data);
-            return;
-        }
 
-        __afl_fork_pid = fork();
-        if(__afl_fork_pid < 0) {
-            printf("Error on fork()\n");
-            return;
-        }
-        if(__afl_fork_pid == 0) {
-            close(FORKSRV_FD);
-            close(FORKSRV_FD+1);
-            break;
-        } else {
-            // parrent stuff
-            n = write(FORKSRV_FD+1,&__afl_fork_pid, 4);
-            pid_t temp_pid = waitpid(__afl_fork_pid,&__afl_temp_data,2);
-            if(temp_pid == 0) {
-                return;
-            }
-            n = write(FORKSRV_FD+1,&__afl_temp_data,4);
-        }
+    __afl_fork_pid = fork();
+    if (__afl_fork_pid < 0) {
+      printf("Error on fork()\n");
+      return;
+    }
+    if (__afl_fork_pid == 0) {
+      close(FORKSRV_FD);
+      close(FORKSRV_FD + 1);
+      break;
+    } else {
+      // parrent stuff
+      n = write(FORKSRV_FD + 1, &__afl_fork_pid, 4);
+      pid_t temp_pid = waitpid(__afl_fork_pid, &__afl_temp_data, 2);
 
+      if (temp_pid == 0) {
+        return;
+      }
+      n = write(FORKSRV_FD + 1, &__afl_temp_data, 4);
     }
+  }
+}
 
+// Should be called on basic block entry
+void bbCallback(unsigned short id) {
+  if (trace_bits) {
+    trace_bits[prev_id ^ id]++;
+    prev_id = id >> 1;
+  }
 }
 
+void save_rdi() {
+  saved_di = rdi;
+/*
+  asm("pop %rax"); // take care of rip
+  asm("push %rdi");
+  asm("push %rax");
+*/
+}
 
-// Should be called on basic block entry
-void bbCallback(unsigned short id)
-{
-    if(trace_bits) {
-        trace_bits[prev_id ^ id]++;
-        prev_id = id >> 1;
-    }
+void restore_rdi() {
+  rdi = saved_di;
+/*
+  asm("pop %rax"); // take care of rip
+  asm("pop %rdi");
+  asm("push %rax");
+*/
 }