fixed map location support for LTO

5 files changed, 135 insertions, 23 deletions

diff --git a/libdislocator/libdislocator.so.c b/libdislocator/libdislocator.so.c
index 8098de91..6c8916d6 100644
--- a/libdislocator/libdislocator.so.c
+++ b/libdislocator/libdislocator.so.c
@@ -466,11 +466,13 @@ void *reallocarray(void *ptr, size_t elem_len, size_t elem_cnt) {
 
 #if !defined(__ANDROID__)
 size_t malloc_usable_size(void *ptr) {
+
 #else
 size_t malloc_usable_size(const void *ptr) {
+
 #endif
 
-   return ptr ? PTR_L(ptr) : 0;
+  return ptr ? PTR_L(ptr) : 0;
 
 }
 
@@ -498,3 +500,4 @@ __attribute__((constructor)) void __dislocator_init(void) {
   align_allocations = !!getenv("AFL_ALIGNED_ALLOC");
 
 }
+
diff --git a/llvm_mode/afl-llvm-lto-instrumentation.so.cc b/llvm_mode/afl-llvm-lto-instrumentation.so.cc
index eae60ccd..dbd4a26e 100644
--- a/llvm_mode/afl-llvm-lto-instrumentation.so.cc
+++ b/llvm_mode/afl-llvm-lto-instrumentation.so.cc
@@ -31,6 +31,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <unistd.h>
+#include <string.h>
 
 #include <list>
 #include <string>
@@ -154,6 +155,7 @@ class AFLLTOPass : public ModulePass {
  protected:
   int      afl_global_id = 1, debug = 0, autodictionary = 0;
   uint32_t be_quiet = 0, inst_blocks = 0, inst_funcs = 0, total_instr = 0;
+  uint64_t map_addr = 0x10000;
 
 };
 
@@ -165,9 +167,11 @@ bool AFLLTOPass::runOnModule(Module &M) {
   std::vector<std::string>         dictionary;
   std::vector<CallInst *>          calls;
   DenseMap<Value *, std::string *> valueMap;
+  char *                           ptr;
 
   IntegerType *Int8Ty = IntegerType::getInt8Ty(C);
   IntegerType *Int32Ty = IntegerType::getInt32Ty(C);
+  IntegerType *Int64Ty = IntegerType::getInt64Ty(C);
 
   if (getenv("AFL_DEBUG")) debug = 1;
 
@@ -186,12 +190,63 @@ bool AFLLTOPass::runOnModule(Module &M) {
       getenv("AFL_LLVM_LTO_AUTODICTIONARY"))
     autodictionary = 1;
 
+  if (getenv("AFL_LLVM_MAP_DYNAMIC")) map_addr = 0;
+
+  if ((ptr = getenv("AFL_LLVM_MAP_ADDR"))) {
+
+    uint64_t val;
+    if (!*ptr || !strcmp(ptr, "0") || !strcmp(ptr, "0x0")) {
+
+      map_addr = 0;
+
+    } else if (map_addr == 0) {
+
+      FATAL(
+          "AFL_LLVM_MAP_ADDR and AFL_LLVM_MAP_DYNAMIC cannot be used together");
+
+    } else if (strncmp(ptr, "0x", 2) != 0) {
+
+      map_addr = 0x10000;  // the default
+
+    } else {
+
+      val = strtoull(ptr, NULL, 16);
+      if (val < 0x100 || val > 0xffffffff00000000) {
+
+        FATAL(
+            "AFL_LLVM_MAP_ADDR must be a value between 0x100 and "
+            "0xffffffff00000000");
+
+      }
+
+      map_addr = val;
+
+    }
+
+  }
+
+  if (debug) { fprintf(stderr, "map address is %lu\n", map_addr); }
+
   /* Get globals for the SHM region and the previous location. Note that
      __afl_prev_loc is thread-local. */
 
-  GlobalVariable *AFLMapPtr =
-      new GlobalVariable(M, PointerType::get(Int8Ty, 0), false,
-                         GlobalValue::ExternalLinkage, 0, "__afl_area_ptr");
+  GlobalVariable *AFLMapPtr = NULL;
+  ;
+  Value *MapPtrFixed = NULL;
+
+  if (!map_addr) {
+
+    AFLMapPtr =
+        new GlobalVariable(M, PointerType::get(Int8Ty, 0), false,
+                           GlobalValue::ExternalLinkage, 0, "__afl_area_ptr");
+
+  } else {
+
+    ConstantInt *MapAddr = ConstantInt::get(Int64Ty, map_addr);
+    MapPtrFixed =
+        ConstantExpr::getIntToPtr(MapAddr, PointerType::getUnqual(Int8Ty));
+
+  }
 
   ConstantInt *Zero = ConstantInt::get(Int8Ty, 0);
   ConstantInt *One = ConstantInt::get(Int8Ty, 1);
@@ -202,7 +257,7 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
   for (auto &F : M) {
 
-    //fprintf(stderr, "DEBUG: Function %s\n", F.getName().str().c_str());
+    // fprintf(stderr, "DEBUG: Function %s\n", F.getName().str().c_str());
 
     if (F.size() < 2) continue;
     if (isBlacklisted(&F)) continue;
@@ -581,10 +636,20 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
           /* Load SHM pointer */
 
-          LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr);
-          MapPtr->setMetadata(M.getMDKindID("nosanitize"),
-                              MDNode::get(C, None));
-          Value *MapPtrIdx = IRB.CreateGEP(MapPtr, CurLoc);
+          Value *MapPtrIdx;
+
+          if (map_addr) {
+
+            MapPtrIdx = IRB.CreateGEP(MapPtrFixed, CurLoc);
+
+          } else {
+
+            LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr);
+            MapPtr->setMetadata(M.getMDKindID("nosanitize"),
+                                MDNode::get(C, None));
+            MapPtrIdx = IRB.CreateGEP(MapPtr, CurLoc);
+
+          }
 
           /* Update bitmap */
 
@@ -629,7 +694,7 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
   }
 
-  if (getenv("AFL_LLVM_LTO_DONTWRITEID") == NULL || dictionary.size()) {
+  if (!getenv("AFL_LLVM_LTO_DONTWRITEID") || dictionary.size() || map_addr) {
 
     // yes we could create our own function, insert it into ctors ...
     // but this would be a pain in the butt ... so we use afl-llvm-rt-lto.o
@@ -658,6 +723,18 @@ bool AFLLTOPass::runOnModule(Module &M) {
     BasicBlock::iterator IP = bb->getFirstInsertionPt();
     IRBuilder<>          IRB(&(*IP));
 
+    if (map_addr) {
+
+      GlobalVariable *AFLMapAddrFixed = new GlobalVariable(
+          M, Int64Ty, true, GlobalValue::ExternalLinkage, 0, "__afl_map_addr",
+          0, GlobalVariable::GeneralDynamicTLSModel, 0, false);
+      ConstantInt *MapAddr = ConstantInt::get(Int64Ty, map_addr);
+      StoreInst *  StoreMapAddr = IRB.CreateStore(MapAddr, AFLMapAddrFixed);
+      StoreMapAddr->setMetadata(M.getMDKindID("nosanitize"),
+                                MDNode::get(C, None));
+
+    }
+
     if (getenv("AFL_LLVM_LTO_DONTWRITEID") == NULL) {
 
       uint32_t write_loc = afl_global_id;
diff --git a/llvm_mode/afl-llvm-rt.o.c b/llvm_mode/afl-llvm-rt.o.c
index f286e66a..6da41192 100644
--- a/llvm_mode/afl-llvm-rt.o.c
+++ b/llvm_mode/afl-llvm-rt.o.c
@@ -52,6 +52,10 @@
 
 #define CONST_PRIO 5
 
+//#ifndef MAP_FIXED_NOREPLACE
+//#define MAP_FIXED_NOREPLACE MAP_FIXED
+//#endif
+
 #include <sys/mman.h>
 #include <fcntl.h>
 
@@ -69,12 +73,14 @@ u32        __afl_final_loc;
 u32        __afl_prev_ctx;
 u32        __afl_cmp_counter;
 u32        __afl_dictionary_len;
+u64        __afl_map_addr;
 #else
 __thread PREV_LOC_T __afl_prev_loc[NGRAM_SIZE_MAX];
 __thread u32        __afl_final_loc;
 __thread u32        __afl_prev_ctx;
 __thread u32        __afl_cmp_counter;
 __thread u32        __afl_dictionary_len;
+__thread u64        __afl_map_addr;
 #endif
 
 struct cmp_map *__afl_cmp_map;
@@ -87,7 +93,11 @@ static u8 is_persistent;
 
 static void __afl_map_shm(void) {
 
-  u8 *id_str = getenv(SHM_ENV_VAR);
+  u8 *         id_str = getenv(SHM_ENV_VAR);
+  unsigned int map_size = MAP_SIZE;
+
+  if (__afl_final_loc > 1 && __afl_final_loc < MAP_SIZE)
+    map_size = __afl_final_loc;
 
   /* If we're running under AFL, attach to the appropriate region, replacing the
      early-stage __afl_area_initial region that is needed to allow some really
@@ -99,10 +109,6 @@ static void __afl_map_shm(void) {
     const char *   shm_file_path = id_str;
     int            shm_fd = -1;
     unsigned char *shm_base = NULL;
-    unsigned int   map_size = MAP_SIZE
-
-        if (__afl_final_loc > 1 && __afl_final_loc < MAP_SIZE) map_size =
-            __afl_final_loc;
 
     /* create the shared memory segment as if it was a file */
     shm_fd = shm_open(shm_file_path, O_RDWR, 0600);
@@ -114,7 +120,18 @@ static void __afl_map_shm(void) {
     }
 
     /* map the shared memory segment to the address space of the process */
-    shm_base = mmap(0, map_size, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0);
+    if (__afl_map_addr) {
+
+      shm_base = mmap((void *)__afl_map_addr, map_size, PROT_READ | PROT_WRITE,
+                      MAP_FIXED | MAP_SHARED, shm_fd, 0);
+
+    } else {
+
+      shm_base =
+          mmap(0, map_size, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0);
+
+    }
+
     if (shm_base == MAP_FAILED) {
 
       close(shm_fd);
@@ -129,7 +146,8 @@ static void __afl_map_shm(void) {
 #else
     u32 shm_id = atoi(id_str);
 
-    __afl_area_ptr = shmat(shm_id, NULL, 0);
+    __afl_area_ptr = shmat(shm_id, (void *)__afl_map_addr, 0);
+
 #endif
 
     /* Whooooops. */
@@ -141,6 +159,19 @@ static void __afl_map_shm(void) {
 
     __afl_area_ptr[0] = 1;
 
+  } else if (__afl_map_addr) {
+
+    __afl_area_ptr =
+        mmap((void *)__afl_map_addr, map_size, PROT_READ | PROT_WRITE,
+             MAP_FIXED | MAP_SHARED | MAP_ANONYMOUS, -1, 0);
+    if (__afl_area_ptr == MAP_FAILED) {
+
+      fprintf(stderr, "can not aquire mmap for address %p\n",
+              (void *)__afl_map_addr);
+      exit(1);
+
+    }
+
   }
 
   id_str = getenv(CMPLOG_SHM_ENV_VAR);
diff --git a/src/afl-common.c b/src/afl-common.c
index 1dae8509..8d444876 100644
--- a/src/afl-common.c
+++ b/src/afl-common.c
@@ -67,10 +67,11 @@ char *afl_environment_variables[] = {
     "AFL_LLVM_INSTRIM_SKIPSINGLEBLOCK", "AFL_LLVM_LAF_SPLIT_COMPARES",
     "AFL_LLVM_LAF_SPLIT_COMPARES_BITW", "AFL_LLVM_LAF_SPLIT_FLOATS",
     "AFL_LLVM_LAF_SPLIT_SWITCHES", "AFL_LLVM_LAF_TRANSFORM_COMPARES",
-    "AFL_LLVM_NGRAM_SIZE", "AFL_NGRAM_SIZE", "AFL_LLVM_NOT_ZERO",
-    "AFL_LLVM_WHITELIST", "AFL_NO_AFFINITY", "AFL_LLVM_LTO_STARTID",
-    "AFL_LLVM_LTO_DONTWRITEID", "AFL_NO_ARITH", "AFL_NO_BUILTIN",
-    "AFL_NO_CPU_RED", "AFL_NO_FORKSRV", "AFL_NO_UI", "AFL_NO_PYTHON",
+    "AFL_LLVM_MAP_ADDR", "AFL_LLVM_MAP_DYNAMIC", "AFL_LLVM_NGRAM_SIZE",
+    "AFL_NGRAM_SIZE", "AFL_LLVM_NOT_ZERO", "AFL_LLVM_WHITELIST",
+    "AFL_NO_AFFINITY", "AFL_LLVM_LTO_STARTID", "AFL_LLVM_LTO_DONTWRITEID",
+    "AFL_NO_ARITH", "AFL_NO_BUILTIN", "AFL_NO_CPU_RED", "AFL_NO_FORKSRV",
+    "AFL_NO_UI", "AFL_NO_PYTHON",
     "AFL_NO_X86",  // not really an env but we dont want to warn on it
     "AFL_MAP_SIZE", "AFL_MAPSIZE", "AFL_PATH", "AFL_PERFORMANCE_FILE",
     //"AFL_PERSISTENT", // not implemented anymore, so warn additionally
diff --git a/src/afl-sharedmem.c b/src/afl-sharedmem.c
index e46c6f50..e024eb18 100644
--- a/src/afl-sharedmem.c
+++ b/src/afl-sharedmem.c
@@ -127,8 +127,8 @@ u8 *afl_shm_init(sharedmem_t *shm, size_t map_size, unsigned char dumb_mode) {
   }
 
   /* map the shared memory segment to the address space of the process */
-  shm->map = mmap(0, map_size, PROT_READ | PROT_WRITE, MAP_SHARED,
-                  shm->g_shm_fd, 0);
+  shm->map =
+      mmap(0, map_size, PROT_READ | PROT_WRITE, MAP_SHARED, shm->g_shm_fd, 0);
   if (shm->map == MAP_FAILED) {
 
     close(shm->g_shm_fd);