[Lsan] Use fp registers to search for pointers

X86 can use xmm registers for pointers operations. e.g. for std::swap.
I don't know yet if it's possible on other platforms.

NT_X86_XSTATE includes all registers from NT_FPREGSET so
the latter used only if the former is not available. I am not sure how
reasonable to expect that but LLD has such fallback in
NativeRegisterContextLinux_x86_64::ReadFPR.

Reviewed By: morehouse

Differential Revision: https://reviews.llvm.org/D87754

compiler-rt/lib/sanitizer_common/sanitizer_stoptheworld_linux_libcdep.cpp

@@ -485,6 +485,9 @@ typedef user_regs_struct regs_struct;
 #else
 #define REG_SP rsp
 #endif
+#define ARCH_IOVEC_FOR_GETREGSET
+// Compiler may use FP registers to store pointers.
+static constexpr uptr kExtraRegs[] = {NT_X86_XSTATE, NT_FPREGSET};
 
 #elif defined(__powerpc__) || defined(__powerpc64__)
 typedef pt_regs regs_struct;
@@ -501,11 +504,13 @@ typedef struct user regs_struct;
 #elif defined(__aarch64__)
 typedef struct user_pt_regs regs_struct;
 #define REG_SP sp
+static constexpr uptr kExtraRegs[] = {};
 #define ARCH_IOVEC_FOR_GETREGSET
 
 #elif defined(__s390__)
 typedef _user_regs_struct regs_struct;
 #define REG_SP gprs[15]
+static constexpr uptr kExtraRegs[] = {};
 #define ARCH_IOVEC_FOR_GETREGSET
 
 #else
