[XRay][compiler-rt] Support sled versioning for custom event sleds

[dberris]

Summary

This change introduces versions to the instrumentation map entries we
emit for XRay instrumentaiton points. The status quo for the version is
currently set to 0 (as emitted by the LLVM back-end), and versions will
count up to 255 (unsigned char).

This change is in preparation for supporting the newer version of the
custom event sleds that will be emitted by the LLVM compiler.

While we're here, we take the opportunity to stash more registers and
align the stack properly in the __xray_CustomEvent trampoline.

[kpw]

lib/xray/xray_x86_64.cc

77

If 0x09eb is Jmp9Seq and 0x14eb is Jmp20Seq, I would expect Jmp16Seq to be 0x10eb.

0x0feb would match this pattern as Jmp15Seq. I think the confusion might come from address being relative to the *next* instruction.

test/xray/TestCases/Linux/custom-event-logging.cc

5

Is the position independent code stuff related to the sled change or just along for the ride?

[dberris]

Address comments.

[dberris]

test/xray/TestCases/Linux/custom-event-logging.cc

5

Ah, right -- it's along for the ride, because the dependent change has to work with PIC and PIE changes.

[dberris]

Rebase

[dberris]

Update comments to be more accurate.

[dberris]

• fixup: address comments, stash more registers in __xray_CustomEvent, and add a test with a function that requires specific alignment

[pcc]

lib/xray/xray_trampoline_x86_64.S

205

Without backend support I'd expect you to need to realign the stack here. You can see how to do that by compiling any C function with -mstackrealign.

test/xray/TestCases/Linux/custom-event-handler-alignment.cc

29

Probably want to copy lines 19 and 20 from the function foo above in order to test that the stack is properly aligned in the event handler.

[dberris]

• fixup: align the stack before calling the custom handler function

[dberris]

lib/xray/xray_trampoline_x86_64.S

205

Awesome stuff, thanks @pcc!

[dberris]

Retitle to include changes to the stack alignment for the __xray_CustomEvent trampoline.

[pcc]

LGTM

lib/xray/xray_interface_internal.h

 extern "C" {
 
 struct XRaySledEntry {
 #if SANITIZER_WORDSIZE == 64
   uint64_t Address;
   uint64_t Function;
   unsigned char Kind;
   unsigned char AlwaysInstrument;
-  unsigned char Padding[14]; // Need 32 bytes
+  unsigned char Version;
+  unsigned char Padding[13]; // Need 32 bytes
 #elif SANITIZER_WORDSIZE == 32
   uint32_t Address;
   uint32_t Function;
   unsigned char Kind;
   unsigned char AlwaysInstrument;
-  unsigned char Padding[6]; // Need 16 bytes
+  unsigned char Version;
+  unsigned char Padding[5]; // Need 16 bytes
 #else
 #error "Unsupported word size."
 #endif
 };
 
 struct XRayFunctionSledIndex {
   const XRaySledEntry* Begin;
   const XRaySledEntry* End;

lib/xray/xray_trampoline_x86_64.S

 
 //===----------------------------------------------------------------------===//
 
   .global __xray_CustomEvent
   .align 16, 0x90
   .type __xray_CustomEvent,@function
 __xray_CustomEvent:
   .cfi_startproc
-  subq $16, %rsp
-  .cfi_def_cfa_offset 24
-  movq %rbp, 8(%rsp)
-  movq %rax, 0(%rsp)
+  SAVE_REGISTERS
 
   // We take two arguments to this trampoline, which should be in rdi and rsi
   // already. We also make sure that we stash %rax because we use that register
   // to call the logging handler.
   movq _ZN6__xray22XRayPatchedCustomEventE(%rip), %rax
   testq %rax,%rax
   je .LcustomEventCleanup
 
   // At this point we know that rcx and rdx already has the data, so we just
-  // call the logging handler.
+  // call the logging handler, after aligning the stack to a 16-byte boundary.
+  // The approach we're taking here uses additional stack space to stash the
+  // stack pointer twice before aligning the pointer to 16-bytes. If the stack
+  // was 8-byte aligned, it will become 16-byte aligned -- when restoring the
+  // pointer, we can always look -8 bytes from the current position to get
+  // either of the values we've stashed in the first place.
+  pushq %rsp
+  pushq (%rsp)
+  andq $-0x10, %rsp
   callq *%rax
+  movq 8(%rsp), %rsp
 
 .LcustomEventCleanup:
-  movq 0(%rsp), %rax
-  movq 8(%rsp), %rbp
-  addq $16, %rsp
+  RESTORE_REGISTERS
   .cfi_def_cfa_offset 8
   retq
 
 .Ltmp8:
   .size __xray_CustomEvent, .Ltmp8-__xray_CustomEvent
   .cfi_endproc
 
 NO_EXEC_STACK_DIRECTIVE

lib/xray/xray_x86_64.cc

uint64_t getTSCFrequency() XRAY_NEVER_INSTRUMENT {
   } else {
     Report("Unable to determine CPU frequency for TSC accounting.\n");
   }
   return TSCFrequency == -1 ? 0 : static_cast<uint64_t>(TSCFrequency);
 }
 
 static constexpr uint8_t CallOpCode = 0xe8;
 static constexpr uint16_t MovR10Seq = 0xba41;
 static constexpr uint16_t Jmp9Seq = 0x09eb;
 static constexpr uint16_t Jmp20Seq = 0x14eb;
+static constexpr uint16_t Jmp15Seq = 0x0feb;
 static constexpr uint8_t JmpOpCode = 0xe9;
 static constexpr uint8_t RetOpCode = 0xc3;
 static constexpr uint16_t NopwSeq = 0x9066;
 
 static constexpr int64_t MinOffset{std::numeric_limits<int32_t>::min()};
 static constexpr int64_t MaxOffset{std::numeric_limits<int32_t>::max()};
 
 bool patchFunctionEntry(const bool Enable, const uint32_t FuncId,
bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
   }
   return true;
 }
 
 bool patchCustomEvent(const bool Enable, const uint32_t FuncId,
                       const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
   // Here we do the dance of replacing the following sled:
   //
+  // In Version 0:
+  //
   // xray_sled_n:
-  //   jmp +19          // 2 bytes
+  //   jmp +20          // 2 bytes
   //   ...
   //
   // With the following:
   //
   //   nopw             // 2 bytes*
   //   ...
   //
-  // We need to do this in the following order:
   //
-  // 1. Overwrite the 5-byte nop with the call (relative), where (relative) is
-  //    the relative offset to the __xray_CustomEvent trampoline.
-  // 2. Do a two-byte atomic write over the 'jmp +24' to turn it into a 'nopw'.
-  //    This allows us to "enable" this code once the changes have committed.
+  // The "unpatch" should just turn the 'nopw' back to a 'jmp +20'.
+  //
+  // ---
+  //
+  // In Version 1:
   //
-  // The "unpatch" should just turn the 'nopw' back to a 'jmp +24'.
+  //   The jump offset is now 15 bytes (0x0f), so when restoring the nopw back
+  //   to a jmp, use 15 bytes instead.
   //
   if (Enable) {
     std::atomic_store_explicit(
         reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), NopwSeq,
         std::memory_order_release);
   } else {
+    switch (Sled.Version) {
+    case 1:
+      std::atomic_store_explicit(
+          reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), Jmp15Seq,
+          std::memory_order_release);
+      break;
+    case 0:
+    default:
       std::atomic_store_explicit(
           reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), Jmp20Seq,
           std::memory_order_release);
