1 files changed, 77 insertions, 15 deletions

diff --git a/utils/aflpp_driver/aflpp_driver.c b/utils/aflpp_driver/aflpp_driver.c
index ff42f3b9..c648674a 100644
--- a/utils/aflpp_driver/aflpp_driver.c
+++ b/utils/aflpp_driver/aflpp_driver.c
@@ -45,6 +45,9 @@ $AFL_HOME/afl-fuzz -i IN -o OUT ./a.out
 #include <sys/stat.h>
 #include <fcntl.h>
 #include <sys/mman.h>
+#ifndef __HAIKU__
+  #include <sys/syscall.h>
+#endif
 
 #include "config.h"
 #include "types.h"
@@ -62,6 +65,27 @@ extern unsigned char *__afl_fuzz_ptr;
 int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size);
 __attribute__((weak)) int LLVMFuzzerInitialize(int *argc, char ***argv);
 
+// Default nop ASan hooks for manual posisoning when not linking the ASan
+// runtime
+// https://github.com/google/sanitizers/wiki/AddressSanitizerManualPoisoning
+__attribute__((weak)) void __asan_poison_memory_region(
+    void const volatile *addr, size_t size) {
+
+  (void)addr;
+  (void)size;
+
+}
+
+__attribute__((weak)) void __asan_unpoison_memory_region(
+    void const volatile *addr, size_t size) {
+
+  (void)addr;
+  (void)size;
+
+}
+
+__attribute__((weak)) void *__asan_region_is_poisoned(void *beg, size_t size);
+
 // Notify AFL about persistent mode.
 static volatile char AFL_PERSISTENT[] = "##SIG_AFL_PERSISTENT##";
 int                  __afl_persistent_loop(unsigned int);
@@ -175,6 +199,9 @@ static int ExecuteFilesOnyByOne(int argc, char **argv) {
 
   unsigned char *buf = (unsigned char *)malloc(MAX_FILE);
 
+  __asan_poison_memory_region(buf, MAX_FILE);
+  ssize_t prev_length = 0;
+
   for (int i = 1; i < argc; i++) {
 
     int fd = 0;
@@ -183,10 +210,26 @@ static int ExecuteFilesOnyByOne(int argc, char **argv) {
 
     if (fd == -1) { continue; }
 
-    ssize_t length = read(fd, buf, MAX_FILE);
+#ifndef __HAIKU__
+    ssize_t length = syscall(SYS_read, fd, buf, MAX_FILE);
+#else
+    ssize_t length = _kern_read(fd, buf, MAX_FILE);
+#endif  // HAIKU
 
     if (length > 0) {
 
+      if (length < prev_length) {
+
+        __asan_poison_memory_region(buf + length, prev_length - length);
+
+      } else {
+
+        __asan_unpoison_memory_region(buf + prev_length, length - prev_length);
+
+      }
+
+      prev_length = length;
+
       printf("Reading %zu bytes from %s\n", length, argv[i]);
       LLVMFuzzerTestOneInput(buf, length);
       printf("Execution successful.\n");
@@ -284,29 +327,48 @@ int main(int argc, char **argv) {
   // on the first execution of LLVMFuzzerTestOneInput is ignored.
   LLVMFuzzerTestOneInput(dummy_input, 1);
 
-  int num_runs = 0;
-  while (__afl_persistent_loop(N)) {
+  __asan_poison_memory_region(__afl_fuzz_ptr, MAX_FILE);
+  size_t prev_length = 0;
 
-#ifdef _DEBUG
-    fprintf(stderr, "CLIENT crc: %016llx len: %u\n",
-            hash64(__afl_fuzz_ptr, *__afl_fuzz_len, 0xa5b35705),
-            *__afl_fuzz_len);
-    fprintf(stderr, "RECV:");
-    for (int i = 0; i < *__afl_fuzz_len; i++)
-      fprintf(stderr, "%02x", __afl_fuzz_ptr[i]);
-    fprintf(stderr, "\n");
-#endif
+  // for speed only insert asan functions if the target is linked with asan
+  if (__asan_region_is_poisoned) {
+
+    while (__afl_persistent_loop(N)) {
+
+      size_t length = *__afl_fuzz_len;
+
+      if (likely(length)) {
+
+        if (length < prev_length) {
+
+          __asan_poison_memory_region(__afl_fuzz_ptr + length,
+                                      prev_length - length);
+
+        } else if (length > prev_length) {
 
-    if (*__afl_fuzz_len) {
+          __asan_unpoison_memory_region(__afl_fuzz_ptr + prev_length,
+                                        length - prev_length);
+
+        }
+
+        prev_length = length;
+        LLVMFuzzerTestOneInput(__afl_fuzz_ptr, length);
+
+      }
+
+    }
+
+  } else {
+
+    while (__afl_persistent_loop(N)) {
 
-      num_runs++;
       LLVMFuzzerTestOneInput(__afl_fuzz_ptr, *__afl_fuzz_len);
 
     }
 
   }
 
-  printf("%s: successfully executed %d input(s)\n", argv[0], num_runs);
+  return 0;
 
 }