[llvm] [LLVM] Make sanitizers respect the disable_santizer_instrumentation attribute. (PR #91732)

[Daniel Kiss via llvm-commits]

https://github.com/DanielKristofKiss updated https://github.com/llvm/llvm-project/pull/91732

>From f14f146cf18a22202659f6661d4b0bd1e138bb19 Mon Sep 17 00:00:00 2001
From: Daniel Kiss <daniel.kiss at arm.com>
Date: Fri, 10 May 2024 13:55:01 +0200
Subject: [PATCH] [LLVM] Make sanitizers respect the
 disable_santizer_instrumentation attribute.

Disable_sanitizer_instrumetation is attached to functions that shall not
be instrumented, e.g. ifunc resolver because those run before everything
is initialized.
Some sanitizer already handles this attribute, this patch adds it to
DataFLow and Coverage.
---
 .../Instrumentation/DataFlowSanitizer.cpp     |  3 +-
 .../Instrumentation/SanitizerCoverage.cpp     |  2 +
 ...aflow-disable-sanitizer-instrumentation.ll | 47 +++++++++++++++++++
 ...erage-disable-sanitizer-instrumentation.ll | 46 ++++++++++++++++++
 4 files changed, 97 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Instrumentation/DataFlowSanitizer/dataflow-disable-sanitizer-instrumentation.ll
 create mode 100644 llvm/test/Instrumentation/SanitizerCoverage/coverage-disable-sanitizer-instrumentation.ll

diff --git a/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
index 851edb4ce829e..20d11e0ab55f2 100644
--- a/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
@@ -1546,7 +1546,8 @@ bool DataFlowSanitizer::runImpl(
   SmallPtrSet<Constant *, 1> PersonalityFns;
   for (Function &F : M)
     if (!F.isIntrinsic() && !DFSanRuntimeFunctions.contains(&F) &&
-        !LibAtomicFunction(F)) {
+        !LibAtomicFunction(F) &&
+        !F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation)) {
       FnsToInstrument.push_back(&F);
       if (F.hasPersonalityFn())
         PersonalityFns.insert(F.getPersonalityFn()->stripPointerCasts());
diff --git a/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp b/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
index 56d4907ae47a9..6a89cee9aaf6c 100644
--- a/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
+++ b/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
@@ -631,6 +631,8 @@ void ModuleSanitizerCoverage::instrumentFunction(Function &F) {
     return;
   if (F.hasFnAttribute(Attribute::NoSanitizeCoverage))
     return;
+  if (F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation))
+    return;
   if (Options.CoverageType >= SanitizerCoverageOptions::SCK_Edge) {
     SplitAllCriticalEdges(
         F, CriticalEdgeSplittingOptions().setIgnoreUnreachableDests());
diff --git a/llvm/test/Instrumentation/DataFlowSanitizer/dataflow-disable-sanitizer-instrumentation.ll b/llvm/test/Instrumentation/DataFlowSanitizer/dataflow-disable-sanitizer-instrumentation.ll
new file mode 100644
index 0000000000000..3fb922736d020
--- /dev/null
+++ b/llvm/test/Instrumentation/DataFlowSanitizer/dataflow-disable-sanitizer-instrumentation.ll
@@ -0,0 +1,47 @@
+
+; This test checks that we are not instrumenting sanitizer code.
+; RUN: opt < %s -passes='module(msan)' -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_memory is instrumented.
+; Function Attrs: nounwind uwtable
+define void @instr_sa(ptr %a) sanitize_memory {
+entry:
+  %tmp1 = load i32, ptr %a, align 4
+  %tmp2 = add i32 %tmp1,  1
+  store i32 %tmp2, ptr %a, align 4
+  ret void
+}
+
+; CHECK-LABEL: @instr_sa
+; CHECK: %0 = load i64, ptr @__msan_param_tls
+
+
+; Function with disable_sanitizer_instrumentation is not instrumented.
+; Function Attrs: nounwind uwtable
+define void @noinstr_dsi(ptr %a) disable_sanitizer_instrumentation {
+entry:
+  %tmp1 = load i32, ptr %a, align 4
+  %tmp2 = add i32 %tmp1,  1
+  store i32 %tmp2, ptr %a, align 4
+  ret void
+}
+
+; CHECK-LABEL: @noinstr_dsi
+; CHECK-NOT: %0 = load i64, ptr @__msan_param_tls
+
+
+; disable_sanitizer_instrumentation takes precedence over sanitize_memory.
+; Function Attrs: nounwind uwtable
+define void @noinstr_dsi_sa(ptr %a) disable_sanitizer_instrumentation sanitize_memory {
+entry:
+  %tmp1 = load i32, ptr %a, align 4
+  %tmp2 = add i32 %tmp1,  1
+  store i32 %tmp2, ptr %a, align 4
+  ret void
+}
+
+; CHECK-LABEL: @noinstr_dsi_sa
+; CHECK-NOT: %0 = load i64, ptr @__msan_param_tls
diff --git a/llvm/test/Instrumentation/SanitizerCoverage/coverage-disable-sanitizer-instrumentation.ll b/llvm/test/Instrumentation/SanitizerCoverage/coverage-disable-sanitizer-instrumentation.ll
new file mode 100644
index 0000000000000..dc3d486220153
--- /dev/null
+++ b/llvm/test/Instrumentation/SanitizerCoverage/coverage-disable-sanitizer-instrumentation.ll
@@ -0,0 +1,46 @@
+; This test checks that we are not instrumenting sanitizer code.
+; RUN: opt < %s -passes='module(sancov-module)' -sanitizer-coverage-level=3 -sanitizer-coverage-control-flow -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
+define void @instr_sa(ptr %a) sanitize_address {
+entry:
+  %tmp1 = load i32, ptr %a, align 4
+  %tmp2 = add i32 %tmp1,  1
+  store i32 %tmp2, ptr %a, align 4
+  ret void
+}
+
+; CHECK-LABEL: @instr_sa
+; CHECK: call void @__sanitizer_cov_trace_pc_guard(
+
+
+; Function with disable_sanitizer_instrumentation is not instrumented.
+; Function Attrs: nounwind uwtable
+define void @noinstr_dsi(ptr %a) disable_sanitizer_instrumentation {
+entry:
+  %tmp1 = load i32, ptr %a, align 4
+  %tmp2 = add i32 %tmp1,  1
+  store i32 %tmp2, ptr %a, align 4
+  ret void
+}
+
+; CHECK-LABEL: @noinstr_dsi
+; CHECK-NOT: call void @__sanitizer_cov_trace_pc_guard(
+
+
+; disable_sanitizer_instrumentation takes precedence over sanitize_address.
+; Function Attrs: nounwind uwtable
+define void @noinstr_dsi_sa(ptr %a) disable_sanitizer_instrumentation sanitize_address {
+entry:
+  %tmp1 = load i32, ptr %a, align 4
+  %tmp2 = add i32 %tmp1,  1
+  store i32 %tmp2, ptr %a, align 4
+  ret void
+}
+
+; CHECK-LABEL: @noinstr_dsi_sa
+; CHECK-NOT: call void @__sanitizer_cov_trace_pc_guard(