Integrate KDAlloc into KLEE

1 files changed, 263 insertions, 58 deletions

diff --git a/lib/Core/MemoryManager.cpp b/lib/Core/MemoryManager.cpp
index f15c0db9..0e47500a 100644
--- a/lib/Core/MemoryManager.cpp
+++ b/lib/Core/MemoryManager.cpp
@@ -10,76 +10,261 @@
 #include "MemoryManager.h"
 
 #include "CoreStats.h"
+#include "ExecutionState.h"
 #include "Memory.h"
 
 #include "klee/Expr/Expr.h"
 #include "klee/Support/ErrorHandling.h"
 
+#include "llvm/IR/GlobalVariable.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/MathExtras.h"
+#if LLVM_VERSION_CODE >= LLVM_VERSION(10, 0)
+#include "llvm/Support/Alignment.h"
+#else
+#include "llvm/Support/MathExtras.h"
+#endif
 
-#include <inttypes.h>
+#include <cinttypes>
+#include <algorithm>
 #include <sys/mman.h>
+#include <tuple>
+#include <string>
 
 using namespace klee;
 
-namespace {
+namespace klee {
+std::uint32_t MemoryManager::quarantine;
+
+std::size_t MemoryManager::pageSize = sysconf(_SC_PAGE_SIZE);
 
+bool MemoryManager::isDeterministic;
+} // namespace klee
+
+namespace {
 llvm::cl::OptionCategory MemoryCat("Memory management options",
                                    "These options control memory management.");
 
-llvm::cl::opt<bool> DeterministicAllocation(
-    "allocate-determ",
+llvm::cl::opt<bool, true> DeterministicAllocation(
+    "kdalloc",
     llvm::cl::desc("Allocate memory deterministically (default=false)"),
-    llvm::cl::init(false), llvm::cl::cat(MemoryCat));
+    llvm::cl::location(MemoryManager::isDeterministic), llvm::cl::init(false),
+    llvm::cl::cat(MemoryCat));
+
+llvm::cl::opt<bool> DeterministicAllocationMarkAsUnneeded(
+    "kdalloc-mark-as-unneeded",
+    llvm::cl::desc("Mark allocations as unneeded after external function calls "
+                   "(default=true)"),
+    llvm::cl::init(true), llvm::cl::cat(MemoryCat));
+
+llvm::cl::opt<unsigned> DeterministicAllocationGlobalsSize(
+    "kdalloc-globals-size",
+    llvm::cl::desc("Reserved memory for globals in GiB (default=10)"),
+    llvm::cl::init(10), llvm::cl::cat(MemoryCat));
+
+llvm::cl::opt<unsigned> DeterministicAllocationConstantsSize(
+    "kdalloc-constants-size",
+    llvm::cl::desc("Reserved memory for constant globals in GiB (default=10)"),
+    llvm::cl::init(10), llvm::cl::cat(MemoryCat));
+
+llvm::cl::opt<unsigned> DeterministicAllocationHeapSize(
+    "kdalloc-heap-size",
+    llvm::cl::desc("Reserved memory for heap in GiB (default=1024)"),
+    llvm::cl::init(1024), llvm::cl::cat(MemoryCat));
+
+llvm::cl::opt<unsigned> DeterministicAllocationStackSize(
+    "kdalloc-stack-size",
+    llvm::cl::desc("Reserved memory for stack in GiB (default=100)"),
+    llvm::cl::init(128), llvm::cl::cat(MemoryCat));
 
-llvm::cl::opt<unsigned> DeterministicAllocationSize(
-    "allocate-determ-size",
+llvm::cl::opt<std::uintptr_t> DeterministicAllocationGlobalsStartAddress(
+    "kdalloc-globals-start-address",
     llvm::cl::desc(
-        "Preallocated memory for deterministic allocation in MB (default=100)"),
-    llvm::cl::init(100), llvm::cl::cat(MemoryCat));
+        "Start address for globals segment (has to be page aligned)"),
+    llvm::cl::cat(MemoryCat));
+
+llvm::cl::opt<std::uintptr_t> DeterministicAllocationConstantsStartAddress(
+    "kdalloc-constants-start-address",
+    llvm::cl::desc(
+        "Start address for constant globals segment (has to be page aligned)"),
+    llvm::cl::cat(MemoryCat));
+
+llvm::cl::opt<std::uintptr_t> DeterministicAllocationHeapStartAddress(
+    "kdalloc-heap-start-address",
+    llvm::cl::desc("Start address for heap segment (has to be page aligned)"),
+    llvm::cl::cat(MemoryCat));
+
+llvm::cl::opt<std::uintptr_t> DeterministicAllocationStackStartAddress(
+    "kdalloc-stack-start-address",
+    llvm::cl::desc("Start address for stack segment (has to be page aligned)"),
+    llvm::cl::cat(MemoryCat));
+
+struct QuarantineSizeParser : public llvm::cl::parser<std::uint32_t> {
+  explicit QuarantineSizeParser(llvm::cl::Option &O)
+      : llvm::cl::parser<std::uint32_t>(O) {}
+
+  bool parse(llvm::cl::Option &O, llvm::StringRef ArgName, llvm::StringRef Arg,
+             std::uint32_t &Val) {
+    if (Arg == "-1") {
+      Val = kdalloc::Allocator::unlimitedQuarantine;
+    } else if (Arg.getAsInteger(0, Val)) {
+      return O.error("'" + Arg + "' value invalid!");
+    }
+
+    return false;
+  }
+};
+
+llvm::cl::opt<std::uint32_t, true, QuarantineSizeParser>
+    DeterministicAllocationQuarantineSize(
+        "kdalloc-quarantine",
+        llvm::cl::desc("Size of quarantine queues in allocator (default=8, "
+                       "disabled=0, unlimited=-1)"),
+        llvm::cl::location(klee::MemoryManager::quarantine),
+        llvm::cl::value_desc("size"), llvm::cl::init(8),
+        llvm::cl::cat(MemoryCat));
 
 llvm::cl::opt<bool> NullOnZeroMalloc(
     "return-null-on-zero-malloc",
     llvm::cl::desc("Returns NULL if malloc(0) is called (default=false)"),
     llvm::cl::init(false), llvm::cl::cat(MemoryCat));
-
-llvm::cl::opt<unsigned> RedzoneSize(
-    "redzone-size",
-    llvm::cl::desc("Set the size of the redzones to be added after each "
-                   "allocation (in bytes). This is important to detect "
-                   "out-of-bounds accesses (default=10)"),
-    llvm::cl::init(10), llvm::cl::cat(MemoryCat));
-
-llvm::cl::opt<unsigned long long> DeterministicStartAddress(
-    "allocate-determ-start-address",
-    llvm::cl::desc("Start address for deterministic allocation. Has to be page "
-                   "aligned (default=0x7ff30000000)"),
-    llvm::cl::init(0x7ff30000000), llvm::cl::cat(MemoryCat));
 } // namespace
 
 /***/
 MemoryManager::MemoryManager(ArrayCache *_arrayCache)
