10 files changed, 117 insertions, 1 deletions
diff --git a/compiler-rt/lib/asan/asan_fuchsia.cc b/compiler-rt/lib/asan/asan_fuchsia.cc
index 0b5bff4f565..34399c92310 100644
--- a/compiler-rt/lib/asan/asan_fuchsia.cc
+++ b/compiler-rt/lib/asan/asan_fuchsia.cc
@@ -190,6 +190,13 @@ static void ThreadExitHook(void *hook, uptr os_id) {
   AsanThread::TSDDtor(per_thread);
 }
 
+bool HandleDlopenInit() {
+  // Not supported on this platform.
+  static_assert(!SANITIZER_SUPPORTS_INIT_FOR_DLOPEN,
+                "Expected SANITIZER_SUPPORTS_INIT_FOR_DLOPEN to be false");
+  return false;
+}
+
 }  // namespace __asan
 
 // These are declared (in extern "C") by <zircon/sanitizer.h>.
diff --git a/compiler-rt/lib/asan/asan_internal.h b/compiler-rt/lib/asan/asan_internal.h
index 654878cd15f..57869497c7d 100644
--- a/compiler-rt/lib/asan/asan_internal.h
+++ b/compiler-rt/lib/asan/asan_internal.h
@@ -111,6 +111,11 @@ void *AsanDlSymNext(const char *sym);
 
 void ReserveShadowMemoryRange(uptr beg, uptr end, const char *name);
 
+// Returns `true` iff most of ASan init process should be skipped due to the
+// ASan library being loaded via `dlopen()`. Platforms may perform any
+// `dlopen()` specific initialization inside this function.
+bool HandleDlopenInit();
+
 // Add convenient macro for interface functions that may be represented as
 // weak hooks.
 #define ASAN_MALLOC_HOOK(ptr, size)                                   \
diff --git a/compiler-rt/lib/asan/asan_linux.cc b/compiler-rt/lib/asan/asan_linux.cc
index 625f32d408d..a150b1955d6 100644
--- a/compiler-rt/lib/asan/asan_linux.cc
+++ b/compiler-rt/lib/asan/asan_linux.cc
@@ -248,6 +248,13 @@ void *AsanDlSymNext(const char *sym) {
   return dlsym(RTLD_NEXT, sym);
 }
 
+bool HandleDlopenInit() {
+  // Not supported on this platform.
+  static_assert(!SANITIZER_SUPPORTS_INIT_FOR_DLOPEN,
+                "Expected SANITIZER_SUPPORTS_INIT_FOR_DLOPEN to be false");
+  return false;
+}
+
 } // namespace __asan
 
 #endif  // SANITIZER_FREEBSD || SANITIZER_LINUX || SANITIZER_NETBSD ||
diff --git a/compiler-rt/lib/asan/asan_malloc_mac.cc b/compiler-rt/lib/asan/asan_malloc_mac.cc
index 733ba2d86e1..27281f1bc83 100644
--- a/compiler-rt/lib/asan/asan_malloc_mac.cc
+++ b/compiler-rt/lib/asan/asan_malloc_mac.cc
@@ -61,4 +61,25 @@ using namespace __asan;
 
 #include "sanitizer_common/sanitizer_malloc_mac.inc"
 
+namespace COMMON_MALLOC_NAMESPACE {
+bool HandleDlopenInit() {
+  static_assert(SANITIZER_SUPPORTS_INIT_FOR_DLOPEN,
+                "Expected SANITIZER_SUPPORTS_INIT_FOR_DLOPEN to be true");
+  // We have no reliable way of knowing how we are being loaded
+  // so make it a requirement on Apple platforms to set this environment
+  // variable to indicate that we want to perform initialization via
+  // dlopen().
+  auto init_str = GetEnv("APPLE_ASAN_INIT_FOR_DLOPEN");
+  if (!init_str)
+    return false;
+  if (internal_strncmp(init_str, "1", 1) != 0)
+    return false;
+  // When we are loaded via `dlopen()` path we still initialize the malloc zone
+  // so Symbolication clients (e.g. `leaks`) that load the ASan allocator can
+  // find an initialized malloc zone.
+  InitMallocZoneFields();
+  return true;
+}
+}  // namespace COMMON_MALLOC_NAMESPACE
+
 #endif
diff --git a/compiler-rt/lib/asan/asan_rtems.cc b/compiler-rt/lib/asan/asan_rtems.cc
index a4af940057e..b48cc6a75d7 100644
--- a/compiler-rt/lib/asan/asan_rtems.cc
+++ b/compiler-rt/lib/asan/asan_rtems.cc
@@ -213,6 +213,12 @@ static void HandleExit() {
   }
 }
 
+bool HandleDlopenInit() {
+  // Not supported on this platform.
+  static_assert(!SANITIZER_SUPPORTS_INIT_FOR_DLOPEN,
+                "Expected SANITIZER_SUPPORTS_INIT_FOR_DLOPEN to be false");
+  return false;
+}
 }  // namespace __asan
 
 // These are declared (in extern "C") by <some_path/sanitizer.h>.
