[TSan] Improve handling of stack pointer mangling in {set,long}jmp, pt.8

[yln]

Summary

Refine longjmp key management. For Linux, re-implement key retrieval in
C (instead of assembly). Removal of InitializeGuardPtr and a final
round of cleanups will be done in the next commit.

[phosek]

Our bots started failing with the following build error after this change landed:

/b/s/w/ir/k/llvm-project/compiler-rt/lib/tsan/rtl/tsan_platform_linux.cc:425:3: error: unknown type name 'jump_buf'; did you mean 'jmp_buf'?
  jump_buf env;
  ^~~~~~~~
  jmp_buf
/b/s/w/ir/k/cipd/linux-arm64/usr/include/setjmp.h:48:30: note: 'jmp_buf' declared here
typedef struct __jmp_buf_tag jmp_buf[1];
                             ^
/b/s/w/ir/k/llvm-project/compiler-rt/lib/tsan/rtl/tsan_platform_linux.cc:426:3: error: use of undeclared identifier 'REAL'
  REAL(setjmp)(env);
  ^
2 errors generated.

Is it possible to fix the issue or revert the change?

compiler-rt/trunk/lib/tsan/rtl/tsan_platform_linux.cc

 #endif
 
 #if SANITIZER_FREEBSD
 extern "C" void *__libc_stack_end;
 void *__libc_stack_end = 0;
 #endif
 
 #if SANITIZER_LINUX && defined(__aarch64__)
