[llvm-branch-commits] [asan] Optimize initialization order checking (PR #101837)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-compiler-rt-sanitizer

Author: Vitaly Buka (vitalybuka)

<details>
<summary>Changes</summary>

Before the patch each TU was executing pair
`__asan_before_dynamic_init` and
`__asan_after_dynamic_init`.

`__asan_before_dynamic_init` supports incremental
poisoning, but `__asan_after_dynamic_init`
unpoisons all globals, so
`__asan_before_dynamic_init` must repoison
everything again.

So it's O(N^2) where N is the number of globals
(or TUs). It can be expensive for large binaries.

The patch will make it O(N).

The patch requires LLD and platform with
SANITIZER_CAN_USE_PREINIT_ARRAY. For the rest it
continue to be O(N^2).


---
Full diff: https://github.com/llvm/llvm-project/pull/101837.diff


3 Files Affected:

- (modified) compiler-rt/lib/asan/asan_globals.cpp (+40) 
- (added) compiler-rt/test/asan/TestCases/initialization-nobug-lld.cpp (+14) 
- (modified) compiler-rt/test/asan/TestCases/initialization-nobug.cpp (+8-5) 


``````````diff
diff --git a/compiler-rt/lib/asan/asan_globals.cpp b/compiler-rt/lib/asan/asan_globals.cpp
index 293e771e1dc06..373b7eaa3c537 100644
--- a/compiler-rt/lib/asan/asan_globals.cpp
+++ b/compiler-rt/lib/asan/asan_globals.cpp
@@ -520,6 +520,44 @@ void __asan_before_dynamic_init(const char *module_name) {
   current_dynamic_init_module_name = module_name;
 }
 
+// Maybe SANITIZER_CAN_USE_PREINIT_ARRAY is to conservative for `.init_array`,
+// however we should not make mistake here. If `AfterDynamicInit` was not
+// executed at all we will have false reports on globals.
+#if SANITIZER_CAN_USE_PREINIT_ARRAY
+// This is optimization. We will ignore all `__asan_after_dynamic_init`, but the
+// last `__asan_after_dynamic_init`. We don't need number of
+// `__asan_{before,after}_dynamic_init` matches, but we need that the last call
+// was to `__asan_after_dynamic_init`, as it will unpoison all global preparing
+// program for `main` execution. To run `__asan_after_dynamic_init` later, we
+// will register in `.init_array`.
+static bool allow_after_dynamic_init SANITIZER_GUARDED_BY(mu_for_globals) =
+    false;
+
+static void __attribute__((used)) AfterDynamicInit(void) {
+  {
+    Lock lock(&mu_for_globals);
+    if (allow_after_dynamic_init)
+      return;
+    allow_after_dynamic_init = true;
+  }
+  if (flags()->report_globals >= 3)
+    Printf("AfterDynamicInit\n");
+  __asan_after_dynamic_init();
+}
+
+// 65537 will make it run after constructors with default priority, but it
+// requires ld.lld. With ld.bfd it can be called to early, and fail the
+// optimization. However, correctness should not be affected, as after the first
+// call all subsequent `__asan_after_dynamic_init` will be allowed.
+__attribute__((section(".init_array.65537"), used)) static void (
+    *asan_after_init_array)(void) = AfterDynamicInit;
+#else
+// Allow all `__asan_after_dynamic_init` if `AfterDynamicInit` is not set.
+// Compiler still generates `__asan_{before,after}_dynamic_init`in pairs, and
+// it's guaranteed that `__asan_after_dynamic_init` will be the last.
+static constexpr bool allow_after_dynamic_init = true;
+#endif  // SANITIZER_CAN_USE_PREINIT_ARRAY
+
 // This method runs immediately after dynamic initialization in each TU, when
 // all dynamically initialized globals except for those defined in the current
 // TU are poisoned.  It simply unpoisons all dynamically initialized globals.
@@ -528,6 +566,8 @@ void __asan_after_dynamic_init() {
     return;
   CHECK(AsanInited());
   Lock lock(&mu_for_globals);
+  if (!allow_after_dynamic_init)
+    return;
   if (!current_dynamic_init_module_name)
     return;
 
diff --git a/compiler-rt/test/asan/TestCases/initialization-nobug-lld.cpp b/compiler-rt/test/asan/TestCases/initialization-nobug-lld.cpp
new file mode 100644
index 0000000000000..ece3118de4ad9
--- /dev/null
+++ b/compiler-rt/test/asan/TestCases/initialization-nobug-lld.cpp
@@ -0,0 +1,14 @@
+// RUN: %clangxx_asan -O3 %p/initialization-nobug.cpp %p/Helpers/initialization-nobug-extra.cpp -fuse-ld=lld -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInit"
+
+// Same as initialization-nobug.cpp, but with lld we expect just one
+// `DynInitUnpoison` executed after `AfterDynamicInit` at the end.
+// REQUIRES: lld-available
+
+// With dynamic runtimes `AfterDynamicInit` will called before `executable`
+// contructors, with constructors of dynamic runtime.
+// XFAIL: asan-dynamic-runtime
+
+// CHECK: DynInitPoison module: {{.*}}initialization-nobug.cpp
+// CHECK: DynInitPoison module: {{.*}}initialization-nobug-extra.cpp
+// CHECK: AfterDynamicInit
+// CHECK: DynInitUnpoison
diff --git a/compiler-rt/test/asan/TestCases/initialization-nobug.cpp b/compiler-rt/test/asan/TestCases/initialization-nobug.cpp
index 5659db088096b..f58f38203d0b2 100644
--- a/compiler-rt/test/asan/TestCases/initialization-nobug.cpp
+++ b/compiler-rt/test/asan/TestCases/initialization-nobug.cpp
@@ -1,10 +1,10 @@
 // A collection of various initializers which shouldn't trip up initialization
 // order checking.  If successful, this will just return 0.
 
-// RUN: %clangxx_asan -O0 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInit"
-// RUN: %clangxx_asan -O1 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInit"
-// RUN: %clangxx_asan -O2 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInit"
-// RUN: %clangxx_asan -O3 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInit"
+// RUN: %clangxx_asan -O0 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInitPoison"
+// RUN: %clangxx_asan -O1 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInitPoison"
+// RUN: %clangxx_asan -O2 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInitPoison"
+// RUN: %clangxx_asan -O3 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInitPoison"
 
 // Simple access:
 // Make sure that accessing a global in the same TU is safe
@@ -44,6 +44,9 @@ int getStructWithDtorValue() { return struct_with_dtor.value; }
 int main() { return 0; }
 
 // CHECK: DynInitPoison module: {{.*}}initialization-nobug.cpp
-// CHECK: DynInitUnpoison
 // CHECK: DynInitPoison module: {{.*}}initialization-nobug-extra.cpp
+
+// In general case we only need the last of DynInit{Poison,Unpoison} is always
+// Unpoison.
+
 // CHECK: DynInitUnpoison

``````````

</details>


https://github.com/llvm/llvm-project/pull/101837