From b177ac4a44ec4024fec6927f4870ee96e9070696 Mon Sep 17 00:00:00 2001
From: davidtrevelyan <davidtrevelyan@users.noreply.github.com>
Date: Mon, 22 Jul 2024 21:35:37 +0100
Subject: [PATCH] [compiler-rt][rtsan] Use sanitizer internal allocator during
 rtsan init to avoid segfault in dlsym (#98679)

Follows https://github.com/llvm/llvm-project/pull/98268 with a fix for a
segfault during preinit on `ubuntu:20.04` environments. Previously,
`rtsan` was not handling the situation where `dlsym` calls `calloc`
during the interceptors initialization, resulting in a call to a
function at a null address.

@cjappl and I took inspiration from the solution in `nsan`, but we
re-used the sanitizer internal allocator instead of our own static
buffer. This PR also re-enables the existing non-instrumented `rtsan`
tests for `x86_64` and `arm64` architectures.

---------

Co-authored-by: Chris Apple <cja-private@pm.me>
---
 .../cmake/Modules/AllSupportedArchDefs.cmake  |  4 +--
 compiler-rt/lib/rtsan/rtsan.cpp               | 15 ++++++++++-
 compiler-rt/lib/rtsan/rtsan.h                 |  7 ++++++
 compiler-rt/lib/rtsan/rtsan_interceptors.cpp  | 25 +++++++++++++++++++
 4 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/compiler-rt/cmake/Modules/AllSupportedArchDefs.cmake b/compiler-rt/cmake/Modules/AllSupportedArchDefs.cmake
index 02ff92f6938101..29e5beb6182ba3 100644
--- a/compiler-rt/cmake/Modules/AllSupportedArchDefs.cmake
+++ b/compiler-rt/cmake/Modules/AllSupportedArchDefs.cmake
@@ -32,9 +32,7 @@ set(ALL_ASAN_SUPPORTED_ARCH ${X86} ${X86_64} ${ARM32} ${ARM64} ${RISCV64}
     ${LOONGARCH64})
 set(ALL_ASAN_ABI_SUPPORTED_ARCH ${X86_64} ${ARM64} ${ARM64_32})
 set(ALL_DFSAN_SUPPORTED_ARCH ${X86_64} ${MIPS64} ${ARM64} ${LOONGARCH64})
-#set(ALL_RTSAN_SUPPORTED_ARCH ${X86} ${X86_64} ${ARM32} ${ARM64} ${RISCV64}
-#    ${MIPS32} ${MIPS64} ${PPC64} ${S390X} ${SPARC} ${SPARCV9} ${HEXAGON}
-#    ${LOONGARCH64})
+set(ALL_RTSAN_SUPPORTED_ARCH ${X86_64} ${ARM64})
 
 if(ANDROID)
   set(OS_NAME "Android")
diff --git a/compiler-rt/lib/rtsan/rtsan.cpp b/compiler-rt/lib/rtsan/rtsan.cpp
index 43a63b4636f1e6..cf7fbddd9eb9c2 100644
--- a/compiler-rt/lib/rtsan/rtsan.cpp
+++ b/compiler-rt/lib/rtsan/rtsan.cpp
@@ -12,10 +12,23 @@
 #include <rtsan/rtsan_context.h>
 #include <rtsan/rtsan_interceptors.h>
 
+using namespace __rtsan;
+
+bool __rtsan::rtsan_initialized;
+bool __rtsan::rtsan_init_is_running;
+
 extern "C" {
 
 SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_init() {
-  __rtsan::InitializeInterceptors();
+  CHECK(!rtsan_init_is_running);
+  if (rtsan_initialized)
+    return;
+  rtsan_init_is_running = true;
+
+  InitializeInterceptors();
+
+  rtsan_init_is_running = false;
+  rtsan_initialized = true;
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_realtime_enter() {
diff --git a/compiler-rt/lib/rtsan/rtsan.h b/compiler-rt/lib/rtsan/rtsan.h
index 8b1219c13cbd34..ccddaf2c893efc 100644
--- a/compiler-rt/lib/rtsan/rtsan.h
+++ b/compiler-rt/lib/rtsan/rtsan.h
@@ -14,6 +14,13 @@
 
 extern "C" {
 
+namespace __rtsan {
+
+extern bool rtsan_initialized;
+extern bool rtsan_init_is_running;
+
+} // namespace __rtsan
+
 // Initialise rtsan interceptors.
 // A call to this method is added to the preinit array on Linux systems.
 SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_init();
diff --git a/compiler-rt/lib/rtsan/rtsan_interceptors.cpp b/compiler-rt/lib/rtsan/rtsan_interceptors.cpp
index 3a65f9d3f779d4..4d5423ec629d22 100644
--- a/compiler-rt/lib/rtsan/rtsan_interceptors.cpp
+++ b/compiler-rt/lib/rtsan/rtsan_interceptors.cpp
@@ -10,10 +10,14 @@
 
 #include "rtsan/rtsan_interceptors.h"
 
+#include "interception/interception.h"
+#include "sanitizer_common/sanitizer_allocator_dlsym.h"
+#include "sanitizer_common/sanitizer_allocator_internal.h"
 #include "sanitizer_common/sanitizer_platform.h"
 #include "sanitizer_common/sanitizer_platform_interceptors.h"
 
 #include "interception/interception.h"
+#include "rtsan/rtsan.h"
 #include "rtsan/rtsan_context.h"
 
 #if SANITIZER_APPLE
@@ -35,6 +39,15 @@
 
 using namespace __sanitizer;
 
+using __rtsan::rtsan_init_is_running;
+using __rtsan::rtsan_initialized;
+
+namespace {
+struct DlsymAlloc : public DlSymAllocator<DlsymAlloc> {
+  static bool UseImpl() { return !rtsan_initialized; }
+};
+} // namespace
+
 void ExpectNotRealtime(const char *intercepted_function_name) {
   __rtsan::GetContextForThisThread().ExpectNotRealtime(
       intercepted_function_name);
@@ -238,11 +251,17 @@ INTERCEPTOR(int, nanosleep, const struct timespec *rqtp,
 // Memory
 
 INTERCEPTOR(void *, calloc, SIZE_T num, SIZE_T size) {
+  if (DlsymAlloc::Use())
+    return DlsymAlloc::Callocate(num, size);
+
   ExpectNotRealtime("calloc");
   return REAL(calloc)(num, size);
 }
 
 INTERCEPTOR(void, free, void *ptr) {
+  if (DlsymAlloc::PointerIsMine(ptr))
+    return DlsymAlloc::Free(ptr);
+
   if (ptr != NULL) {
     ExpectNotRealtime("free");
   }
@@ -250,11 +269,17 @@ INTERCEPTOR(void, free, void *ptr) {
 }
 
 INTERCEPTOR(void *, malloc, SIZE_T size) {
+  if (DlsymAlloc::Use())
+    return DlsymAlloc::Allocate(size);
+
   ExpectNotRealtime("malloc");
   return REAL(malloc)(size);
 }
 
 INTERCEPTOR(void *, realloc, void *ptr, SIZE_T size) {
+  if (DlsymAlloc::Use() || DlsymAlloc::PointerIsMine(ptr))
+    return DlsymAlloc::Realloc(ptr, size);
+
   ExpectNotRealtime("realloc");
   return REAL(realloc)(ptr, size);
 }