[msan] Implement -msan-pass-caller-to-runtime.

[glider]

Summary

Linux kernel has a concept of noinstr code, which is used to prevent
all kinds of instrumentation for annotated functions.
In particular, syscall and IRQ entry functions are implemented as
noinstr.

When these functions call KMSAN-instrumented functions, they fail to
properly set up the metadata for function arguments, potentially leading
to false positive reports.

In order to detect transitions from noinstr to instrumented code, we
introduce the -msan-pass-caller-to-runtime flag, which allows KMSAN to
call msan_get_context_state_caller() at the beginning of functions
that take one or more parameters. msan_get_context_state_caller()
accepts the caller address passed to it by the instrumentation code.
That address can be used by the runtime to figure out whether a call
happened from a noinstr function, and wipe the context state, preventing
the error reports.

For backward compatibility with BSD systems that use KMSAN, we keep
-msan-pass-caller-to-runtime=0 a default value.

[vitalybuka]

Do you have estimate of how often this happend? How many different instrumented functions which can be called from uninstrumented code?

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1279

Can you avoid ClPassCallerToRuntime and always pass the argument? I guess function will just not use it.

[glider]

In D126385#3542914, @vitalybuka wrote:

Do you have estimate of how often this happend? How many different instrumented functions which can be called from uninstrumented code?

There are ~200 noinstr annotations in the Linux kernel. Some are used in macros (e.g. there are several dozens of interrupt descriptor table entries that are implemented as noinstr functions).
It's hard to tell how many instrumented functions end up being called with consequences, but certainly annotating them manually is not an option.

We could probably optimize by not passing builtin_return_address() to msan_get_context_state() if the instrumented function doesn't take any arguments. Not sure it is worth the hassle (we'll need a separate version of __msan_get_context_state())

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1279

This would cost us a register spill, sounds like a regression for *BSD systems that don't need the argument.

[glider]

Introduced __msan_get_context_state_caller()

[glider]

Updated a comment

[glider]

FWIW in my current configuration out of 128K instrumented functions there are 120K calls to msan_get_context_state_caller() and 8K calls to msan_get_context_state().

[vitalybuka]

I assume it's abandoned

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

 // The instrumentation is only enabled in KMSAN builds, and only if
 // -msan-handle-asm-conservative is on. This is done because we may want to
 // quickly disable assembly instrumentation when it breaks.
 static cl::opt<bool> ClHandleAsmConservative(
     "msan-handle-asm-conservative",
     cl::desc("conservative handling of inline assembly"), cl::Hidden,
     cl::init(true));
 
+static cl::opt<bool> ClPassCallerToRuntime(
+    "msan-pass-caller-to-runtime",
+    cl::desc("(KMSAN only) pass caller address to the runtime to detect calls "
+             "from non-instrumented code"), cl::Hidden, cl::init(false));
+
 // This flag controls whether we check the shadow of the address
 // operand of load or store. Such bugs are very rare, since load from
 // a garbage address typically results in SEGV, but still happen
 // (e.g. only lower bits of address are garbage, or the access happens
 // early at program startup where malloc-ed memory is more likely to
 // be zeroed. As of 2012-08-28 this flag adds 20% slowdown.
 static cl::opt<bool> ClCheckAccessAddress("msan-check-access-address",
        cl::desc("report accesses through a pointer which has poisoned shadow"),
private:
   FunctionCallee MsanSetOriginFn;
 
   /// MSan runtime replacements for memmove, memcpy and memset.
   FunctionCallee MemmoveFn, MemcpyFn, MemsetFn;
 
   /// KMSAN callback for task-local function argument shadow.
   StructType *MsanContextStateTy;
   FunctionCallee MsanGetContextStateFn;
+  FunctionCallee MsanGetContextStateCallerFn;
 
   /// Functions for poisoning/unpoisoning local variables
   FunctionCallee MsanPoisonAllocaFn, MsanUnpoisonAllocaFn;
 
   /// Each of the MsanMetadataPtrXxx functions returns a pair of shadow/origin
   /// pointers.
   FunctionCallee MsanMetadataPtrForLoadN, MsanMetadataPtrForStoreN;
   FunctionCallee MsanMetadataPtrForLoad_1_8[4];
void MemorySanitizer::createKernelApi(Module &M) {
   ParamTLS = nullptr;
   ParamOriginTLS = nullptr;
   VAArgTLS = nullptr;
   VAArgOriginTLS = nullptr;
   VAArgOverflowSizeTLS = nullptr;
 
   WarningFn = M.getOrInsertFunction("__msan_warning", IRB.getVoidTy(),
                                     IRB.getInt32Ty());
-  // Requests the per-task context state (kmsan_context_state*) from the
-  // runtime library.
+  // Per-task context state (kmsan_context_state*) provided by the runtime
+  // library.
   MsanContextStateTy = StructType::get(
       ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8),
       ArrayType::get(IRB.getInt64Ty(), kRetvalTLSSize / 8),
       ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8),
       ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8), /* va_arg_origin */
       IRB.getInt64Ty(), ArrayType::get(OriginTy, kParamTLSSize / 4), OriginTy,
       OriginTy);
