Implementing page permission setting in MCJIT unit test SectionMemoryManager.cpp

This commit is primarily here for the revision history.  I'm about to move the SectionMemoryManager into the RuntimeDyld library, but I wanted to check the changes in here so people could see the differences in the updated implementation.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@168718 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 205 insertions, 63 deletions

diff --git a/unittests/ExecutionEngine/MCJIT/MCJITTest.cpp b/unittests/ExecutionEngine/MCJIT/MCJITTest.cpp
index 6b79a683bc..92230f3f5e 100644
--- a/unittests/ExecutionEngine/MCJIT/MCJITTest.cpp
+++ b/unittests/ExecutionEngine/MCJIT/MCJITTest.cpp
@@ -12,10 +12,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "gtest/gtest.h"
 #include "llvm/ExecutionEngine/MCJIT.h"
 #include "MCJITTestBase.h"
-#include "SectionMemoryManager.h"
-#include "gtest/gtest.h"
 
 using namespace llvm;
 
@@ -47,6 +46,7 @@ TEST_F(MCJITTest, global_variable) {
   GlobalValue *Global = insertGlobalInt32(M.get(), "test_global", initialValue);
   createJIT(M.take());
   void *globalPtr =  TheJIT->getPointerToGlobal(Global);
+  MM->applyPermissions();
   static_cast<SectionMemoryManager*>(MM)->invalidateInstructionCache();
   EXPECT_TRUE(0 != globalPtr)
     << "Unable to get pointer to global value from JIT";
@@ -61,6 +61,7 @@ TEST_F(MCJITTest, add_function) {
   Function *F = insertAddFunction(M.get());
   createJIT(M.take());
   void *addPtr = TheJIT->getPointerToFunction(F);
+  MM->applyPermissions();
   static_cast<SectionMemoryManager*>(MM)->invalidateInstructionCache();
   EXPECT_TRUE(0 != addPtr)
     << "Unable to get pointer to function from JIT";
@@ -78,6 +79,7 @@ TEST_F(MCJITTest, run_main) {
   Function *Main = insertMainFunction(M.get(), 6);
   createJIT(M.take());
   void *vPtr = TheJIT->getPointerToFunction(Main);
+  MM->applyPermissions();
   static_cast<SectionMemoryManager*>(MM)->invalidateInstructionCache();
   EXPECT_TRUE(0 != vPtr)
     << "Unable to get pointer to main() from JIT";
@@ -100,6 +102,7 @@ TEST_F(MCJITTest, return_global) {
 
   createJIT(M.take());
   void *rgvPtr = TheJIT->getPointerToFunction(ReturnGlobal);
+  MM->applyPermissions();
   static_cast<SectionMemoryManager*>(MM)->invalidateInstructionCache();
   EXPECT_TRUE(0 != rgvPtr);
 
@@ -169,6 +172,7 @@ TEST_F(MCJITTest, multiple_functions) {
 
   createJIT(M.take());
   void *vPtr = TheJIT->getPointerToFunction(Outer);
+  MM->applyPermissions();
   static_cast<SectionMemoryManager*>(MM)->invalidateInstructionCache();
   EXPECT_TRUE(0 != vPtr)
     << "Unable to get pointer to outer function from JIT";
diff --git a/unittests/ExecutionEngine/MCJIT/SectionMemoryManager.cpp b/unittests/ExecutionEngine/MCJIT/SectionMemoryManager.cpp
index 225106ecab..e8aaa82a5f 100644
--- a/unittests/ExecutionEngine/MCJIT/SectionMemoryManager.cpp
+++ b/unittests/ExecutionEngine/MCJIT/SectionMemoryManager.cpp
@@ -1,4 +1,4 @@
-//===-- SectionMemoryManager.cpp - The memory manager for MCJIT -----------===//
+//===- SectionMemoryManager.cpp - Memory manager for MCJIT/RtDyld *- C++ -*-==//
 //
 //                     The LLVM Compiler Infrastructure
 //
@@ -7,25 +7,24 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file defines the implementation of the section-based memory manager
-// used by MCJIT.
+// This file implements the section-based memory manager used by the MCJIT
+// execution engine and RuntimeDyld
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Config/config.h"
 #include "llvm/Support/DynamicLibrary.h"
 #include "llvm/Support/MathExtras.h"
-
 #include "SectionMemoryManager.h"
 
 #ifdef __linux__
-// These includes used by SectionMemoryManager::getPointerToNamedFunction()
-// for Glibc trickery. Look comments in this function for more information.
-#ifdef HAVE_SYS_STAT_H
-#include <sys/stat.h>
-#endif
-#include <fcntl.h>
-#include <unistd.h>
+  // These includes used by SectionMemoryManager::getPointerToNamedFunction()
+  // for Glibc trickery. See comments in this function for more information.
+  #ifdef HAVE_SYS_STAT_H
+    #include <sys/stat.h>
+  #endif
+  #include <fcntl.h>
+  #include <unistd.h>
 #endif
 
 namespace llvm {
@@ -34,65 +33,137 @@ uint8_t *SectionMemoryManager::allocateDataSection(uintptr_t Size,
                                                     unsigned Alignment,
                                                     unsigned SectionID,
                                                     bool IsReadOnly) {
-  if (!Alignment)
-    Alignment = 16;
-  // Ensure that enough memory is requested to allow aligning.
-  size_t NumElementsAligned = 1 + (Size + Alignment - 1)/Alignment;
-  uint8_t *Addr = (uint8_t*)calloc(NumElementsAligned, Alignment);
-
-  // Honour the alignment requirement.
-  uint8_t *AlignedAddr = (uint8_t*)RoundUpToAlignment((uint64_t)Addr, Alignment);
-
-  // Store the original address from calloc so we can free it later.
-  AllocatedDataMem.push_back(sys::MemoryBlock(Addr, NumElementsAligned*Alignment));
-  return AlignedAddr;
+  if (IsReadOnly)
+    return allocateSection(RODataMem, Size, Alignment);
+  return allocateSection(RWDataMem, Size, Alignment);
 }
 
 uint8_t *SectionMemoryManager::allocateCodeSection(uintptr_t Size,
-                                                    unsigned Alignment,
-                                                    unsigned SectionID) {
+                                                   unsigned Alignment,
+                                                   unsigned SectionID) {
+  return allocateSection(CodeMem, Size, Alignment);
+}
+
+uint8_t *SectionMemoryManager::allocateSection(MemoryGroup &MemGroup,
+                                               uintptr_t Size,
+                                               unsigned Alignment) {
   if (!Alignment)
     Alignment = 16;
-  unsigned NeedAllocate = Alignment * ((Size + Alignment - 1)/Alignment + 1);
+
+  assert(!(Alignment & (Alignment - 1)) && "Alignment must be a power of two.");
+
+  uintptr_t RequiredSize = Alignment * ((Size + Alignment - 1)/Alignment + 1);
   uintptr_t Addr = 0;
-  // Look in the list of free code memory regions and use a block there if one
+
+  // Look in the list of free memory regions and use a block there if one
   // is available.
-  for (int i = 0, e = FreeCodeMem.size(); i != e; ++i) {
-    sys::MemoryBlock &MB = FreeCodeMem[i];
-    if (MB.size() >= NeedAllocate) {
+  for (int i = 0, e = MemGroup.FreeMem.size(); i != e; ++i) {
+    sys::MemoryBlock &MB = MemGroup.FreeMem[i];
+    if (MB.size() >= RequiredSize) {
       Addr = (uintptr_t)MB.base();
       uintptr_t EndOfBlock = Addr + MB.size();
       // Align the address.
       Addr = (Addr + Alignment - 1) & ~(uintptr_t)(Alignment - 1);
       // Store cutted free memory block.
-      FreeCodeMem[i] = sys::MemoryBlock((void*)(Addr + Size),
-                                        EndOfBlock - Addr - Size);
+      MemGroup.FreeMem[i] = sys::MemoryBlock((void*)(Addr + Size),
+                                             EndOfBlock - Addr - Size);
       return (uint8_t*)Addr;
     }
   }
 
   // No pre-allocated free block was large enough. Allocate a new memory region.
-  sys::MemoryBlock MB = sys::Memory::AllocateRWX(NeedAllocate, 0, 0);
+  // Note that all sections get allocated as read-write.  The permissions will
+  // be updated later based on memory group.
+  //
+  // FIXME: It would be useful to define a default allocation size (or add
+  // it as a constructor parameter) to minimize the number of allocations.
+  // 
+  // FIXME: Initialize the Near member for each memory group to avoid
+  // interleaving.
+  error_code ec;
+  sys::MemoryBlock MB = sys::Memory::allocateMappedMemory(RequiredSize,
+                                                          &MemGroup.Near,
+                                                          sys::Memory::MF_READ |
+                                                            sys::Memory::MF_WRITE,
+                                                          ec);
+  if (ec) {
+    // FIXME: Add error propogation to the interface.
+    return NULL;
+  }
 
-  AllocatedCodeMem.push_back(MB);
+  // Save this address as the basis for our next request
+  MemGroup.Near = MB;
+
+  MemGroup.AllocatedMem.push_back(MB);
   Addr = (uintptr_t)MB.base();
   uintptr_t EndOfBlock = Addr + MB.size();
+
   // Align the address.
   Addr = (Addr + Alignment - 1) & ~(uintptr_t)(Alignment - 1);
-  // The AllocateRWX may allocate much more memory than we need. In this case,
-  // we store the unused memory as a free memory block.
+
+  // The allocateMappedMemory may allocate much more memory than we need. In
+  // this case, we store the unused memory as a free memory block.
   unsigned FreeSize = EndOfBlock-Addr-Size;
   if (FreeSize > 16)
-    FreeCodeMem.push_back(sys::MemoryBlock((void*)(Addr + Size), FreeSize));
+    MemGroup.FreeMem.push_back(sys::MemoryBlock((void*)(Addr + Size), FreeSize));
 
   // Return aligned address
   return (uint8_t*)Addr;
 }
 
+bool SectionMemoryManager::applyPermissions(std::string *ErrMsg)
+{
+  // FIXME: Should in-progress permissions be reverted if an error occurs?
+  error_code ec;
+
+  // Make code memory executable.
+  ec = applyMemoryGroupPermissions(CodeMem,
+                                   sys::Memory::MF_READ | sys::Memory::MF_EXEC);
+  if (ec) {
+    if (ErrMsg) {
+      *ErrMsg = ec.message();
+    }
+    return true;
+  }
+
+  // Make read-only data memory read-only.
+  ec = applyMemoryGroupPermissions(RODataMem,
+                                   sys::Memory::MF_READ | sys::Memory::MF_EXEC);
+  if (ec) {
+    if (ErrMsg) {
+      *ErrMsg = ec.message();
+    }
+    return true;
+  }
+
+  // Read-write data memory already has the correct permissions
+
+  return false;
+}
+
+error_code SectionMemoryManager::applyMemoryGroupPermissions(MemoryGroup &MemGroup,
+                                                             unsigned Permissions) {
+
+  for (int i = 0, e = MemGroup.AllocatedMem.size(); i != e; ++i) {
+      error_code ec;
+      ec = sys::Memory::protectMappedMemory(MemGroup.AllocatedMem[i],
+                                            Permissions);
+      if (ec) {
+        return ec;
+      }
+  }
+
+  return error_code::success();
+}
+
 void SectionMemoryManager::invalidateInstructionCache() {
-  for (int i = 0, e = AllocatedCodeMem.size(); i != e; ++i)
-    sys::Memory::InvalidateInstructionCache(AllocatedCodeMem[i].base(),
-                                            AllocatedCodeMem[i].size());
+  for (int i = 0, e = CodeMem.AllocatedMem.size(); i != e; ++i)
+    sys::Memory::InvalidateInstructionCache(CodeMem.AllocatedMem[i].base(),
+                                            CodeMem.AllocatedMem[i].size());
+}
+
+static int jit_noop() {
+  return 0;
 }
 
 void *SectionMemoryManager::getPointerToNamedFunction(const std::string &Name,
@@ -117,6 +188,14 @@ void *SectionMemoryManager::getPointerToNamedFunction(const std::string &Name,
   if (Name == "mknod") return (void*)(intptr_t)&mknod;
 #endif // __linux__
 
+  // We should not invoke parent's ctors/dtors from generated main()!
+  // On Mingw and Cygwin, the symbol __main is resolved to
+  // callee's(eg. tools/lli) one, to invoke wrong duplicated ctors
+  // (and register wrong callee's dtors with atexit(3)).
+  // We expect ExecutionEngine::runStaticConstructorsDestructors()
+  // is called before ExecutionEngine::runFunctionAsMain() is called.
+  if (Name == "__main") return (void*)(intptr_t)&jit_noop;
+
   const char *NameStr = Name.c_str();
   void *Ptr = sys::DynamicLibrary::SearchForAddressOfSymbol(NameStr);
   if (Ptr) return Ptr;
@@ -135,10 +214,13 @@ void *SectionMemoryManager::getPointerToNamedFunction(const std::string &Name,
 }
 
 SectionMemoryManager::~SectionMemoryManager() {
-  for (unsigned i = 0, e = AllocatedCodeMem.size(); i != e; ++i)
-    sys::Memory::ReleaseRWX(AllocatedCodeMem[i]);
-  for (unsigned i = 0, e = AllocatedDataMem.size(); i != e; ++i)
-    free(AllocatedDataMem[i].base());
+  for (unsigned i = 0, e = CodeMem.AllocatedMem.size(); i != e; ++i)
+    sys::Memory::releaseMappedMemory(CodeMem.AllocatedMem[i]);
+  for (unsigned i = 0, e = RWDataMem.AllocatedMem.size(); i != e; ++i)
+    sys::Memory::releaseMappedMemory(RWDataMem.AllocatedMem[i]);
+  for (unsigned i = 0, e = RODataMem.AllocatedMem.size(); i != e; ++i)
+    sys::Memory::releaseMappedMemory(RODataMem.AllocatedMem[i]);
 }
 
 } // namespace llvm
+
diff --git a/unittests/ExecutionEngine/MCJIT/SectionMemoryManager.h b/unittests/ExecutionEngine/MCJIT/SectionMemoryManager.h
index 968ee63ffd..7d5a5ff3f3 100644
--- a/unittests/ExecutionEngine/MCJIT/SectionMemoryManager.h
+++ b/unittests/ExecutionEngine/MCJIT/SectionMemoryManager.h
@@ -1,4 +1,4 @@
-//===-- SectionMemoryManager.h - Memory allocator for MCJIT -----*- C++ -*-===//
+//===- SectionMemoryManager.h - Memory manager for MCJIT/RtDyld -*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
@@ -8,7 +8,7 @@
 //===----------------------------------------------------------------------===//
 //
 // This file contains the declaration of a section-based memory manager used by
-// the MCJIT execution engine.
+// the MCJIT execution engine and RuntimeDyld.
 //
 //===----------------------------------------------------------------------===//
 
@@ -22,42 +22,97 @@
 
 namespace llvm {
 
-// Section-based memory manager for MCJIT
+/// This is a simple memory manager which implements the methods called by 
+/// the RuntimeDyld class to allocate memory for section-based loading of
+/// objects, usually those generated by the MCJIT execution engine.
+/// 
+/// This memory manager allocates all section memory as read-write.  The 
+/// RuntimeDyld will copy JITed section memory into these allocated blocks
+/// and perform any necessary linking and relocations.
+/// 
+/// Any client using this memory manager MUST ensure that section-specific
+/// page permissions have been applied before attempting to execute functions
+/// in the JITed object.  Permissions can be applied either by calling
+/// MCJIT::finalizeObject or by calling SectionMemoryManager::applyPermissions
+/// directly.  Clients of MCJIT should call MCJIT::finalizeObject.
 class SectionMemoryManager : public JITMemoryManager {
+  SectionMemoryManager(const SectionMemoryManager&) LLVM_DELETED_FUNCTION;
+  void operator=(const SectionMemoryManager&) LLVM_DELETED_FUNCTION;
 
 public:
-
   SectionMemoryManager() { }
-  ~SectionMemoryManager();
+  virtual ~SectionMemoryManager();
 
+  /// \brief Allocates a memory block of (at least) the given size suitable for
+  /// executable code.
+  /// 
+  /// The value of \p Alignment must be a power of two.  If \p Alignment is zero
+  /// a default alignment of 16 will be used.
   virtual uint8_t *allocateCodeSection(uintptr_t Size, unsigned Alignment,
                                        unsigned SectionID);
 
+  /// \brief Allocates a memory block of (at least) the given size suitable for
+  /// executable code.
+  /// 
+  /// The value of \p Alignment must be a power of two.  If \p Alignment is zero
+  /// a default alignment of 16 will be used.
   virtual uint8_t *allocateDataSection(uintptr_t Size, unsigned Alignment,
-                                       unsigned SectionID, bool IsReadOnly);
+                                       unsigned SectionID,
+                                       bool isReadOnly);
 
-  virtual bool applyPermissions(std::string *ErrMsg) { return false; }
+  /// \brief Applies section-specific memory permissions.
+  ///
+  /// This method is called when object loading is complete and section page
+  /// permissions can be applied.  It is up to the memory manager implementation
+  /// to decide whether or not to act on this method.  The memory manager will
+  /// typically allocate all sections as read-write and then apply specific
+  /// permissions when this method is called.  Code sections cannot be executed
+  /// until this function has been called.
+  ///
+  /// \returns true if an error occurred, false otherwise.
+  virtual bool applyPermissions(std::string *ErrMsg = 0);
 
+  /// This method returns the address of the specified function. As such it is
+  /// only useful for resolving library symbols, not code generated symbols.
+  ///
+  /// If \p AbortOnFailure is false and no function with the given name is
+  /// found, this function returns a null pointer. Otherwise, it prints a
+  /// message to stderr and aborts.
   virtual void *getPointerToNamedFunction(const std::string &Name,
                                           bool AbortOnFailure = true);
 
-  // Invalidate instruction cache for code sections. Some platforms with
-  // separate data cache and instruction cache require explicit cache flush,
-  // otherwise JIT code manipulations (like resolved relocations) will get to
-  // the data cache but not to the instruction cache.
+  /// \brief Invalidate instruction cache for code sections.
+  ///
+  /// Some platforms with separate data cache and instruction cache require
+  /// explicit cache flush, otherwise JIT code manipulations (like resolved
+  /// relocations) will get to the data cache but not to the instruction cache.
+  /// 
+  /// This method is not called by RuntimeDyld or MCJIT during the load
+  /// process.  Clients may call this function when needed.  See the lli
+  /// tool for example use.
   virtual void invalidateInstructionCache();
 
 private:
+  struct MemoryGroup {
+      SmallVector<sys::MemoryBlock, 16> AllocatedMem;
+      SmallVector<sys::MemoryBlock, 16> FreeMem;
+      sys::MemoryBlock Near;
+  };
 
-  SmallVector<sys::MemoryBlock, 16> AllocatedDataMem;
-  SmallVector<sys::MemoryBlock, 16> AllocatedCodeMem;
-  SmallVector<sys::MemoryBlock, 16> FreeCodeMem;
+  uint8_t *allocateSection(MemoryGroup &MemGroup, uintptr_t Size,
+                           unsigned Alignment);
 
-public:
+  error_code applyMemoryGroupPermissions(MemoryGroup &MemGroup,
+                                         unsigned Permissions);
+
+  MemoryGroup CodeMem;
+  MemoryGroup RWDataMem;
+  MemoryGroup RODataMem;
 
+public:
   ///
-  /// Functions below are not used by MCJIT, but must be implemented because
-  /// they are declared as pure virtuals in the base class.
+  /// Functions below are not used by MCJIT or RuntimeDyld, but must be
+  /// implemented because they are declared as pure virtuals in the base class.
   ///
 
   virtual void setMemoryWritable() {
@@ -118,3 +173,4 @@ public:
 }
 
 #endif // LLVM_EXECUTION_ENGINE_SECTION_MEMORY_MANAGER_H
+