diff --git a/compiler-rt/lib/asan/asan_rtl.cc b/compiler-rt/lib/asan/asan_rtl.cc
index 0c6390352a5..0ecbcd59487 100644
--- a/compiler-rt/lib/asan/asan_rtl.cc
+++ b/compiler-rt/lib/asan/asan_rtl.cc
@@ -396,6 +396,14 @@ static void AsanInitInternal() {
   // initialization steps look at flags().
   InitializeFlags();
 
+  // Stop performing init at this point if we are being loaded via
+  // dlopen() and the platform supports it.
+  if (SANITIZER_SUPPORTS_INIT_FOR_DLOPEN && UNLIKELY(HandleDlopenInit())) {
+    asan_init_is_running = false;
+    VReport(1, "AddressSanitizer init is being performed for dlopen().\n");
+    return;
+  }
+
   AsanCheckIncompatibleRT();
   AsanCheckDynamicRTPrereqs();
   AvoidCVE_2016_2143();
diff --git a/compiler-rt/lib/asan/asan_win.cc b/compiler-rt/lib/asan/asan_win.cc
index 5661d911cab..068f4a5d247 100644
--- a/compiler-rt/lib/asan/asan_win.cc
+++ b/compiler-rt/lib/asan/asan_win.cc
@@ -322,6 +322,13 @@ int __asan_set_seh_filter() {
   return 0;
 }
 
+bool HandleDlopenInit() {
+  // Not supported on this platform.
+  static_assert(!SANITIZER_SUPPORTS_INIT_FOR_DLOPEN,
+                "Expected SANITIZER_SUPPORTS_INIT_FOR_DLOPEN to be false");
+  return false;
+}
+
 #if !ASAN_DYNAMIC
 // The CRT runs initializers in this order:
 // - C initializers, from XIA to XIZ
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_platform.h b/compiler-rt/lib/sanitizer_common/sanitizer_platform.h
index 106a147e573..fd7228c3b33 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_platform.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_platform.h
@@ -342,4 +342,13 @@
 #define SANITIZER_SYMBOLIZER_MARKUP 0
 #endif
 
+// Enable ability to support sanitizer initialization that is
+// compatible with the sanitizer library being loaded via
+// `dlopen()`.
+#if SANITIZER_MAC
+#define SANITIZER_SUPPORTS_INIT_FOR_DLOPEN 1
+#else
+#define SANITIZER_SUPPORTS_INIT_FOR_DLOPEN 0
+#endif
+
 #endif // SANITIZER_PLATFORM_H
diff --git a/compiler-rt/test/asan/TestCases/Darwin/init_for_dlopen.cc b/compiler-rt/test/asan/TestCases/Darwin/init_for_dlopen.cc
new file mode 100644
index 00000000000..8a0fbf943b6
--- /dev/null
+++ b/compiler-rt/test/asan/TestCases/Darwin/init_for_dlopen.cc
@@ -0,0 +1,46 @@
+// RUN: %clangxx -g -O0 %s -o %t
+
+// Check that trying to dlopen() the ASan dylib fails.
+// We explictly set `abort_on_error=0` because
+// - By default the lit config sets this but we don't want this
+//   test to implicitly depend on this.
+// - It avoids requiring `--crash` to be passed to `not`.
+// RUN: APPLE_ASAN_INIT_FOR_DLOPEN=0 %env_asan_opts=abort_on_error=0 not \
+// RUN:   %run %t %shared_libasan 2>&1 | \
+// RUN:   FileCheck -check-prefix=CHECK-DL-OPEN-FAIL %s
+// RUN: env -u APPLE_ASAN_INIT_FOR_DLOPEN %env_asan_opts=abort_on_error=0 not \
+// RUN:   %run %t %shared_libasan 2>&1 | \
+// RUN:   FileCheck -check-prefix=CHECK-DL-OPEN-FAIL %s
+
+// Check that we can successfully dlopen the ASan dylib when we set the right
+// environment variable.
+// RUN: env APPLE_ASAN_INIT_FOR_DLOPEN=1 %run %t %shared_libasan 2>&1 | \
+// RUN:   FileCheck -check-prefix=CHECK-DL-OPEN-SUCCESS %s
+
+#include <dlfcn.h>
+#include <stdio.h>
+
+// CHECK-DL-OPEN-FAIL: ERROR: Interceptors are not working
+
+int main(int argc, char **argv) {
+  if (argc != 2) {
+    fprintf(stderr, "Usage: %s <dylib_path>\n", argv[0]);
+    return 1;
+  }
+  const char *dylib_path = argv[1];
+  void *handle = dlopen(dylib_path, RTLD_LAZY);
+  if (!handle) {
+    fprintf(stderr, "Failed to dlopen: %s\n", dlerror());
+    return 1;
+  }
+  // Make sure we can find a function we expect to be in the dylib.
+  void *fn = dlsym(handle, "__sanitizer_mz_size");
+  if (!fn) {
+    fprintf(stderr, "Failed to get symbol: %s\n", dlerror());
+    return 1;
+  }
+  // TODO(dliew): Actually call a function from the dylib that is safe to call.
+  // CHECK-DL-OPEN-SUCCESS: DONE
+  printf("DONE\n");
+  return 0;
+}
diff --git a/compiler-rt/test/sanitizer_common/ios_commands/iossim_run.py b/compiler-rt/test/sanitizer_common/ios_commands/iossim_run.py
index 0d9ca935ccc..e1f633e8746 100755
--- a/compiler-rt/test/sanitizer_common/ios_commands/iossim_run.py
+++ b/compiler-rt/test/sanitizer_common/ios_commands/iossim_run.py
@@ -8,7 +8,7 @@ if not "SANITIZER_IOSSIM_TEST_DEVICE_IDENTIFIER" in os.environ:
 
 device_id = os.environ["SANITIZER_IOSSIM_TEST_DEVICE_IDENTIFIER"]
 
-for e in ["ASAN_OPTIONS", "TSAN_OPTIONS", "UBSAN_OPTIONS"]:
+for e in ["ASAN_OPTIONS", "TSAN_OPTIONS", "UBSAN_OPTIONS", "APPLE_ASAN_INIT_FOR_DLOPEN"]:
   if e in os.environ:
     os.environ["SIMCTL_CHILD_" + e] = os.environ[e]