1 files changed, 27 insertions, 13 deletions

diff --git a/lib/xray/xray_x86_64.cc b/lib/xray/xray_x86_64.cc
index e34806fa1cea..e17f00ac3a62 100644
--- a/lib/xray/xray_x86_64.cc
+++ b/lib/xray/xray_x86_64.cc
@@ -76,6 +76,7 @@ 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;
@@ -207,8 +208,10 @@ 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:
@@ -216,24 +219,35 @@ bool patchCustomEvent(const bool Enable, const uint32_t FuncId,
   //   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'.
+  //
+  // ---
   //
-  // The "unpatch" should just turn the 'nopw' back to a 'jmp +24'.
+  // In Version 1:
+  //
+  //   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 {
-    std::atomic_store_explicit(
-        reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), Jmp20Seq,
-        std::memory_order_release);
-  }
+    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;
 }
 
@@ -244,9 +258,9 @@ bool probeRequiredCPUFeatures() XRAY_NEVER_INSTRUMENT {
 
   // We check whether rdtscp support is enabled. According to the x86_64 manual,
   // level should be set at 0x80000001, and we should have a look at bit 27 in
-  // EDX. That's 0x8000000 (or 1u << 26).
+  // EDX. That's 0x8000000 (or 1u << 27).
   __get_cpuid(0x80000001, &EAX, &EBX, &ECX, &EDX);
-  if (!(EDX & (1u << 26))) {
+  if (!(EDX & (1u << 27))) {
     Report("Missing rdtscp support.\n");
     return false;
   }