1 files changed, 96 insertions, 4 deletions

diff --git a/lib/Core/MemoryManager.cpp b/lib/Core/MemoryManager.cpp
index 02bbe678..560bb6f6 100644
--- a/lib/Core/MemoryManager.cpp
+++ b/lib/Core/MemoryManager.cpp
@@ -11,25 +11,84 @@
 #include "Memory.h"
 #include "MemoryManager.h"
 
-#include "klee/ExecutionState.h"
 #include "klee/Expr.h"
-#include "klee/Solver.h"
 #include "klee/Internal/Support/ErrorHandling.h"
 
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/MathExtras.h"
 
+#include <sys/mman.h>
 using namespace klee;
 
+namespace {
+llvm::cl::opt<bool> DeterministicAllocation(
+    "allocate-determ",
+    llvm::cl::desc("Allocate memory deterministically(default=off)"),
+    llvm::cl::init(false));
+
+llvm::cl::opt<unsigned> DeterministicAllocationSize(
+    "allocate-determ-size",
+    llvm::cl::desc(
+        "Preallocated memory for deterministic allocation in MB (default=100)"),
+    llvm::cl::init(100));
+
+llvm::cl::opt<bool>
+    NullOnZeroMalloc("return-null-on-zero-malloc",
+                     llvm::cl::desc("Returns NULL in case malloc(size) was "
+                                    "called with size 0 (default=off)."),
+                     llvm::cl::init(false));
+
+llvm::cl::opt<unsigned> RedZoneSpace(
+    "red-zone-space",
+    llvm::cl::desc("Set the amount of free space between allocations. This is "
+                   "important to detect out-of-bound accesses (default=10)."),
+    llvm::cl::init(10));
+
+llvm::cl::opt<uint64_t> DeterministicStartAddress(
+    "allocate-determ-start-address",
+    llvm::cl::desc("Start address for deterministic allocation. Has to be page "
+                   "aligned (default=0x7ff30000000)."),
+    llvm::cl::init(0x7ff30000000));
+}
+
 /***/
+MemoryManager::MemoryManager(ArrayCache *_arrayCache)
+    : arrayCache(_arrayCache), deterministicSpace(0),
+      nextFreeSlot(0),
+      spaceSize(DeterministicAllocationSize.getValue() * 1024 * 1024) {
+  if (DeterministicAllocation) {
+    // Page boundary
+    void *expectedAddress = (void *)DeterministicStartAddress.getValue();
+
+    char *newSpace =
+        (char *)mmap(expectedAddress, spaceSize, PROT_READ | PROT_WRITE,
+                     MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
+
+    if (newSpace == MAP_FAILED) {
+      klee_error("Couldn't mmap() memory for deterministic allocations");
+    }
+    if (expectedAddress != newSpace) {
+      klee_error("Could not allocate memory deterministically");
+    }
+
+    klee_message("Deterministic memory allocation starting from %p",
+                 expectedAddress);
+    deterministicSpace = newSpace;
+    nextFreeSlot = newSpace;
+  }
+}
 
 MemoryManager::~MemoryManager() { 
   while (!objects.empty()) {
     MemoryObject *mo = *objects.begin();
-    if (!mo->isFixed)
+    if (!mo->isFixed && !DeterministicAllocation)
       free((void *)mo->address);
     objects.erase(mo);
     delete mo;
   }
+
+  if (DeterministicAllocation)
+    munmap(deterministicSpace, spaceSize);
 }
 
 MemoryObject *MemoryManager::allocate(uint64_t size, bool isLocal,
@@ -40,7 +99,34 @@ MemoryObject *MemoryManager::allocate(uint64_t size, bool isLocal,
     klee_warning_once(0, "Large alloc: %lu bytes.  KLEE may run out of memory.",
                       size);
 
+  // Return NULL if size is zero, this is equal to error during allocation
+  if (NullOnZeroMalloc && size == 0)
+    return 0;
+
+  if (!llvm::isPowerOf2_64(alignment)) {
+    klee_warning("Only alignment of power of two is supported");
+    return 0;
+  }
+
   uint64_t address = 0;
+  if (DeterministicAllocation) {
+
+    address = llvm::RoundUpToAlignment((uint64_t)nextFreeSlot + alignment - 1,
+                                       alignment);
+
+    // Handle the case of 0-sized allocations as 1-byte allocations.
+    // This way, we make sure we have this allocation between its own red zones
+    size_t alloc_size = std::max(size, (uint64_t)1);
+    if ((char *)address + alloc_size < deterministicSpace + spaceSize) {
+      nextFreeSlot = (char *)address + alloc_size + RedZoneSpace;
+    } else {
+      klee_warning_once(
+          0,
+          "Couldn't allocate %lu bytes. Not enough deterministic space left.",
+          size);
+      address = 0;
+    }
+  } else {
     // Use malloc for the standard case
     if (alignment <= 8)
       address = (uint64_t)malloc(size);
@@ -51,6 +137,8 @@ MemoryObject *MemoryManager::allocate(uint64_t size, bool isLocal,
         address = 0;
       }
     }
+  }
+
   if (!address)
     return 0;
   
@@ -86,8 +174,12 @@ void MemoryManager::deallocate(const MemoryObject *mo) {
 void MemoryManager::markFreed(MemoryObject *mo) {
   if (objects.find(mo) != objects.end())
   {
-    if (!mo->isFixed)
+    if (!mo->isFixed && !DeterministicAllocation)
       free((void *)mo->address);
     objects.erase(mo);
   }
 }
+
+size_t MemoryManager::getUsedDeterministicSize() {
+  return nextFreeSlot - deterministicSpace;
+}