1 files changed, 183 insertions, 129 deletions

diff --git a/lib/xray/xray_interface.cc b/lib/xray/xray_interface.cc
index 7ad6a9b84..766313e85 100644
--- a/lib/xray/xray_interface.cc
+++ b/lib/xray/xray_interface.cc
@@ -23,12 +23,15 @@
 
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray_defs.h"
+#include "xray_flags.h"
+
+extern __sanitizer::SpinMutex XRayInstrMapMutex;
+extern __sanitizer::atomic_uint8_t XRayInitialized;
+extern __xray::XRaySledMap XRayInstrMap;
 
 namespace __xray {
 
 #if defined(__x86_64__)
-// FIXME: The actual length is 11 bytes. Why was length 12 passed to mprotect()
-// ?
 static const int16_t cSledLength = 12;
 #elif defined(__aarch64__)
 static const int16_t cSledLength = 32;
@@ -53,6 +56,10 @@ __sanitizer::atomic_uintptr_t XRayArgLogger{0};
 // This is the function to call when we encounter a custom event log call.
 __sanitizer::atomic_uintptr_t XRayPatchedCustomEvent{0};
 
+// This is the global status to determine whether we are currently
+// patching/unpatching.
+__sanitizer::atomic_uint8_t XRayPatching{0};
+
 // MProtectHelper is an RAII wrapper for calls to mprotect(...) that will undo
 // any successful mprotect(...) changes. This is used to make a page writeable
 // and executable, and upon destruction if it was successful in doing so returns
@@ -88,85 +95,10 @@ public:
   }
 };
 
