[llvm] [ASAN] Make asan pass idempotent. (PR #99439)

[via llvm-commits]

https://github.com/skc7 created https://github.com/llvm/llvm-project/pull/99439

When address sanitizer pass is run after it has already been run before, double instrumentation is seen in the resulting IR. This happens because there is no check in the pass, to verify if IR has been asan instrumented before.

This PR checks if asan_module_ctor is already present in the module and if present, return early without running the pass again.

>From 04d5f9bf9af4d751d0044a8dfc1f244094893c79 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Thu, 18 Jul 2024 11:49:24 +0530
Subject: [PATCH] [ASAN] Make asan pass idempotent.

---
 .../Instrumentation/AddressSanitizer.cpp      | 14 +++++
 .../AddressSanitizer/asan-pass-second-run.ll  | 56 +++++++++++++++++++
 2 files changed, 70 insertions(+)
 create mode 100644 llvm/test/Instrumentation/AddressSanitizer/asan-pass-second-run.ll

diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
index adf77f20cb1c7..30752e23073f7 100644
--- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
@@ -1249,8 +1249,22 @@ AddressSanitizerPass::AddressSanitizerPass(
       UseOdrIndicator(UseOdrIndicator), DestructorKind(DestructorKind),
       ConstructorKind(ConstructorKind) {}
 
+static bool hasAsanModuleCtor(Module &M) {
+  for (Function &F : M) {
+    if (F.getName() == kAsanModuleCtorName) {
+      return true;
+    }
+  }
+  return false;
+}
+
 PreservedAnalyses AddressSanitizerPass::run(Module &M,
                                             ModuleAnalysisManager &MAM) {
+  // Return early if asan.module_ctor is already present in the module.
+  // This implies that asan pass has already run before.
+  if (hasAsanModuleCtor(M))
+    return PreservedAnalyses::all();
+
   ModuleAddressSanitizer ModuleSanitizer(
       M, Options.InsertVersionCheck, Options.CompileKernel, Options.Recover,
       UseGlobalGC, UseOdrIndicator, DestructorKind, ConstructorKind);
diff --git a/llvm/test/Instrumentation/AddressSanitizer/asan-pass-second-run.ll b/llvm/test/Instrumentation/AddressSanitizer/asan-pass-second-run.ll
new file mode 100644
index 0000000000000..65d72a72f605b
--- /dev/null
+++ b/llvm/test/Instrumentation/AddressSanitizer/asan-pass-second-run.ll
@@ -0,0 +1,56 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 5
+; This test checks in the second run, function is not instrumented again.
+; RUN: opt < %s -passes=asan,asan -S | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; Function with sanitize_address is instrumented.
+; Function Attrs: nounwind uwtable
+;.
+; CHECK: @llvm.used = appending global [1 x ptr] [ptr @asan.module_ctor], section "llvm.metadata"
+; CHECK: @___asan_globals_registered = common hidden global i64 0
+; CHECK: @__start_asan_globals = extern_weak hidden global i64
+; CHECK: @__stop_asan_globals = extern_weak hidden global i64
+; CHECK: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 1, ptr @asan.module_ctor, ptr @asan.module_ctor }]
+;.
+define void @instr_sa(ptr %a) sanitize_address {
+; CHECK-LABEL: define void @instr_sa(
+; CHECK-SAME: ptr [[A:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = ptrtoint ptr [[A]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i64 [[TMP0]], 3
+; CHECK-NEXT:    [[TMP2:%.*]] = add i64 [[TMP1]], 2147450880
+; CHECK-NEXT:    [[TMP3:%.*]] = inttoptr i64 [[TMP2]] to ptr
+; CHECK-NEXT:    [[TMP4:%.*]] = load i8, ptr [[TMP3]], align 1
+; CHECK-NEXT:    [[TMP5:%.*]] = icmp ne i8 [[TMP4]], 0
+; CHECK-NEXT:    br i1 [[TMP5]], label %[[BB6:.*]], label %[[BB12:.*]], !prof [[PROF0:![0-9]+]]
+; CHECK:       [[BB6]]:
+; CHECK-NEXT:    [[TMP7:%.*]] = and i64 [[TMP0]], 7
+; CHECK-NEXT:    [[TMP8:%.*]] = add i64 [[TMP7]], 3
+; CHECK-NEXT:    [[TMP9:%.*]] = trunc i64 [[TMP8]] to i8
+; CHECK-NEXT:    [[TMP10:%.*]] = icmp sge i8 [[TMP9]], [[TMP4]]
+; CHECK-NEXT:    br i1 [[TMP10]], label %[[BB11:.*]], label %[[BB12]]
+; CHECK:       [[BB11]]:
+; CHECK-NEXT:    call void @__asan_report_load4(i64 [[TMP0]]) #[[ATTR3:[0-9]+]]
+; CHECK-NEXT:    unreachable
+; CHECK:       [[BB12]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = load i32, ptr [[A]], align 4
+; CHECK-NEXT:    [[TMP2:%.*]] = add i32 [[TMP1]], 1
+; CHECK-NEXT:    store i32 [[TMP2]], ptr [[A]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %tmp1 = load i32, ptr %a, align 4
+  %tmp2 = add i32 %tmp1,  1
+  store i32 %tmp2, ptr %a, align 4
+  ret void
+}
+;.
+; CHECK: attributes #[[ATTR0]] = { sanitize_address }
+; CHECK: attributes #[[ATTR1:[0-9]+]] = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
+; CHECK: attributes #[[ATTR2:[0-9]+]] = { nounwind }
+; CHECK: attributes #[[ATTR3]] = { nomerge }
+;.
+; CHECK: [[PROF0]] = !{!"branch_weights", i32 1, i32 1048575}
+;.