1 files changed, 32 insertions, 14 deletions

diff --git a/lib/xray/xray_x86_64.cc b/lib/xray/xray_x86_64.cc
index 3ee91896c6e0..8c2a4e313e3a 100644
--- a/lib/xray/xray_x86_64.cc
+++ b/lib/xray/xray_x86_64.cc
@@ -1,6 +1,8 @@
+#include "cpuid.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "xray_defs.h"
 #include "xray_interface_internal.h"
+
 #include <atomic>
 #include <cstdint>
 #include <errno.h>
@@ -42,9 +44,9 @@ static bool readValueFromFile(const char *Filename,
   ssize_t BytesRead;
   bool Success;
   std::tie(BytesRead, Success) = retryingReadSome(Fd, Line, Line + BufSize);
+  close(Fd);
   if (!Success)
     return false;
-  close(Fd);
   char *End = nullptr;
   long long Tmp = internal_simple_strtoll(Line, &End, 10);
   bool Result = false;
@@ -55,19 +57,19 @@ static bool readValueFromFile(const char *Filename,
   return Result;
 }
 
-uint64_t cycleFrequency() XRAY_NEVER_INSTRUMENT {
-  long long CPUFrequency = -1;
+uint64_t getTSCFrequency() XRAY_NEVER_INSTRUMENT {
+  long long TSCFrequency = -1;
   if (readValueFromFile("/sys/devices/system/cpu/cpu0/tsc_freq_khz",
-                        &CPUFrequency)) {
-    CPUFrequency *= 1000;
+                        &TSCFrequency)) {
+    TSCFrequency *= 1000;
   } else if (readValueFromFile(
-      "/sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_max_freq",
-      &CPUFrequency)) {
-    CPUFrequency *= 1000;
+                 "/sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_max_freq",
+                 &TSCFrequency)) {
+    TSCFrequency *= 1000;
   } else {
     Report("Unable to determine CPU frequency for TSC accounting.\n");
   }
-  return CPUFrequency == -1 ? 0 : static_cast<uint64_t>(CPUFrequency);
+  return TSCFrequency == -1 ? 0 : static_cast<uint64_t>(TSCFrequency);
 }
 
 static constexpr uint8_t CallOpCode = 0xe8;
@@ -80,7 +82,8 @@ 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,
-                        const XRaySledEntry &Sled) XRAY_NEVER_INSTRUMENT {
+                        const XRaySledEntry &Sled,
+                        void (*Trampoline)()) XRAY_NEVER_INSTRUMENT {
   // Here we do the dance of replacing the following sled:
   //
   // xray_sled_n:
@@ -101,13 +104,12 @@ bool patchFunctionEntry(const bool Enable, const uint32_t FuncId,
   // 4. Do an atomic write over the jmp instruction for the "mov r10d"
   // opcode and first operand.
   //
-  // Prerequisite is to compute the relative offset to the
-  // __xray_FunctionEntry function's address.
-  int64_t TrampolineOffset = reinterpret_cast<int64_t>(__xray_FunctionEntry) -
+  // Prerequisite is to compute the relative offset to the trampoline's address.
+  int64_t TrampolineOffset = reinterpret_cast<int64_t>(Trampoline) -
                              (static_cast<int64_t>(Sled.Address) + 11);
   if (TrampolineOffset < MinOffset || TrampolineOffset > MaxOffset) {
     Report("XRay Entry trampoline (%p) too far from sled (%p)\n",
-           __xray_FunctionEntry, reinterpret_cast<void *>(Sled.Address));
+           Trampoline, reinterpret_cast<void *>(Sled.Address));
     return false;
   }
   if (Enable) {
@@ -199,4 +201,20 @@ bool patchFunctionTailExit(const bool Enable, const uint32_t FuncId,
   return true;
 }
 
+// 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;
+
+  // 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).
+  __get_cpuid(0x80000001, &EAX, &EBX, &ECX, &EDX);
+  if (!(EDX & (1u << 26))) {
+    Report("Missing rdtscp support.\n");
+    return false;
+  }
+  return true;
+}
+
 } // namespace __xray