-} // namespace __xray
-
-extern __sanitizer::SpinMutex XRayInstrMapMutex;
-extern __sanitizer::atomic_uint8_t XRayInitialized;
-extern __xray::XRaySledMap XRayInstrMap;
-
-int __xray_set_handler(void (*entry)(int32_t,
-                                     XRayEntryType)) XRAY_NEVER_INSTRUMENT {
-  if (__sanitizer::atomic_load(&XRayInitialized,
-                               __sanitizer::memory_order_acquire)) {
-
-    __sanitizer::atomic_store(&__xray::XRayPatchedFunction,
-                              reinterpret_cast<uintptr_t>(entry),
-                              __sanitizer::memory_order_release);
-    return 1;
-  }
-  return 0;
-}
-
-int __xray_set_customevent_handler(void (*entry)(void *, size_t))
-    XRAY_NEVER_INSTRUMENT {
-  if (__sanitizer::atomic_load(&XRayInitialized,
-                               __sanitizer::memory_order_acquire)) {
-    __sanitizer::atomic_store(&__xray::XRayPatchedCustomEvent,
-                              reinterpret_cast<uintptr_t>(entry),
-                              __sanitizer::memory_order_release);
-    return 1;
-  }
-  return 0;
-}
-
-
-int __xray_remove_handler() XRAY_NEVER_INSTRUMENT {
-  return __xray_set_handler(nullptr);
-}
-
-int __xray_remove_customevent_handler() XRAY_NEVER_INSTRUMENT {
-  return __xray_set_customevent_handler(nullptr);
-}
-
-__sanitizer::atomic_uint8_t XRayPatching{0};
-
-using namespace __xray;
-
-// FIXME: Figure out whether we can move this class to sanitizer_common instead
-// as a generic "scope guard".
-template <class Function> class CleanupInvoker {
-  Function Fn;
-
-public:
-  explicit CleanupInvoker(Function Fn) XRAY_NEVER_INSTRUMENT : Fn(Fn) {}
-  CleanupInvoker(const CleanupInvoker &) XRAY_NEVER_INSTRUMENT = default;
-  CleanupInvoker(CleanupInvoker &&) XRAY_NEVER_INSTRUMENT = default;
-  CleanupInvoker &
-  operator=(const CleanupInvoker &) XRAY_NEVER_INSTRUMENT = delete;
-  CleanupInvoker &operator=(CleanupInvoker &&) XRAY_NEVER_INSTRUMENT = delete;
-  ~CleanupInvoker() XRAY_NEVER_INSTRUMENT { Fn(); }
-};
-
-template <class Function>
-CleanupInvoker<Function> scopeCleanup(Function Fn) XRAY_NEVER_INSTRUMENT {
-  return CleanupInvoker<Function>{Fn};
-}
-
-inline bool patchSled(const XRaySledEntry &Sled, bool Enable,
-                      int32_t FuncId) XRAY_NEVER_INSTRUMENT {
-  // While we're here, we should patch the nop sled. To do that we mprotect
-  // the page containing the function to be writeable.
-  const uint64_t PageSize = GetPageSizeCached();
-  void *PageAlignedAddr =
-      reinterpret_cast<void *>(Sled.Address & ~(PageSize - 1));
-  std::size_t MProtectLen = (Sled.Address + cSledLength) -
-                            reinterpret_cast<uint64_t>(PageAlignedAddr);
-  MProtectHelper Protector(PageAlignedAddr, MProtectLen);
-  if (Protector.MakeWriteable() == -1) {
-    printf("Failed mprotect: %d\n", errno);
-    return XRayPatchingStatus::FAILED;
-  }
+namespace {
 
+bool patchSled(const XRaySledEntry &Sled, bool Enable,
+               int32_t FuncId) XRAY_NEVER_INSTRUMENT {
   bool Success = false;
   switch (Sled.Kind) {
   case XRayEntryType::ENTRY:
@@ -191,6 +123,55 @@ inline bool patchSled(const XRaySledEntry &Sled, bool Enable,
   return Success;
 }
 
+XRayPatchingStatus patchFunction(int32_t FuncId,
+                                 bool Enable) XRAY_NEVER_INSTRUMENT {
+  if (!__sanitizer::atomic_load(&XRayInitialized,
+                                __sanitizer::memory_order_acquire))
+    return XRayPatchingStatus::NOT_INITIALIZED; // Not initialized.
+
+  uint8_t NotPatching = false;
+  if (!__sanitizer::atomic_compare_exchange_strong(
+          &XRayPatching, &NotPatching, true, __sanitizer::memory_order_acq_rel))
+    return XRayPatchingStatus::ONGOING; // Already patching.
+
+  // Next, we look for the function index.
+  XRaySledMap InstrMap;
+  {
+    __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
+    InstrMap = XRayInstrMap;
+  }
+
+  // If we don't have an index, we can't patch individual functions.
+  if (InstrMap.Functions == 0)
+    return XRayPatchingStatus::NOT_INITIALIZED;
+
+  // FuncId must be a positive number, less than the number of functions
+  // instrumented.
+  if (FuncId <= 0 || static_cast<size_t>(FuncId) > InstrMap.Functions) {
+    Report("Invalid function id provided: %d\n", FuncId);
+    return XRayPatchingStatus::FAILED;
+  }
+
+  // Now we patch ths sleds for this specific function.
+  auto SledRange = InstrMap.SledsIndex[FuncId - 1];
+  auto *f = SledRange.Begin;
+  auto *e = SledRange.End;
+
+  bool SucceedOnce = false;
+  while (f != e)
+    SucceedOnce |= patchSled(*f++, Enable, FuncId);
+
+  __sanitizer::atomic_store(&XRayPatching, false,
+                            __sanitizer::memory_order_release);
+
+  if (!SucceedOnce) {
+    Report("Failed patching any sled for function '%d'.", FuncId);
+    return XRayPatchingStatus::FAILED;
+  }
+
+  return XRayPatchingStatus::SUCCESS;
+}
+
 // controlPatching implements the common internals of the patching/unpatching
 // implementation. |Enable| defines whether we're enabling or disabling the
 // runtime XRay instrumentation.
@@ -205,14 +186,13 @@ XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
     return XRayPatchingStatus::ONGOING; // Already patching.
 
   uint8_t PatchingSuccess = false;
-  auto XRayPatchingStatusResetter = scopeCleanup([&PatchingSuccess] {
-    if (!PatchingSuccess)
-      __sanitizer::atomic_store(&XRayPatching, false,
-                                __sanitizer::memory_order_release);
-  });
-
-  // Step 1: Compute the function id, as a unique identifier per function in the
-  // instrumentation map.
+  auto XRayPatchingStatusResetter =
+      __sanitizer::at_scope_exit([&PatchingSuccess] {
+        if (!PatchingSuccess)
+          __sanitizer::atomic_store(&XRayPatching, false,
+                                    __sanitizer::memory_order_release);
+      });
+
   XRaySledMap InstrMap;
   {
     __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
@@ -221,16 +201,47 @@ XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
   if (InstrMap.Entries == 0)
     return XRayPatchingStatus::NOT_INITIALIZED;
 
-  const uint64_t PageSize = GetPageSizeCached();
+  uint32_t FuncId = 1;
+  uint64_t CurFun = 0;
+
+  // First we want to find the bounds for which we have instrumentation points,
+  // and try to get as few calls to mprotect(...) as possible. We're assuming
+  // that all the sleds for the instrumentation map are contiguous as a single
+  // set of pages. When we do support dynamic shared object instrumentation,
+  // we'll need to do this for each set of page load offsets per DSO loaded. For
+  // now we're assuming we can mprotect the whole section of text between the
+  // minimum sled address and the maximum sled address (+ the largest sled
+  // size).
+  auto MinSled = InstrMap.Sleds[0];
+  auto MaxSled = InstrMap.Sleds[InstrMap.Entries - 1];
+  for (std::size_t I = 0; I < InstrMap.Entries; I++) {
+    const auto &Sled = InstrMap.Sleds[I];
+    if (Sled.Address < MinSled.Address)
+      MinSled = Sled;
+    if (Sled.Address > MaxSled.Address)
+      MaxSled = Sled;
+  }
+
+  const size_t PageSize = flags()->xray_page_size_override > 0
+                              ? flags()->xray_page_size_override
+                              : GetPageSizeCached();
   if ((PageSize == 0) || ((PageSize & (PageSize - 1)) != 0)) {
     Report("System page size is not a power of two: %lld\n", PageSize);
     return XRayPatchingStatus::FAILED;
   }
 
-  uint32_t FuncId = 1;
-  uint64_t CurFun = 0;
-  for (std::size_t I = 0; I < InstrMap.Entries; I++) {
-    auto Sled = InstrMap.Sleds[I];
+  void *PageAlignedAddr =
+      reinterpret_cast<void *>(MinSled.Address & ~(PageSize - 1));
+  size_t MProtectLen =
+      (MaxSled.Address - reinterpret_cast<uptr>(PageAlignedAddr)) + cSledLength;
+  MProtectHelper Protector(PageAlignedAddr, MProtectLen);
+  if (Protector.MakeWriteable() == -1) {
+    Report("Failed mprotect: %d\n", errno);
+    return XRayPatchingStatus::FAILED;
+  }
+
+  for (std::size_t I = 0; I < InstrMap.Entries; ++I) {
+    auto &Sled = InstrMap.Sleds[I];
     auto F = Sled.Function;
     if (CurFun == 0)
       CurFun = F;
@@ -246,36 +257,14 @@ XRayPatchingStatus controlPatching(bool Enable) XRAY_NEVER_INSTRUMENT {
   return XRayPatchingStatus::SUCCESS;
 }
 
-XRayPatchingStatus __xray_patch() XRAY_NEVER_INSTRUMENT {
-  return controlPatching(true);
-}
-
-XRayPatchingStatus __xray_unpatch() XRAY_NEVER_INSTRUMENT {
-  return controlPatching(false);
-}
-
-XRayPatchingStatus patchFunction(int32_t FuncId,
-                                 bool Enable) XRAY_NEVER_INSTRUMENT {
-  if (!__sanitizer::atomic_load(&XRayInitialized,
-                                __sanitizer::memory_order_acquire))
-    return XRayPatchingStatus::NOT_INITIALIZED; // Not initialized.
-
-  uint8_t NotPatching = false;
-  if (!__sanitizer::atomic_compare_exchange_strong(
-          &XRayPatching, &NotPatching, true, __sanitizer::memory_order_acq_rel))
-    return XRayPatchingStatus::ONGOING; // Already patching.
-
-  // Next, we look for the function index.
+XRayPatchingStatus mprotectAndPatchFunction(int32_t FuncId,
+                                            bool Enable) XRAY_NEVER_INSTRUMENT {
   XRaySledMap InstrMap;
   {
     __sanitizer::SpinMutexLock Guard(&XRayInstrMapMutex);
     InstrMap = XRayInstrMap;
   }
 
-  // If we don't have an index, we can't patch individual functions.
-  if (InstrMap.Functions == 0)
-    return XRayPatchingStatus::NOT_INITIALIZED;
-
   // FuncId must be a positive number, less than the number of functions
   // instrumented.
   if (FuncId <= 0 || static_cast<size_t>(FuncId) > InstrMap.Functions) {
@@ -283,33 +272,98 @@ XRayPatchingStatus patchFunction(int32_t FuncId,
     return XRayPatchingStatus::FAILED;
   }
 
-  // Now we patch ths sleds for this specific function.
+  const size_t PageSize = flags()->xray_page_size_override > 0
+                              ? flags()->xray_page_size_override
+                              : GetPageSizeCached();
+  if ((PageSize == 0) || ((PageSize & (PageSize - 1)) != 0)) {
+    Report("Provided page size is not a power of two: %lld\n", PageSize);
+    return XRayPatchingStatus::FAILED;
+  }
+
+  // Here we compute the minumum sled and maximum sled associated with a
+  // particular function ID.
   auto SledRange = InstrMap.SledsIndex[FuncId - 1];
   auto *f = SledRange.Begin;
   auto *e = SledRange.End;
+  auto MinSled = *f;
+  auto MaxSled = *(SledRange.End - 1);
+  while (f != e) {
+    if (f->Address < MinSled.Address)
+      MinSled = *f;
+    if (f->Address > MaxSled.Address)
+      MaxSled = *f;
+    ++f;
+  }
 
-  bool SucceedOnce = false;
-  while (f != e)
-    SucceedOnce |= patchSled(*f++, Enable, FuncId);
+  void *PageAlignedAddr =
+      reinterpret_cast<void *>(MinSled.Address & ~(PageSize - 1));
+  size_t MProtectLen =
+      (MaxSled.Address - reinterpret_cast<uptr>(PageAlignedAddr)) + cSledLength;
+  MProtectHelper Protector(PageAlignedAddr, MProtectLen);
+  if (Protector.MakeWriteable() == -1) {
+    Report("Failed mprotect: %d\n", errno);
+    return XRayPatchingStatus::FAILED;
+  }
+  return patchFunction(FuncId, Enable);
+}
 
-  __sanitizer::atomic_store(&XRayPatching, false,
-                            __sanitizer::memory_order_release);
+} // namespace
 
-  if (!SucceedOnce) {
-    Report("Failed patching any sled for function '%d'.", FuncId);
-    return XRayPatchingStatus::FAILED;
+} // namespace __xray
+
+using namespace __xray;
+
+// The following functions are declared `extern "C" {...}` in the header, hence
+// they're defined in the global namespace.
+
+int __xray_set_handler(void (*entry)(int32_t,
+                                     XRayEntryType)) XRAY_NEVER_INSTRUMENT {
+  if (__sanitizer::atomic_load(&XRayInitialized,
+                               __sanitizer::memory_order_acquire)) {
+
+    __sanitizer::atomic_store(&__xray::XRayPatchedFunction,
+                              reinterpret_cast<uintptr_t>(entry),
+                              __sanitizer::memory_order_release);
+    return 1;
   }
+  return 0;
+}
 
-  return XRayPatchingStatus::SUCCESS;
+int __xray_set_customevent_handler(void (*entry)(void *, size_t))
+    XRAY_NEVER_INSTRUMENT {
+  if (__sanitizer::atomic_load(&XRayInitialized,
+                               __sanitizer::memory_order_acquire)) {
+    __sanitizer::atomic_store(&__xray::XRayPatchedCustomEvent,
+                              reinterpret_cast<uintptr_t>(entry),
+                              __sanitizer::memory_order_release);
+    return 1;
+  }
+  return 0;
+}
+
+int __xray_remove_handler() XRAY_NEVER_INSTRUMENT {
+  return __xray_set_handler(nullptr);
+}
+
+int __xray_remove_customevent_handler() XRAY_NEVER_INSTRUMENT {
+  return __xray_set_customevent_handler(nullptr);
+}
+
+XRayPatchingStatus __xray_patch() XRAY_NEVER_INSTRUMENT {
+  return controlPatching(true);
+}
+
+XRayPatchingStatus __xray_unpatch() XRAY_NEVER_INSTRUMENT {
+  return controlPatching(false);
 }
 
 XRayPatchingStatus __xray_patch_function(int32_t FuncId) XRAY_NEVER_INSTRUMENT {
-  return patchFunction(FuncId, true);
+  return mprotectAndPatchFunction(FuncId, true);
 }
 
 XRayPatchingStatus
 __xray_unpatch_function(int32_t FuncId) XRAY_NEVER_INSTRUMENT {
-  return patchFunction(FuncId, false);
+  return mprotectAndPatchFunction(FuncId, false);
 }
 
 int __xray_set_handler_arg1(void (*entry)(int32_t, XRayEntryType, uint64_t)) {