-    : arrayCache(_arrayCache), deterministicSpace(0), nextFreeSlot(0),
-      spaceSize(DeterministicAllocationSize.getValue() * 1024 * 1024) {
+    : arrayCache(_arrayCache) {
   if (DeterministicAllocation) {
-    // Page boundary
-    void *expectedAddress = (void *)DeterministicStartAddress.getValue();
+    if (DeterministicAllocationQuarantineSize ==
+        kdalloc::Allocator::unlimitedQuarantine) {
+      klee_message("Deterministic allocator: Using unlimited quarantine");
+    } else if (DeterministicAllocationQuarantineSize != 0) {
+      klee_message("Deterministic allocator: Using quarantine queue size %u",
+                   DeterministicAllocationQuarantineSize.getValue());
+    }
 
-    char *newSpace =
-        (char *)mmap(expectedAddress, spaceSize, PROT_READ | PROT_WRITE,
-                     MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
+    std::vector<std::tuple<std::string,
+                           std::uintptr_t, // start address (0 if none
+                                           // requested)
+                           std::size_t,    // size of segment
+                           std::reference_wrapper<kdalloc::AllocatorFactory>,
+                           kdalloc::Allocator *>>
+        requestedSegments;
+    requestedSegments.emplace_back(
+        "globals",
+        DeterministicAllocationGlobalsStartAddress
+            ? DeterministicAllocationGlobalsStartAddress.getValue()
+            : 0,
+        static_cast<std::size_t>(
+            DeterministicAllocationGlobalsSize.getValue()) *
+            1024 * 1024 * 1024,
+        globalsFactory, &globalsAllocator);
+    requestedSegments.emplace_back(
+        "constants",
+        DeterministicAllocationConstantsStartAddress
+            ? DeterministicAllocationConstantsStartAddress.getValue()
+            : 0,
+        static_cast<std::size_t>(
+            DeterministicAllocationConstantsSize.getValue()) *
+            1024 * 1024 * 1024,
+        constantsFactory, &constantsAllocator);
+    requestedSegments.emplace_back(
+        "heap",
+        DeterministicAllocationHeapStartAddress
+            ? DeterministicAllocationHeapStartAddress.getValue()
+            : 0,
+        static_cast<std::size_t>(DeterministicAllocationHeapSize.getValue()) *
+            1024 * 1024 * 1024,
+        heapFactory, nullptr);
+    requestedSegments.emplace_back(
+        "stack",
+        DeterministicAllocationStackStartAddress
+            ? DeterministicAllocationStackStartAddress.getValue()
+            : 0,
+        static_cast<std::size_t>(DeterministicAllocationStackSize.getValue()) *
+            1024 * 1024 * 1024,
+        stackFactory, nullptr);
 
-    if (newSpace == MAP_FAILED) {
-      klee_error("Couldn't mmap() memory for deterministic allocations");
-    }
-    if (expectedAddress != newSpace && expectedAddress != 0) {
-      klee_error("Could not allocate memory deterministically");
+    // check invariants
+#if LLVM_VERSION_CODE >= LLVM_VERSION(10, 0)
+    llvm::Align pageAlignment(pageSize);
+#endif
+    for (auto &requestedSegment : requestedSegments) {
+      auto &segment1 = std::get<0>(requestedSegment);
+      auto &start1 = std::get<1>(requestedSegment);
+      auto &size1 = std::get<2>(requestedSegment);
+      // check for page alignment
+      // NOTE: sizes are assumed to be page aligned due to multiplication
+#if LLVM_VERSION_CODE >= LLVM_VERSION(10, 0)
+      if (start1 != 0 && !llvm::isAligned(pageAlignment, start1)) {
+        klee_error("Deterministic allocator: Requested start address for %s "
+                   "is not page aligned (page size: %zu B)",
+                   segment1.c_str(), pageAlignment.value());
+      }
+#else
+      if (start1 != 0 && llvm::OffsetToAlignment(start1, pageSize) != 0) {
+        klee_error("Deterministic allocator: Requested start address for %s "
+                   "is not page aligned (page size: %zu B)",
+                   segment1.c_str(), pageSize);
+      }
+#endif
+
+      // check for overlap of segments
+      std::uintptr_t end1 = start1 + size1;
+      for (auto &requestedSegment : requestedSegments) {
+        auto &segment2 = std::get<0>(requestedSegment);
+        auto &start2 = std::get<1>(requestedSegment);
+        auto &size2 = std::get<2>(requestedSegment);
+        if (start1 != 0 && start2 != 0 && segment1 != segment2) {
+          std::uintptr_t end2 = start2 + size2;
+          if (!(end1 <= start2 || start1 >= end2)) {
+            klee_error("Deterministic allocator: Requested mapping for %s "
+                       "(start-address=0x%" PRIxPTR " size=%zu GiB) "
+                       "overlaps with that for %s "
+                       "(start-address=0x%" PRIxPTR " size=%zu GiB)",
+                       segment1.c_str(), start1, size1 / (1024 * 1024 * 1024),
+                       segment2.c_str(), start2, size2 / (1024 * 1024 * 1024));
+          }
+        }
+      }
     }
 
-    klee_message("Deterministic memory allocation starting from %p", newSpace);
-    deterministicSpace = newSpace;
-    nextFreeSlot = newSpace;
+    // initialize factories and allocators
+    for (auto &requestedSegment : requestedSegments) {
+      auto &segment = std::get<0>(requestedSegment);
+      auto &start = std::get<1>(requestedSegment);
+      auto &size = std::get<2>(requestedSegment);
+      auto &factory = std::get<3>(requestedSegment);
+      auto &allocator = std::get<4>(requestedSegment);
+      factory.get() = kdalloc::AllocatorFactory(
+          start, size, DeterministicAllocationQuarantineSize);
+
+      if (!factory.get()) {
+        klee_error(
+            "Deterministic allocator: Could not allocate mapping for %s: %s",
+            segment.c_str(), strerror(errno));
+      }
+      if (start && factory.get().getMapping().getBaseAddress() !=
+                       reinterpret_cast<void *>(start)) {
+        klee_error("Deterministic allocator: Could not allocate mapping for %s "
+                   "at requested address",
+                   segment.c_str());
+      }
+      if (factory.get().getMapping().getSize() != size) {
+        klee_error("Deterministic allocator: Could not allocate mapping for %s "
+                   "with the requested size",
+                   segment.c_str());
+      }
+
+      klee_message("Deterministic allocator: %s "
+                   "(start-address=0x%" PRIxPTR " size=%zu GiB)",
+                   segment.c_str(),
+                   reinterpret_cast<std::uintptr_t>(
+                       factory.get().getMapping().getBaseAddress()),
+                   size / (1024 * 1024 * 1024));
+      if (allocator) {
+        *allocator = factory.get().makeAllocator();
+      }
+    }
   }
 }
 
@@ -91,13 +276,10 @@ MemoryManager::~MemoryManager() {
     objects.erase(mo);
     delete mo;
   }
-
-  if (DeterministicAllocation)
-    munmap(deterministicSpace, spaceSize);
 }
 
 MemoryObject *MemoryManager::allocate(uint64_t size, bool isLocal,
-                                      bool isGlobal,
+                                      bool isGlobal, ExecutionState *state,
                                       const llvm::Value *allocSite,
                                       size_t alignment) {
   if (size > 10 * 1024 * 1024)
@@ -116,19 +298,29 @@ MemoryObject *MemoryManager::allocate(uint64_t size, bool isLocal,
 
   uint64_t address = 0;
   if (DeterministicAllocation) {
-    address = llvm::alignTo((uint64_t)nextFreeSlot + alignment - 1, alignment);
+    void *allocAddress;
 
-    // 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 + RedzoneSize;
+    if (isGlobal) {
+      const llvm::GlobalVariable *gv =
+          dyn_cast<llvm::GlobalVariable>(allocSite);
+      if (isa<llvm::Function>(allocSite) || (gv && gv->isConstant())) {
+        allocAddress = constantsAllocator.allocate(
+            std::max(size, static_cast<std::uint64_t>(alignment)));
+      } else {
+        allocAddress = globalsAllocator.allocate(
+            std::max(size, static_cast<std::uint64_t>(alignment)));
+      }
     } else {
-      klee_warning_once(0, "Couldn't allocate %" PRIu64
-                           " bytes. Not enough deterministic space left.",
-                        size);
-      address = 0;
+      if (isLocal) {
+        allocAddress = state->stackAllocator.allocate(
+            std::max(size, static_cast<std::uint64_t>(alignment)));
+      } else {
+        allocAddress = state->heapAllocator.allocate(
+            std::max(size, static_cast<std::uint64_t>(alignment)));
+      }
     }
+
+    address = reinterpret_cast<std::uint64_t>(allocAddress);
   } else {
     // Use malloc for the standard case
     if (alignment <= 8)
@@ -146,8 +338,8 @@ MemoryObject *MemoryManager::allocate(uint64_t size, bool isLocal,
     return 0;
 
   ++stats::allocations;
-  MemoryObject *res = new MemoryObject(address, size, isLocal, isGlobal, false,
-                                       allocSite, this);
+  MemoryObject *res = new MemoryObject(address, size, alignment, isLocal,
+                                       isGlobal, false, allocSite, this);
   objects.insert(res);
   return res;
 }
@@ -165,13 +357,11 @@ MemoryObject *MemoryManager::allocateFixed(uint64_t address, uint64_t size,
 
   ++stats::allocations;
   MemoryObject *res =
-      new MemoryObject(address, size, false, true, true, allocSite, this);
+      new MemoryObject(address, size, 0, false, true, true, allocSite, this);
   objects.insert(res);
   return res;
 }
 
-void MemoryManager::deallocate(const MemoryObject *mo) { assert(0); }
-
 void MemoryManager::markFreed(MemoryObject *mo) {
   if (objects.find(mo) != objects.end()) {
     if (!mo->isFixed && !DeterministicAllocation)
@@ -180,6 +370,21 @@ void MemoryManager::markFreed(MemoryObject *mo) {
   }
 }
 
-size_t MemoryManager::getUsedDeterministicSize() {
-  return nextFreeSlot - deterministicSpace;
+bool MemoryManager::markMappingsAsUnneeded() {
+  if (!DeterministicAllocation)
+    return false;
+
+  if (!DeterministicAllocationMarkAsUnneeded)
+    return false;
+
+  globalsFactory.getMapping().clear();
+  heapFactory.getMapping().clear();
+  stackFactory.getMapping().clear();
+
+  return true;
 }
+
+size_t MemoryManager::getUsedDeterministicSize() {
+  // TODO: implement
+  return 0;
+}
\ No newline at end of file