[msan] Allow KMSAN to use -fsanitize-memory-param-retval

Let -fsanitize-memory-param-retval be used together with
-fsanitize=kernel-memory, so that it can be applied when building the
Linux kernel.

Also add clang/test/CodeGen/kmsan-param-retval.c to ensure that
-fsanitize-memory-param-retval eliminates shadow accesses for parameters
marked as undef.

Reviewed By: eugenis, vitalybuka

Differential Revision: https://reviews.llvm.org/D127860

3 files changed, 43 insertions, 0 deletions

diff --git a/clang/lib/Driver/SanitizerArgs.cpp b/clang/lib/Driver/SanitizerArgs.cpp
index edb1bfb..55d4469 100644
--- a/clang/lib/Driver/SanitizerArgs.cpp
+++ b/clang/lib/Driver/SanitizerArgs.cpp
@@ -648,6 +648,11 @@ SanitizerArgs::SanitizerArgs(const ToolChain &TC,
         options::OPT_fno_sanitize_memory_param_retval, MsanParamRetval);
     NeedPIE |= !(TC.getTriple().isOSLinux() &&
                  TC.getTriple().getArch() == llvm::Triple::x86_64);
+  } else if (AllAddedKinds & SanitizerKind::KernelMemory) {
+    MsanUseAfterDtor = false;
+    MsanParamRetval = Args.hasFlag(
+        options::OPT_fsanitize_memory_param_retval,
+        options::OPT_fno_sanitize_memory_param_retval, MsanParamRetval);
   } else {
     MsanUseAfterDtor = false;
     MsanParamRetval = false;
diff --git a/clang/test/CodeGen/kmsan-param-retval.c b/clang/test/CodeGen/kmsan-param-retval.c
new file mode 100644
index 0000000..3d952c0
--- /dev/null
+++ b/clang/test/CodeGen/kmsan-param-retval.c
@@ -0,0 +1,36 @@
+// RUN: %clang_cc1 -triple x86_64-linux-gnu -S -emit-llvm -O2 -fsanitize=kernel-memory -no-enable-noundef-analysis -o - %s | \
+// RUN:     FileCheck %s --check-prefix=CLEAN
+// RUN: %clang_cc1 -triple x86_64-linux-gnu -S -emit-llvm -O2 -fsanitize=kernel-memory -o - %s | \
+// RUN:     FileCheck %s --check-prefixes=NOUNDEF,NOUNDEF_ONLY
+// RUN: %clang_cc1 -triple x86_64-linux-gnu -S -emit-llvm -O2 -fsanitize=kernel-memory -mllvm -msan-eager-checks -o - %s | \
+// RUN:     FileCheck %s --check-prefixes=NOUNDEF,EAGER
+// RUN: %clang_cc1 -triple x86_64-linux-gnu -S -emit-llvm -O2 -fsanitize=kernel-memory -no-enable-noundef-analysis -fsanitize-memory-param-retval -o - %s | \
+// RUN:     FileCheck %s --check-prefixes=CLEAN
+// RUN: %clang_cc1 -triple x86_64-linux-gnu -S -emit-llvm -O2 -fsanitize=kernel-memory -fsanitize-memory-param-retval -o - %s | \
+// RUN:     FileCheck %s --check-prefixes=NOUNDEF,EAGER
+
+void foo();
+
+void bar(int x) {
+  if (x)
+    foo();
+}
+
+
+// CLEAN:   define dso_local void @bar(i32 %x)
+// NOUNDEF: define dso_local void @bar(i32 noundef %x)
+//
+// %param_shadow assignment gets optimized away with -O2, because it is at the beginning of the
+// struct returned by __msan_get_context_state(). Use %param_origin as a sign that the shadow of
+// the first argument is being used.
+//
+// Without noundef analysis, KMSAN emits metadata checks for the function parameter.
+// CLEAN:        load i32, ptr %param_origin
+//
+// With noundef analysis enabled, but without eager checks, KMSAN still emits metadata checks,
+// although the parameter is known to be defined.
+// NOUNDEF_ONLY: load i32, ptr %param_origin
+//
+// With noundef analysis and eager checks enabled, KMSAN won't emit metadata checks for function
+// parameters.
+// EAGER-NOT:    load i32, ptr %param_origin
diff --git a/clang/test/Driver/fsanitize-memory-param-retval.c b/clang/test/Driver/fsanitize-memory-param-retval.c
index 98ca16e..d82d208 100644
--- a/clang/test/Driver/fsanitize-memory-param-retval.c
+++ b/clang/test/Driver/fsanitize-memory-param-retval.c
@@ -3,6 +3,8 @@
 // RUN: %clang -target aarch64-linux-gnu %s -fsanitize=memory -fsanitize-memory-param-retval -c -### 2>&1 | FileCheck %s
 // RUN: %clang -target riscv32-linux-gnu %s -fsanitize=memory -fsanitize-memory-param-retval -c -### 2>&1 | FileCheck %s
 // RUN: %clang -target riscv64-linux-gnu %s -fsanitize=memory -fsanitize-memory-param-retval -c -### 2>&1 | FileCheck %s
+// RUN: %clang -target x86_64-linux-gnu %s -fsanitize=kernel-memory -fsanitize-memory-param-retval -c -### 2>&1 | FileCheck %s
+
 // CHECK: "-fsanitize-memory-param-retval"
 
 // RUN: %clang -target aarch64-linux-gnu -fsyntax-only %s -fsanitize=memory -fsanitize-memory-param-retval -c -### 2>&1 | FileCheck --check-prefix=11 %s