+
+  // __msan_get_context_state() requests the context state from the runtime.
+  // For BSD kernels it is inserted into every instrumented function.
+  // For Linux kernels - into every instrumented function without parameters.
   MsanGetContextStateFn = M.getOrInsertFunction(
       "__msan_get_context_state", PointerType::get(MsanContextStateTy, 0));
 
+  if (ClPassCallerToRuntime) {
+    // __msan_get_context_state_caller() takes the result of __builtin_return_address(0).
+    // Used when building Linux with KMSAN.
+    MsanGetContextStateCallerFn = M.getOrInsertFunction(
+        "__msan_get_context_state_caller", PointerType::get(MsanContextStateTy, 0),
+  PointerType::get(IRB.getInt8Ty(), 0));
+  } else {
+    MsanGetContextStateCallerFn = nullptr;
+  }
+
   Type *RetTy = StructType::get(PointerType::get(IRB.getInt8Ty(), 0),
                                 PointerType::get(IRB.getInt32Ty(), 0));
 
   for (int ind = 0, size = 1; ind < 4; ind++, size <<= 1) {
     std::string name_load =
         "__msan_metadata_ptr_for_load_" + std::to_string(size);
     std::string name_store =
         "__msan_metadata_ptr_for_store_" + std::to_string(size);
for (const auto &ShadowData : InstrumentationList) {
       Value *Origin = ShadowData.Origin;
       materializeOneCheck(OrigIns, Shadow, Origin, InstrumentWithCalls);
     }
     LLVM_DEBUG(dbgs() << "DONE:\n" << F);
   }
 
   // Returns the last instruction in the new prologue
   void insertKmsanPrologue(IRBuilder<> &IRB) {
-    Value *ContextState = IRB.CreateCall(MS.MsanGetContextStateFn, {});
+    Value *ContextState = nullptr;
+    if (MS.MsanGetContextStateCallerFn && F.arg_size() > 0) {
+        Value *RetAddr = IRB.CreateCall(
+            Intrinsic::getDeclaration(F.getParent(), Intrinsic::returnaddress),
+            IRB.getInt32(0));
+        ContextState = IRB.CreateCall(MS.MsanGetContextStateCallerFn, RetAddr);
+    } else {
+        ContextState = IRB.CreateCall(MS.MsanGetContextStateFn, {});
+    }
+
     Constant *Zero = IRB.getInt32(0);
     MS.ParamTLS = IRB.CreateGEP(MS.MsanContextStateTy, ContextState,
                                 {Zero, IRB.getInt32(0)}, "param_shadow");
     MS.RetvalTLS = IRB.CreateGEP(MS.MsanContextStateTy, ContextState,
                                  {Zero, IRB.getInt32(1)}, "retval_shadow");
     MS.VAArgTLS = IRB.CreateGEP(MS.MsanContextStateTy, ContextState,
                                 {Zero, IRB.getInt32(2)}, "va_arg_shadow");
     MS.VAArgOriginTLS = IRB.CreateGEP(MS.MsanContextStateTy, ContextState,

llvm/test/Instrumentation/MemorySanitizer/msan_kernel_pass-caller-to-runtime.ll

+; Test KMSAN behavior with and without -msan-pass-caller-to-runtime.
+; RUN: opt < %s -msan-kernel=1 -S -passes=msan 2>&1 | FileCheck %s -check-prefixes=CHECK-NOPASS
+; RUN: opt < %s -msan-kernel=1 -S -passes=msan -msan-pass-caller-to-runtime=1 2>&1 | FileCheck %s -check-prefixes=CHECK-PASS
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; Check the instrumentation prologue for a function without arguments.
+define void @FnNoArgs() nounwind uwtable sanitize_memory {
+entry:
+  ret void
+}
+
+; CHECK-LABEL: @FnNoArgs
+; CHECK: entry:
+; BSD and Linux systems: instrument with __msan_get_context_state().
+; CHECK: @__msan_get_context_state()
+; %param_shadow:
+; CHECK: getelementptr {{.*}} i32 0, i32 0
+
+; Check the instrumentation prologue for a function with arguments.
+define void @FnOneArg(i32 %input) nounwind uwtable sanitize_memory {
+entry:
+  ret void
+}
+
+; CHECK-LABEL: @FnOneArg
+; CHECK: entry:
+; BSD systems instrument with __msan_get_context_state() that does not accept parameters.
+; CHECK-NOPASS: @__msan_get_context_state()
+; Linux systems: pass __builtin_return_address(0) to __msan_get_context_state_caller().
+; CHECK-PASS: @llvm.returnaddress
+; CHECK-PASS: @__msan_get_context_state_caller({{.*}})
+; %param_shadow:
+; CHECK: getelementptr {{.*}} i32 0, i32 0