[XRay] [compiler-rt] fix heap overflow by computing record pointers correctly

[pelikan]

Summary

While there, unify InMemoryRawLog and InMemoryRawLogWithArg's coding style:

• swap libc's memcpy(3) for sanitizer's internal memcpy
• use basic pointer arithmetics to compute offsets from the first record entry in the pre-allocated buffer, which is always the appropriate type for the given function
• lose the local variable references as the TLD.* names fit just as well

[dberris]

Good catch, thanks!

lib/xray/xray_inmemory_log.cc

ThreadLocalData &getThreadLocalData() XRAY_NEVER_INSTRUMENT {
   }();
   return TLD;
 }
 
 template <class RDTSC>
 void InMemoryRawLog(int32_t FuncId, XRayEntryType Type,
                     RDTSC ReadTSC) XRAY_NEVER_INSTRUMENT {
   auto &TLD = getThreadLocalData();
-  auto &InMemoryBuffer = TLD.InMemoryBuffer;
   int Fd = getGlobalFd();
   if (Fd == -1)
     return;
 
   // Use a simple recursion guard, to handle cases where we're already logging
   // and for one reason or another, this function gets called again in the same
   // thread.
   if (RecursionGuard)
void InMemoryRawLog(int32_t FuncId, XRayEntryType Type,
   // the exit record is higher than the threshold.
   __xray::XRayRecord R;
   R.RecordType = RecordTypes::NORMAL;
   R.CPU = CPU;
   R.TSC = TSC;
   R.TId = TLD.TID;
   R.Type = Type;
   R.FuncId = FuncId;
-  auto EntryPtr = static_cast<char *>(InMemoryBuffer) +
-                  (sizeof(__xray::XRayRecord) * TLD.BufferOffset);
-  __sanitizer::internal_memcpy(EntryPtr, &R, sizeof(R));
+  auto FirstEntry = reinterpret_cast<__xray::XRayRecord *>(TLD.InMemoryBuffer);
+  __sanitizer::internal_memcpy(FirstEntry + TLD.BufferOffset, &R, sizeof(R));
   if (++TLD.BufferOffset == TLD.BufferSize) {
     __sanitizer::SpinMutexLock L(&LogMutex);
-    auto RecordBuffer = reinterpret_cast<__xray::XRayRecord *>(InMemoryBuffer);
-    retryingWriteAll(Fd, reinterpret_cast<char *>(RecordBuffer),
-                     reinterpret_cast<char *>(RecordBuffer + TLD.BufferOffset));
+    retryingWriteAll(Fd, reinterpret_cast<char *>(FirstEntry),
+                     reinterpret_cast<char *>(FirstEntry + TLD.BufferOffset));
     TLD.BufferOffset = 0;
     TLD.StackEntries = 0;
   }
 }
 
 template <class RDTSC>
 void InMemoryRawLogWithArg(int32_t FuncId, XRayEntryType Type, uint64_t Arg1,
                            RDTSC ReadTSC) XRAY_NEVER_INSTRUMENT {
   auto &TLD = getThreadLocalData();
-  auto &InMemoryBuffer = TLD.InMemoryBuffer;
-  auto &Offset = TLD.BufferOffset;
+  auto FirstEntry =
+      reinterpret_cast<__xray::XRayArgPayload *>(TLD.InMemoryBuffer);
   const auto &BuffLen = TLD.BufferSize;
   int Fd = getGlobalFd();
   if (Fd == -1)
     return;
 
   // First we check whether there's enough space to write the data consecutively
   // in the thread-local buffer. If not, we first flush the buffer before
   // attempting to write the two records that must be consecutive.
-  if (Offset + 2 > BuffLen) {
+  if (TLD.BufferOffset + 2 > BuffLen) {
     __sanitizer::SpinMutexLock L(&LogMutex);
-    auto RecordBuffer = reinterpret_cast<__xray::XRayRecord *>(InMemoryBuffer);
-    retryingWriteAll(Fd, reinterpret_cast<char *>(RecordBuffer),
-                     reinterpret_cast<char *>(RecordBuffer + Offset));
-    Offset = 0;
+    retryingWriteAll(Fd, reinterpret_cast<char *>(FirstEntry),
+                     reinterpret_cast<char *>(FirstEntry + TLD.BufferOffset));
+    TLD.BufferOffset = 0;
     TLD.StackEntries = 0;
   }
 
   // Then we write the "we have an argument" record.
   InMemoryRawLog(FuncId, Type, ReadTSC);
 
   if (RecursionGuard)
     return;
   RecursionGuard = true;
   auto ExitGuard = __sanitizer::at_scope_exit([] { RecursionGuard = false; });
 
   // And from here on write the arg payload.
   __xray::XRayArgPayload R;
   R.RecordType = RecordTypes::ARG_PAYLOAD;
   R.FuncId = FuncId;
   R.TId = TLD.TID;
   R.Arg = Arg1;
-  auto EntryPtr =
-      &reinterpret_cast<__xray::XRayArgPayload *>(&InMemoryBuffer)[Offset];
-  std::memcpy(EntryPtr, &R, sizeof(R));
-  if (++Offset == BuffLen) {
+  __sanitizer::internal_memcpy(FirstEntry + TLD.BufferOffset, &R, sizeof(R));
+  if (++TLD.BufferOffset == BuffLen) {
     __sanitizer::SpinMutexLock L(&LogMutex);
-    auto RecordBuffer = reinterpret_cast<__xray::XRayRecord *>(InMemoryBuffer);
-    retryingWriteAll(Fd, reinterpret_cast<char *>(RecordBuffer),
-                     reinterpret_cast<char *>(RecordBuffer + Offset));
-    Offset = 0;
+    retryingWriteAll(Fd, reinterpret_cast<char *>(FirstEntry),
+                     reinterpret_cast<char *>(FirstEntry + TLD.BufferOffset));
+    TLD.BufferOffset = 0;
     TLD.StackEntries = 0;
   }
 }
 
 void basicLoggingHandleArg0RealTSC(int32_t FuncId,
                                    XRayEntryType Type) XRAY_NEVER_INSTRUMENT {
   InMemoryRawLog(FuncId, Type, __xray::readTSC);
 }