[compiler-rt][XRay] Implement __xray_unpatch() and __xray_remove_handler()

[dberris]

Summary

We also add one test (and the XRay testing infrastructure) to exercise
the patching and unpatching code. This uses the XRay API exported
through the headers as well, installing a custom log handler.

Depends on D23101 for the updated emitted code alignment for the
return/entry sleds.

[rSerge]

As I understood, __xray_remove_handler() is not yet implemented (just declared). Shouldn't it go to this fix too? I think it's about a line of somewhat related code.

[dberris]

Good point. I'll add it in the next update, it should be easy to do.

[dberris]

• Implement __xray_remove_handler() and test it too

[rSerge]

lib/xray/xray_interface.cc

115

It may need a comment on what it does and what the parameter means. Or better, to make the code self-explaining, can it be renamed to e.g. ControlPatching or ControlTracing with a parameter like "bool Enabled"?

[dberris]

• Rename PatchCommon to ControlPatching
• Add comment to ControlPatching

lib/xray/xray_interface.cc

115

Thanks, did both -- added a comment and renamed as you suggested.

[rSerge]

Looks good.

lib/xray/xray_interface.cc

 int __xray_set_handler(void (*entry)(int32_t, XRayEntryType)) {
   if (XRayInitialized.load(std::memory_order_acquire)) {
     __xray::XRayPatchedFunction.store(entry, std::memory_order_release);
     return 1;
   }
   return 0;
 }
 
+int __xray_remove_handler() {
+  if (XRayInitialized.load(std::memory_order_acquire)) {
+    __xray::XRayPatchedFunction.store(nullptr, std::memory_order_release);
+    return 1;
+  }
+  return 0;
+}
+
 std::atomic<bool> XRayPatching{false};
 
 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) : Fn(Fn) {}
   CleanupInvoker(const CleanupInvoker &) = default;
   CleanupInvoker(CleanupInvoker &&) = default;
   CleanupInvoker &operator=(const CleanupInvoker &) = delete;
   CleanupInvoker &operator=(CleanupInvoker &&) = delete;
   ~CleanupInvoker() { Fn(); }
 };
 
 template <class Function> CleanupInvoker<Function> ScopeCleanup(Function Fn) {
   return CleanupInvoker<Function>{Fn};
 }
 