+      break;
+    }
     }
   return false;
 }
 
 // We determine whether the CPU we're running on has the correct features we
 // need. In x86_64 this will be rdtscp support.
 bool probeRequiredCPUFeatures() XRAY_NEVER_INSTRUMENT {
   unsigned int EAX, EBX, ECX, EDX;
 

test/xray/TestCases/Linux/custom-event-handler-alignment.cc

+// Make sure we're aligning the stack properly when lowering the custom event
+// calls.
+//
+// RUN: %clangxx_xray -std=c++11 %s -o %t
+// RUN: XRAY_OPTIONS="patch_premain=false verbosity=1 xray_naive_log=false" \
+// RUN:     %run %t 2>&1
+// REQUIRES: x86_64-linux
+// REQUIRES: built-in-llvm-tree
+#include <xmmintrin.h>
+#include <stdio.h>
+#include "xray/xray_interface.h"
+
+[[clang::xray_never_instrument]] __attribute__((weak)) __m128 f(__m128 *i) {
+  return *i;
+}
+
+[[clang::xray_always_instrument]] void foo() {
+  __xray_customevent(0, 0);
+  __m128 v = {};
+  f(&v);
+}
+
+[[clang::xray_always_instrument]] void bar() {
+  __xray_customevent(0, 0);
+}
+
+void printer(void* ptr, size_t size) {
+  printf("handler called\n");
+  __m128 v = {};
+  f(&v);
+}
+
+int main(int argc, char* argv[]) {
+  __xray_set_customevent_handler(printer);
+  __xray_patch();
+  foo();  // CHECK: handler called
+  bar();  // CHECK: handler called
+  __xray_unpatch();
+  __xray_remove_customevent_handler();
+  foo();
+  bar();
+}

test/xray/TestCases/Linux/custom-event-logging.cc

 // Use the clang feature for custom xray event logging.
 //
 // RUN: %clangxx_xray -std=c++11 %s -o %t
 // RUN: XRAY_OPTIONS="patch_premain=false verbosity=1 xray_naive_log=false xray_logfile_base=custom-event-logging.xray-" %run %t 2>&1 | FileCheck %s
+// RUN: %clangxx_xray -std=c++11 -fpic -fpie %s -o %t
+// RUN: XRAY_OPTIONS="patch_premain=false verbosity=1 xray_naive_log=false xray_logfile_base=custom-event-logging.xray-" %run %t 2>&1 | FileCheck %s
 // FIXME: Support this in non-x86_64 as well
 // REQUIRES: x86_64-linux
 // REQUIRES: built-in-llvm-tree
 #include <cstdio>
 #include "xray/xray_interface.h"
 
 [[clang::xray_always_instrument]] void foo() {
   static constexpr char CustomLogged[] = "hello custom logging!";