-void InitializeGuardPtr() __attribute__((visibility("hidden")));
+__tsan::uptr InitializeGuardPtr() __attribute__((visibility("hidden")));
 extern "C" __tsan::uptr _tsan_pointer_chk_guard;
 #endif
 
 namespace __tsan {
 
+#if SANITIZER_LINUX && defined(__aarch64__)
+static void InitializeLongjmpXorKey();
+static uptr longjmp_xor_key;
+#endif
+
 #ifdef TSAN_RUNTIME_VMA
 // Runtime detected VMA size.
 uptr vmaSize;
 #endif
 
 enum {
   MemTotal  = 0,
   MemShadow = 1,
#if SANITIZER_LINUX && defined(__aarch64__)
     if (old_personality != -1 && (old_personality & ADDR_NO_RANDOMIZE) == 0) {
       VReport(1, "WARNING: Program is run with randomized virtual address "
               "space, which wouldn't work with ThreadSanitizer.\n"
               "Re-execing with fixed virtual address space.\n");
       CHECK_NE(personality(old_personality | ADDR_NO_RANDOMIZE), -1);
       reexec = true;
     }
     // Initialize the guard pointer used in {sig}{set,long}jump.
-    InitializeGuardPtr();
+    longjmp_xor_key = InitializeGuardPtr();
+    uptr old_value = longjmp_xor_key;
+    InitializeLongjmpXorKey();
+    CHECK_EQ(longjmp_xor_key, old_value);
+    // If the above check fails for you, please contact me ([e-mail])
+    // and let me know the values of the two differing keys.  Please also set a
+    // breakpoint on `InitializeGuardPtr` and `InitializeLongjmpXorKey` and tell
+    // me the stack pointer (SP) values that go into the XOR operation (where we
+    // derive the key):
+    //
+    //   InitializeLongjmpXorKey:
+    //     uptr sp = (uptr)__builtin_frame_address(0);
+    //
+    //   InitializeGuardPtr (in tsan_rtl_aarch64.S):
+    //       mov  x0, sp
+    //       ...
+    //       eor  x0, x0, x1
+    //
+    // Then feel free to comment out the call to `InitializeLongjmpXorKey`.
 #endif
     if (reexec)
       ReExec();
   }
 
 #if !SANITIZER_GO
   CheckAndProtect();
   InitTlsSize();
# else  // Linux
   asm("ror  $0x11,     %0 \n"
       "xor  %%fs:0x30, %0 \n"
       : "=r" (sp)
       : "0" (mangled_sp));
   return sp;
 # endif
 #elif defined(__aarch64__)
 # if SANITIZER_LINUX
-  // TODO(yln): fix this
-  // return mangled_sp ^ _tsan_pointer_chk_guard;
-  return mangled_sp;
+  return mangled_sp ^ longjmp_xor_key;
 # else
   return mangled_sp;
 # endif
 #elif defined(__powerpc64__)
   // Reverse of:
   //  ld   r4, -28696(r13)
   //  xor  r4, r3, r4
   uptr xor_guard;

 # endif
 #endif
 
 uptr ExtractLongJmpSp(uptr *env) {
   uptr mangled_sp = env[LONG_JMP_SP_ENV_SLOT];
   return UnmangleLongJmpSp(mangled_sp);
 }
 
+#if SANITIZER_LINUX && defined(__aarch64__)
+// GLIBC mangles the function pointers in jmp_buf (used in {set,long}*jmp
+// functions) by XORing them with a random key.  For AArch64 it is a global
+// variable rather than a TCB one (as for x86_64/powerpc).  We obtain the key by
+// issuing a setjmp and XORing the SP pointer values to derive the key.
+static void InitializeLongjmpXorKey() {
+  // 1. Call REAL(setjmp), which stores the mangled SP in env.
+  jump_buf env;
+  REAL(setjmp)(env);
+
+  // 2. Retrieve mangled/vanilla SP.
+  uptr mangled_sp = ((uptr *)&env)[LONG_JMP_SP_ENV_SLOT];
+  uptr sp = (uptr)__builtin_frame_address(0);
+
+  // 3. xor SPs to obtain key.
+  longjmp_xor_key = mangled_sp ^ sp;
+}
+#endif
+
 void ImitateTlsWrite(ThreadState *thr, uptr tls_addr, uptr tls_size) {
   // Check that the thr object is in tls;
   const uptr thr_beg = (uptr)thr;
   const uptr thr_end = (uptr)thr + sizeof(*thr);
   CHECK_GE(thr_beg, tls_addr);
   CHECK_LE(thr_beg, tls_addr + tls_size);
   CHECK_GE(thr_end, tls_addr);
   CHECK_LE(thr_end, tls_addr + tls_size);

compiler-rt/trunk/lib/tsan/rtl/tsan_platform_mac.cc

#if defined(__aarch64__)
   if (max_vm != Mapping::kHiAppMemEnd) {
     Printf("ThreadSanitizer: unsupported vm address limit %p, expected %p.\n",
            max_vm, Mapping::kHiAppMemEnd);
     Die();
   }
 #endif
 }
 
-static const uptr kPthreadSetjmpXorKeySlot = 0x7;
-extern "C" uptr __tsan_darwin_setjmp_xor_key = 0;
+static uptr longjmp_xor_key = 0;
 
 void InitializePlatform() {
   DisableCoreDumperIfNecessary();
 #if !SANITIZER_GO
   CheckAndProtect();
 
   CHECK_EQ(main_thread_identity, 0);
   main_thread_identity = (uptr)pthread_self();
 
   prev_pthread_introspection_hook =
       pthread_introspection_hook_install(&my_pthread_introspection_hook);
 #endif
 
   if (GetMacosVersion() >= MACOS_VERSION_MOJAVE) {
-    __tsan_darwin_setjmp_xor_key =
-        (uptr)pthread_getspecific(kPthreadSetjmpXorKeySlot);
+    // Libsystem currently uses a process-global key; this might change.
+    const unsigned kTLSLongjmpXorKeySlot = 0x7;
+    longjmp_xor_key = (uptr)pthread_getspecific(kTLSLongjmpXorKeySlot);
   }
 }
 
 #ifdef __aarch64__
 # define LONG_JMP_SP_ENV_SLOT \
     ((GetMacosVersion() >= MACOS_VERSION_MOJAVE) ? 12 : 13)
 #else
 # define LONG_JMP_SP_ENV_SLOT 2
 #endif
 
 uptr ExtractLongJmpSp(uptr *env) {
   uptr mangled_sp = env[LONG_JMP_SP_ENV_SLOT];
-  return mangled_sp ^ __tsan_darwin_setjmp_xor_key;
+  uptr sp = mangled_sp ^ longjmp_xor_key;
+  return sp;
 }
 
 #if !SANITIZER_GO
 void ImitateTlsWrite(ThreadState *thr, uptr tls_addr, uptr tls_size) {
   // The pointer to the ThreadState object is stored in the shadow memory
   // of the tls.
   uptr tls_end = tls_addr + tls_size;
   uptr thread_identity = (uptr)pthread_self();