-XRayPatchingStatus __xray_patch() {
+XRayPatchingStatus PatchCommon(bool Patching) {
   if (!XRayInitialized.load(std::memory_order_acquire))
     return XRayPatchingStatus::NOT_INITIALIZED; // Not initialized.
 
   static bool NotPatching = false;
   if (!XRayPatching.compare_exchange_strong(NotPatching, true,
                                             std::memory_order_acq_rel,
                                             std::memory_order_acquire)) {
     return XRayPatchingStatus::ONGOING; // Already patching.
XRayPatchingStatus PatchCommon(bool Patching) {
   // instrumentation map.
   XRaySledMap InstrMap = XRayInstrMap.load(std::memory_order_acquire);
   if (InstrMap.Entries == 0)
     return XRayPatchingStatus::NOT_INITIALIZED;
 
   int32_t FuncId = 1;
   static constexpr uint8_t CallOpCode = 0xe8;
   static constexpr uint16_t MovR10Seq = 0xba41;
+  static constexpr uint16_t Jmp9Seq = 0x09eb;
   static constexpr uint8_t JmpOpCode = 0xe9;
+  static constexpr uint8_t RetOpCode = 0xc3;
   uint64_t CurFun = 0;
   for (std::size_t I = 0; I < InstrMap.Entries; I++) {
     auto Sled = InstrMap.Sleds[I];
     auto F = Sled.Function;
     if (CurFun == 0)
       CurFun = F;
     if (F != CurFun) {
       ++FuncId;
if (Sled.Kind == XRayEntryType::ENTRY) {
           (static_cast<int64_t>(Sled.Address) + 11);
       if (TrampolineOffset < MinOffset || TrampolineOffset > MaxOffset) {
         Report("XRay Entry trampoline (%p) too far from sled (%p); distance = "
                "%ld\n",
                __xray_FunctionEntry, reinterpret_cast<void *>(Sled.Address),
                TrampolineOffset);
         continue;
       }
+      if (Patching) {
         *reinterpret_cast<uint32_t *>(Sled.Address + 2) = FuncId;
         *reinterpret_cast<uint8_t *>(Sled.Address + 6) = CallOpCode;
         *reinterpret_cast<uint32_t *>(Sled.Address + 7) = TrampolineOffset;
         std::atomic_store_explicit(
             reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), MovR10Seq,
             std::memory_order_release);
+      } else {
+        std::atomic_store_explicit(
+            reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), Jmp9Seq,
+            std::memory_order_release);
+        // FIXME: Write out the nops still?
+      }
     }
 
     if (Sled.Kind == XRayEntryType::EXIT) {
       // FIXME: Implement this in a more extensible manner, per-platform.
       // Here we do the dance of replacing the following sled:
       //
       // xray_sled_n:
       //   ret
if (Sled.Kind == XRayEntryType::EXIT) {
           (static_cast<int64_t>(Sled.Address) + 11);
       if (TrampolineOffset < MinOffset || TrampolineOffset > MaxOffset) {
         Report("XRay Exit trampoline (%p) too far from sled (%p); distance = "
                "%ld\n",
                __xray_FunctionExit, reinterpret_cast<void *>(Sled.Address),
                TrampolineOffset);
         continue;
       }
+      if (Patching) {
         *reinterpret_cast<uint32_t *>(Sled.Address + 2) = FuncId;
         *reinterpret_cast<uint8_t *>(Sled.Address + 6) = JmpOpCode;
         *reinterpret_cast<uint32_t *>(Sled.Address + 7) = TrampolineOffset;
         std::atomic_store_explicit(
             reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), MovR10Seq,
             std::memory_order_release);
+      } else {
+        std::atomic_store_explicit(
+            reinterpret_cast<std::atomic<uint8_t> *>(Sled.Address), RetOpCode,
+            std::memory_order_release);
+        // FIXME: Write out the nops still?
+      }
     }
   }
   XRayPatching.store(false, std::memory_order_release);
   PatchingSuccess = true;
   return XRayPatchingStatus::SUCCESS;
 }
+
+XRayPatchingStatus __xray_patch() { return PatchCommon(true); }
+
+XRayPatchingStatus __xray_unpatch() { return PatchCommon(false); }

test/CMakeLists.txt

if(COMPILER_RT_HAS_SAFESTACK)
     add_subdirectory(safestack)
   endif()
   if(COMPILER_RT_HAS_ESAN)
     add_subdirectory(esan)
   endif()
   if(COMPILER_RT_HAS_SCUDO)
     add_subdirectory(scudo)
   endif()
+  if(COMPILER_RT_HAS_XRAY)
+    add_subdirectory(xray)
+  endif()
 endif()
 
 if(COMPILER_RT_STANDALONE_BUILD)
   # Now that we've traversed all the directories and know all the lit testsuites,
   # introduce a rule to run to run all of them.
   get_property(LLVM_LIT_TESTSUITES GLOBAL PROPERTY LLVM_LIT_TESTSUITES)
   get_property(LLVM_LIT_DEPENDS GLOBAL PROPERTY LLVM_LIT_DEPENDS)
   add_lit_target(check-all
     "Running all regression tests"
     ${LLVM_LIT_TESTSUITES}
     DEPENDS ${LLVM_LIT_DEPENDS})
 endif()

test/xray/CMakeLists.txt

+set(XRAY_LIT_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR})
+
+set(XRAY_TESTSUITES)
+
+set(XRAY_TEST_DEPS ${SANITIZER_COMMON_LIT_TEST_DEPS})
+if(NOT COMPILER_RT_STANDALONE_BUILD)
+  list(APPEND XRAY_TEST_DEPS xray)
+endif()
+
+set(XRAY_TEST_ARCH ${XRAY_SUPPORTED_ARCH})
+foreach(arch ${XRAY_TEST_ARCH})
+  set(XRAY_TEST_TARGET_ARCH ${arch})
+  string(TOLOWER "-${arch}-${OS_NAME}" XRAY_TEST_CONFIG_SUFFIX)
+
+  if(ANDROID OR ${arch} MATCHES "arm|aarch64")
+    # This is only true if we are cross-compiling.
+    # Build all tests with host compiler and use host tools.
+    set(XRAY_TEST_TARGET_CC ${COMPILER_RT_TEST_COMPILER})
+    set(XRAY_TEST_TARGET_CFLAGS ${COMPILER_RT_TEST_COMPILER_CFLAGS})
+  else()
+    get_target_flags_for_arch(${arch} XRAY_TEST_TARGET_CFLAGS)
+    string(REPLACE ";" " " XRAY_TEST_TARGET_CFLAGS "${XRAY_TEST_TARGET_CFLAGS}")
+  endif()
+
+  string(TOUPPER ${arch} ARCH_UPPER_CASE)
+  set(CONFIG_NAME ${ARCH_UPPER_CASE}${OS_NAME}Config)
+
+  configure_lit_site_cfg(
+    ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.in
+    ${CMAKE_CURRENT_BINARY_DIR}/${CONFIG_NAME}/lit.site.cfg)
+  list(APPEND XRAY_TESTSUITES ${CMAKE_CURRENT_BINARY_DIR}/${CONFIG_NAME})
+endforeach()
+
+add_lit_testsuite(check-xray "Running the XRay tests"
+  ${XRAY_TESTSUITES}
+  DEPENDS ${XRAY_TEST_DEPS})
+set_target_properties(check-xray PROPERTIES FOLDER "Compiler-RT Misc")

test/xray/TestCases/Linux/patching-unpatching.cc

+// Check that we can patch and un-patch on demand, and that logging gets invoked
+// appropriately.
+//
+// RUN: %clangxx_xray -fxray-instrument -std=c++11 %s -o %t
+// RUN: XRAY_OPTIONS="patch_premain=false" %run %t 2>&1 | FileCheck %s
+
+#include "xray/xray_interface.h"
+
+#include <cstdio>
+
+bool called = false;
+
+void test_handler(int32_t fid, XRayEntryType type) {
+  printf("called: %d, type=%d\n", fid, static_cast<int32_t>(type));
+  called = true;
+}
+
+[[clang::xray_always_instrument]] void always_instrument() {
+  printf("always instrumented called\n");
+}
+
+int main() {
+  __xray_set_handler(test_handler);
+  always_instrument();
+  // CHECK: always instrumented called
+  auto status = __xray_patch();
+  printf("patching status: %d\n", static_cast<int32_t>(status));
+  // CHECK-NEXT: patching status: 1
+  always_instrument();
+  // CHECK-NEXT: called: {{.*}}, type=0
+  // CHECK-NEXT: always instrumented called
+  // CHECK-NEXT: called: {{.*}}, type=1
+  status = __xray_unpatch();
+  printf("patching status: %d\n", static_cast<int32_t>(status));
+  // CHECK-NEXT: patching status: 1
+  always_instrument();
+  // CHECK-NEXT: always instrumented called
+  status = __xray_patch();
+  printf("patching status: %d\n", static_cast<int32_t>(status));
+  // CHECK-NEXT: patching status: 1
+  __xray_remove_handler();
+  always_instrument();
+  // CHECK-NEXT: always instrumented called
+  status = __xray_unpatch();
+  printf("patching status: %d\n", static_cast<int32_t>(status));
+  // CHECK-NEXT: patching status: 1
+}

test/xray/lit.cfg

+# -*- Python -*-
+
+import os
+
+# Setup config name.
+config.name = 'XRay' + config.name_suffix
+
+# Setup source root.
+config.test_source_root = os.path.dirname(__file__)
+
+# Setup default compiler flags use with -fxray-instrument option.
+clang_xray_cflags = (['-fxray-instrument', config.target_cflags])
+clang_xray_cxxflags = config.cxx_mode_flags + clang_xray_cflags
+
+
+def build_invocation(compile_flags):
+  return ' ' + ' '.join([config.clang] + compile_flags) + ' '
+
+# Setup substitutions.
+config.substitutions.append(
+    ('%clang ', build_invocation([config.target_cflags])))
+config.substitutions.append(
+    ('%clangxx ',
+     build_invocation(config.cxx_mode_flags + [config.target_cflags])))
+config.substitutions.append(
+    ('%clang_xray ', build_invocation(clang_xray_cflags)))
+config.substitutions.append(
+    ('%clangxx_xray', build_invocation(clang_xray_cxxflags)))
+
+# Default test suffixes.
+config.suffixes = ['.c', '.cc', '.cpp']
+
+if config.host_os not in ['Linux'] or config.host_arch.find('64') == -1:
+  config.unsupported = True

test/xray/lit.site.cfg.in

+@LIT_SITE_CFG_IN_HEADER@
+
+# Tool-specific config options.
+config.name_suffix = "@XRAY_TEST_CONFIG_SUFFIX@"
+config.xray_lit_source_dir = "@XRAY_LIT_SOURCE_DIR@"
+config.target_cflags = "@XRAY_TEST_TARGET_CFLAGS@"
+config.target_arch = "@XRAY_TEST_TARGET_ARCH@"
+
+# Load common config for all compiler-rt lit tests
+lit_config.load_config(config, "@COMPILER_RT_BINARY_DIR@/test/lit.common.configured")
+
+# Load tool-specific config that would do the real work.
+lit_config.load_config(config, "@XRAY_LIT_SOURCE_DIR@/lit.cfg")