aarch64 support for afl-untracer

3 files changed, 83 insertions, 9 deletions

diff --git a/examples/afl_untracer/Makefile b/examples/afl_untracer/Makefile
index 5c525877..ad67e8bc 100644
--- a/examples/afl_untracer/Makefile
+++ b/examples/afl_untracer/Makefile
@@ -1,10 +1,16 @@
+ifdef DEBUG
+  OPT=-O0
+else
+  OPT=-O3
+endif
+
 all:	afl-untracer libtestinstr.so
 
 afl-untracer:	afl-untracer.c
-	$(CC) -I../../include -g -o afl-untracer afl-untracer.c -ldl -pthread
+	$(CC) $(OPT) -I../../include -g -o afl-untracer afl-untracer.c -ldl -pthread
 
 libtestinstr.so:	libtestinstr.c
-	$(CC) -fPIC -o libtestinstr.so -shared libtestinstr.c
+	$(CC) -g -O0 -fPIC -o libtestinstr.so -shared libtestinstr.c
 
 clean:
 	rm -f afl-untracer libtestinstr.so *~ core
diff --git a/examples/afl_untracer/README.md b/examples/afl_untracer/README.md
index 4ff96423..8e24c2a4 100644
--- a/examples/afl_untracer/README.md
+++ b/examples/afl_untracer/README.md
@@ -14,6 +14,21 @@ To generate the `patches.txt` file for your target library use the
 `ida_get_patchpoints.py` script for IDA Pro or
 `ghidra_get_patchpoints.java` for Ghidra.
 
+The patches.txt file has to pointed to by `AFL_UNTRACER_FILE`.
+
+Example (after modfying afl-untracer.c to your needs, compiling and creating
+patches.txt):
+```
+AFL_UNTRACER_FILE=./patches.txt afl-fuzz -i in -o out -- ./afl-untracer
+```
+
+To testing/debugging you can try:
+```
+make DEBUG=1
+AFL_UNTRACER_FILE=./patches.txt AFL_DEBUG=1 gdb ./afl-untracer
+```
+and then you can easily set breakpoints to "breakpoint" and "fuzz".
+
 This idea is based on [UnTracer](https://github.com/FoRTE-Research/UnTracer-AFL)
 and modified by [Trapfuzz](https://github.com/googleprojectzero/p0tools/tree/master/TrapFuzz).
 This implementation is slower because the traps are not patched out with each
diff --git a/examples/afl_untracer/afl-untracer.c b/examples/afl_untracer/afl-untracer.c
index d319b530..ed1a8239 100644
--- a/examples/afl_untracer/afl-untracer.c
+++ b/examples/afl_untracer/afl-untracer.c
@@ -156,7 +156,7 @@ void read_library_information() {
         liblist[liblist_cnt].addr_end = strtoull(m, NULL, 16);
         if (debug)
           fprintf(
-              stderr, "%s:%x (%lx-%lx)\n", liblist[liblist_cnt].name,
+              stderr, "%s:%llx (%llx-%llx)\n", liblist[liblist_cnt].name,
               liblist[liblist_cnt].addr_end - liblist[liblist_cnt].addr_start,
               liblist[liblist_cnt].addr_start,
               liblist[liblist_cnt].addr_end - 1);
@@ -276,6 +276,17 @@ library_list_t *find_library(char *name) {
 
 }
 
+/* for having an easy breakpoint after load the shared library */
+// this seems to work for clang too. nice :) requires gcc 4.4+
+#pragma GCC push_options
+#pragma GCC optimize ("O0")
+void breakpoint() {
+
+  if (debug) fprintf(stderr, "Breakpoint function \"breakpoint\" reached.\n");
+
+}
+#pragma GCC pop_options
+
 /* Error reporting to forkserver controller */
 
 void send_forkserver_error(int error) {
@@ -433,10 +444,17 @@ static void __afl_end_testcase(int status) {
 
 }
 
+#ifdef __aarch64__
+#define SHADOW(addr)                                     \
+  ((uint64_t *)(((uintptr_t)addr & 0xfffffffffffffff8) - \
+                MEMORY_MAP_DECREMENT -                   \
+                ((uintptr_t)addr & 0x7) * 0x10000000000))
+#else
 #define SHADOW(addr)                                     \
   ((uint32_t *)(((uintptr_t)addr & 0xfffffffffffffffc) - \
                 MEMORY_MAP_DECREMENT -                   \
                 ((uintptr_t)addr & 0x3) * 0x10000000000))
+#endif
 
 void setup_trap_instrumentation() {
 
@@ -453,7 +471,11 @@ void setup_trap_instrumentation() {
   if (!patches) FATAL("Couldn't open AFL_UNTRACER_FILE file %s", filename);
 
   // Index into the coverage bitmap for the current trap instruction.
-  int bitmap_index = 0;
+#ifdef __aarch64__
+  uint64_t bitmap_index = 0;
+#else
+  uint32_t bitmap_index = 0;
+#endif
 
   while ((nread = getline(&line, &len, patches)) != -1) {
 
@@ -486,7 +508,11 @@ void setup_trap_instrumentation() {
         FATAL("Failed to mprotect library %s writable", line);
 
       // Create shadow memory.
+#ifdef __aarch64__
+      for (int i = 0; i < 8; i++) {
+#else
       for (int i = 0; i < 4; i++) {
+#endif
 
         void *shadow_addr = SHADOW(lib_addr + i);
         void *shadow = mmap(shadow_addr, lib_size, PROT_READ | PROT_WRITE,
@@ -513,12 +539,17 @@ void setup_trap_instrumentation() {
     if (bitmap_index >= __afl_map_size)
       FATAL("Too many basic blocks to instrument");
 
-    uint32_t *shadow = SHADOW(lib_addr + offset);
+#ifdef __arch64__
+    uint64_t 
+#else
+    uint32_t 
+#endif
+    *shadow = SHADOW(lib_addr + offset);
     if (*shadow != 0) continue;  // skip duplicates
 
       // Make lookup entry in shadow memory.
 
-#if ((defined(__APPLE__) && defined(__LP64__)) || defined(__x86_64__))
+#if ((defined(__APPLE__) && defined(__LP64__)) || defined(__x86_64__) || defined(__i386__))
 
     // this is for Intel x64
 
@@ -531,10 +562,24 @@ void setup_trap_instrumentation() {
               lib_addr, offset, lib_addr + offset, orig_byte, shadow,
               bitmap_index, *shadow);
 
+#elif defined(__aarch64__)
+
+    // this is for aarch64
+
+    uint32_t *patch_bytes = (uint32_t*)(lib_addr + offset);
+    uint32_t orig_bytes = *patch_bytes;
+    *shadow = (bitmap_index << 32) | orig_bytes;
+    *patch_bytes = 0xd4200000; // replace instruction with debug trap
+    if (debug)
+      fprintf(stderr,
+              "Patch entry: %p[%x] = %p = %02x -> SHADOW(%p) #%d -> %016x\n",
+              lib_addr, offset, lib_addr + offset, orig_bytes, shadow,
+              bitmap_index, *shadow);
+
 #else
       // this will be ARM and AARCH64
       // for ARM we will need to identify if the code is in thumb or ARM
-#error "non x86_64 not supported yet"
+#error "non x86_64/aarch64 not supported yet"
       //__arm__:
       // linux thumb: 0xde01
       // linux arm: 0xe7f001f0
@@ -573,8 +618,15 @@ static void sigtrap_handler(int signum, siginfo_t *si, void *context) {
   ctx->uc_mcontext->__ss.__rip -= 1;
   addr = ctx->uc_mcontext->__ss.__rip;
 #elif defined(__linux__)
-  ctx->uc_mcontext.gregs[REG_RIP] -= 1;
-  addr = ctx->uc_mcontext.gregs[REG_RIP];
+#if defined(__x86_64__) || defined(__i386__)
+  ctx->uc_mcontext.mc_rip -= 1;
+  addr = ctx->uc_mcontext.mc_rip;
+#elif defined(__aarch64__)
+  ctx->uc_mcontext.pc -= 4;
+  addr = ctx->uc_mcontext.pc;
+#else
+#error "Unsupported processor"
+#endif
 #elif defined(__FreeBSD__) && defined(__LP64__)
   ctx->uc_mcontext.mc_rip -= 1;
   addr = ctx->uc_mcontext.mc_rip;
@@ -642,6 +694,7 @@ int main(int argc, char *argv[]) {
   // END STEP 2
 
   /* setup instrumentation, shared memory and forkserver */
+  breakpoint();
   read_library_information();
   setup_trap_instrumentation();
   __afl_map_shm();