Little additions from feedback

3 files changed, 23 insertions, 1 deletions

diff --git a/libdislocator/Makefile b/libdislocator/Makefile
index 05ba26b3..216d2862 100644
--- a/libdislocator/Makefile
+++ b/libdislocator/Makefile
@@ -21,6 +21,10 @@ VERSION     = $(shell grep '^\#define VERSION ' ../config.h | cut -d '"' -f2)
 CFLAGS      ?= -O3 -funroll-loops -I ../include/
 CFLAGS      += -Wall -D_FORTIFY_SOURCE=2 -g -Wno-pointer-sign
 
+ifdef USEHUGEPAGE
+	CFLAGS += -DUSEHUGEPAGE
+endif
+
 all: libdislocator.so
 
 libdislocator.so: libdislocator.so.c ../config.h
diff --git a/libdislocator/README.dislocator.md b/libdislocator/README.dislocator.md
index 5d5a1464..77626901 100644
--- a/libdislocator/README.dislocator.md
+++ b/libdislocator/README.dislocator.md
@@ -25,6 +25,8 @@ heap-related security bugs in several ways:
   - It checks for calloc() overflows and can cause soft or hard failures
     of alloc requests past a configurable memory limit (AFL_LD_LIMIT_MB,
     AFL_LD_HARD_FAIL).
+  - Optionally, in platforms supporting it, huge pages can be used by passing
+    USEHUGEPAGE=1 to make.
 
 Basically, it is inspired by some of the non-default options available for the
 OpenBSD allocator - see malloc.conf(5) on that platform for reference. It is
diff --git a/libdislocator/libdislocator.so.c b/libdislocator/libdislocator.so.c
index eb8e9de3..0268cc52 100644
--- a/libdislocator/libdislocator.so.c
+++ b/libdislocator/libdislocator.so.c
@@ -111,6 +111,7 @@ static __thread u32 call_depth;         /* To avoid recursion via fprintf() */
 static void* __dislocator_alloc(size_t len) {
 
   void* ret;
+  size_t tlen;
   int flags, fd, sp;
 
   if (total_mem + len > max_mem || total_mem + len < total_mem) {
@@ -123,6 +124,7 @@ static void* __dislocator_alloc(size_t len) {
 
   }
 
+  tlen = (1 + PG_COUNT(len + 8)) * PAGE_SIZE;
   flags = MAP_PRIVATE | MAP_ANONYMOUS;
   fd = -1;
 #if defined(USEHUGEPAGE)
@@ -142,8 +144,22 @@ static void* __dislocator_alloc(size_t len) {
   /* We will also store buffer length and a canary below the actual buffer, so
      let's add 8 bytes for that. */
 
-  ret = mmap(NULL, (1 + PG_COUNT(len + 8)) * PAGE_SIZE, PROT_READ | PROT_WRITE,
+  ret = mmap(NULL, tlen, PROT_READ | PROT_WRITE,
              flags, fd, 0);
+#if defined(USEHUGEPAGE)
+  /* We try one more time with regular call */
+  if (ret == MAP_FAILED) {
+#if defined(__APPLE__)
+  fd = -1;
+#elif defined(__linux__)
+  flags &= -MAP_HUGETLB;
+#elif defined(__FreeBSD__)
+  flags &= -MAP_ALIGNED_SUPER;
+#endif
+     ret = mmap(NULL, tlen, PROT_READ | PROT_WRITE,
+                flags, fd, 0);
+  }
+#endif
 
   if (ret == MAP_FAILED) {