@@ -535,22 +540,56 @@ void SuspendedThreadsListLinux::Append(tid_t tid) {
 PtraceRegistersStatus SuspendedThreadsListLinux::GetRegistersAndSP(
     uptr index, InternalMmapVector<uptr> *buffer, uptr *sp) const {
   pid_t tid = GetThreadID(index);
-  regs_struct regs;
+  constexpr uptr uptr_sz = sizeof(uptr);
   int pterrno;
 #ifdef ARCH_IOVEC_FOR_GETREGSET
-  struct iovec regset_io;
-  regset_io.iov_base = &regs;
-  regset_io.iov_len = sizeof(regs_struct);
-  bool isErr = internal_iserror(internal_ptrace(PTRACE_GETREGSET, tid,
-                                (void*)NT_PRSTATUS, (void*)&regset_io),
-                                &pterrno);
+  auto append = [&](uptr regset) {
+    uptr size = buffer->size();
+    // NT_X86_XSTATE requires 64bit alignment.
+    uptr size_up = RoundUpTo(size, 8 / uptr_sz);
+    buffer->reserve(Max<uptr>(1024, size_up));
+    struct iovec regset_io;
+    for (;; buffer->resize(buffer->capacity() * 2)) {
+      buffer->resize(buffer->capacity());
+      uptr available_bytes = (buffer->size() - size_up) * uptr_sz;
+      regset_io.iov_base = buffer->data() + size_up;
+      regset_io.iov_len = available_bytes;
+      bool fail =
+          internal_iserror(internal_ptrace(PTRACE_GETREGSET, tid,
+                                           (void *)regset, (void *)&regset_io),
+                           &pterrno);
+      if (fail) {
+        VReport(1, "Could not get regset %p from thread %d (errno %d).\n",
+                regset, tid, pterrno);
+        buffer->resize(size);
+        return false;
+      }
+
+      // Far enough from the buffer size, no need to resize and repeat.
+      if (regset_io.iov_len + 64 < available_bytes)
+        break;
+    }
+    buffer->resize(size_up + RoundUpTo(regset_io.iov_len, uptr_sz) / uptr_sz);
+    return true;
+  };
+
+  buffer->clear();
+  bool fail = !append(NT_PRSTATUS);
+  if (!fail) {
+    // Accept the first available and do not report errors.
+    for (uptr regs : kExtraRegs)
+      if (append(regs))
+        break;
+  }
 #else
-  bool isErr = internal_iserror(internal_ptrace(PTRACE_GETREGS, tid, nullptr,
-                                &regs), &pterrno);
-#endif
-  if (isErr) {
+  buffer->resize(RoundUpTo(sizeof(regs_struct), uptr_sz) / uptr_sz);
+  bool fail = internal_iserror(
+      internal_ptrace(PTRACE_GETREGS, tid, nullptr, buffer->data()), &pterrno);
+  if (fail)
     VReport(1, "Could not get registers from thread %d (errno %d).\n", tid,
             pterrno);
+#endif
+  if (fail) {
     // ESRCH means that the given thread is not suspended or already dead.
     // Therefore it's unsafe to inspect its data (e.g. walk through stack) and
     // we should notify caller about this.
@@ -558,9 +597,7 @@ PtraceRegistersStatus SuspendedThreadsListLinux::GetRegistersAndSP(
                             : REGISTERS_UNAVAILABLE;
   }
 
-  *sp = regs.REG_SP;
-  buffer->resize(RoundUpTo(sizeof(regs), sizeof(uptr)) / sizeof(uptr));
-  internal_memcpy(buffer->data(), &regs, sizeof(regs));
+  *sp = reinterpret_cast<regs_struct *>(buffer->data())[0].REG_SP;
   return REGISTERS_AVAILABLE;
 }
 

compiler-rt/test/lsan/TestCases/use_registers.cpp

@@ -16,6 +16,9 @@ extern "C"
 void *registers_thread_func(void *arg) {
   int *sync = reinterpret_cast<int *>(arg);
   void *p = malloc(1337);
+  print_address("Test alloc: ", 1, p);
+  fflush(stderr);
+
   // To store the pointer, choose a register which is unlikely to be reused by
   // a function call.
 #if defined(__i386__)
@@ -50,8 +53,6 @@ void *registers_thread_func(void *arg) {
 #else
 #error "Test is not supported on this architecture."
 #endif
-  print_address("Test alloc: ", 1, p);
-  fflush(stderr);
   __sync_fetch_and_xor(sync, 1);
   while (true)
     sched_yield();

compiler-rt/test/lsan/TestCases/use_registers_extra.cpp

@@ -0,0 +1,61 @@
+// Test that registers of running threads are included in the root set.
+// RUN: LSAN_BASE="report_objects=1:use_stacks=0"
+// RUN: %clangxx_lsan -pthread %s -o %t
+// RUN: %env_lsan_opts=$LSAN_BASE:"use_registers=0" not %run %t 2>&1 | FileCheck %s
+// RUN: %env_lsan_opts=$LSAN_BASE:"use_registers=1" %run %t 2>&1
+// RUN: %env_lsan_opts="" %run %t 2>&1
+
+// FIXME: Support more platforms.
+// REQUIRES: x86-target-arch
+
+#include "sanitizer_common/print_address.h"
+#include <assert.h>
+#include <pthread.h>
+#include <sched.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+extern "C" void *registers_thread_func(void *arg) {
+  int *sync = reinterpret_cast<int *>(arg);
+  void *p = malloc(1337);
+  print_address("Test alloc: ", 1, p);
+  fflush(stderr);
+
+  // To store the pointer, choose a register which is unlikely to be reused by
+  // a function call.
+#if defined(__i386__)
+  asm(R"(
+    movd %0, %%xmm0
+    mov $0, %0
+  )"
+      :
+      : "r"(p));
+#elif defined(__x86_64__)
+  asm(R"(
+    movq %0, %%xmm0
+    mov $0, %0
+  )"
+      :
+      : "r"(p));
+#else
+#error "Test is not supported on this architecture."
+#endif
+
+  __sync_fetch_and_xor(sync, 1);
+  while (true)
+    sched_yield();
+}
+
+int main() {
+  int sync = 0;
+  pthread_t thread_id;
+  int res = pthread_create(&thread_id, 0, registers_thread_func, &sync);
+  assert(res == 0);
+  while (!__sync_fetch_and_xor(&sync, 0))
+    sched_yield();
+  return 0;
+}
+// CHECK: Test alloc: [[ADDR:0x[0-9,a-f]+]]
+// CHECK: LeakSanitizer: detected memory leaks
+// CHECK: [[ADDR]] (1337 bytes)
+// CHECK: SUMMARY: {{(Leak|Address)}}Sanitizer: