hwasan: Implement lazy thread initialization for the interceptor ABI.

[pcc]

Summary

The problem is similar to D55986 but for threads: a process with the
interceptor hwasan library loaded might have some threads started by
instrumented libraries and some by uninstrumented libraries, and we
need to be able to run instrumented code on the latter.

The solution is to perform per-thread initialization lazily. If a
function needs to access shadow memory or add itself to the per-thread
ring buffer its prologue checks to see whether the value in the
sanitizer TLS slot is null, and if so it calls __hwasan_thread_enter
and reloads from the TLS slot. The runtime does the same thing if it
needs to access this data structure.

This change means that the code generator needs to know whether we
are targeting the interceptor runtime, since we don't want to pay
the cost of lazy initialization when targeting a platform with native
hwasan support. A flag -fsanitize-hwaddress-abi={interceptor,platform}
has been introduced for selecting the runtime ABI to target. The
default ABI is set to interceptor since it's assumed that it will
be more common that users will be compiling application code than
platform code.

Because we can no longer assume that the TLS slot is initialized,
the pthread_create interceptor is no longer necessary, so it has
been removed.

Ideally, lazy initialization should only cost one instruction in the
hot path, but at present the call may cause us to spill arguments
to the stack, which means more instructions in the hot path (or
theoretically in the cold path if the spills are moved with shrink
wrapping). With an appropriately chosen calling convention for
the per-thread initialization function (TODO) the hot path should
always need just one instruction and the cold path should need two
instructions with no spilling required.

[eugenis]

LGTM

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

827

This could jump back to the earlier load to save an instruction (at the cost of one extra not-taken branch on the slow path).

Or even better, return the new value from the custom calling convention variant of __hwasan_thread_enter. Not in this change, of course.

[pcc]

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

827

You're right that it could jump back, but I think that I'd prefer not to add this complexity just to remove it later when we switch to the new calling convention.

And yes, the function with the new calling convention would return the thread long value. On ARM64 the code would look like this:

ldr x9, [x9, #64]
cbz .Linit
.Lcont:
...

.Linit:
bl __hwasan_lazy_thread_enter # returns in x9
b .Lcont

clang/include/clang/Basic/CodeGenOptions.h

public:
   /// The preferred width for auto-vectorization transforms. This is intended to
   /// override default transforms based on the width of the architected vector
   /// registers.
   std::string PreferVectorWidth;
 
   /// Set of XRay instrumentation kinds to emit.
   XRayInstrSet XRayInstrumentationBundle;
 
+  std::vector<std::string> DefaultFunctionAttrs;
+
 public:
   // Define accessors/mutators for code generation options of enumeration type.
 #define CODEGENOPT(Name, Bits, Default)
 #define ENUM_CODEGENOPT(Name, Type, Bits, Default) \
   Type get##Name() const { return static_cast<Type>(Name); } \
   void set##Name(Type Value) { Name = static_cast<unsigned>(Value); }
 #include "clang/Basic/CodeGenOptions.def"
 

clang/include/clang/Driver/CC1Options.td

 //===----------------------------------------------------------------------===//
 // CodeGen Options
 //===----------------------------------------------------------------------===//
 
 let Flags = [CC1Option, CC1AsOption, NoDriverOption] in {
 def debug_info_kind_EQ : Joined<["-"], "debug-info-kind=">;
 def debug_info_macro : Flag<["-"], "debug-info-macro">,
   HelpText<"Emit macro debug information">;
+def default_function_attr : Separate<["-"], "default-function-attr">,
+  HelpText<"Apply given attribute to all functions">;
 def dwarf_version_EQ : Joined<["-"], "dwarf-version=">;
 def debugger_tuning_EQ : Joined<["-"], "debugger-tuning=">;
 def fdebug_compilation_dir : Separate<["-"], "fdebug-compilation-dir">,
   HelpText<"The compilation directory to embed in the debug info.">;
 def dwarf_debug_flags : Separate<["-"], "dwarf-debug-flags">,
   HelpText<"The string to embed in the Dwarf debug flags record.">;
 def record_command_line : Separate<["-"], "record-command-line">,
   HelpText<"The string to embed in the .LLVM.command.line section.">;

clang/include/clang/Driver/Options.td

 def fsanitize_address_use_odr_indicator
     : Flag<["-"], "fsanitize-address-use-odr-indicator">,
       Group<f_clang_Group>,
       HelpText<"Enable ODR indicator globals to avoid false ODR violation reports in partially sanitized programs at the cost of an increase in binary size">;
 def fno_sanitize_address_use_odr_indicator
     : Flag<["-"], "fno-sanitize-address-use-odr-indicator">,
       Group<f_clang_Group>,
       HelpText<"Disable ODR indicator globals">;
+def fsanitize_hwaddress_abi_EQ
+    : Joined<["-"], "fsanitize-hwaddress-abi=">,
+      Group<f_clang_Group>,
+      HelpText<"Select the HWAddressSanitizer ABI to target (interceptor or platform, default interceptor)">;
 def fsanitize_recover : Flag<["-"], "fsanitize-recover">, Group<f_clang_Group>;
 def fno_sanitize_recover : Flag<["-"], "fno-sanitize-recover">,
                            Flags<[CoreOption, DriverOption]>,
                            Group<f_clang_Group>;
 def fsanitize_recover_EQ : CommaJoined<["-"], "fsanitize-recover=">,
                            Group<f_clang_Group>,
                            HelpText<"Enable recovery for specified sanitizers">;
 def fno_sanitize_recover_EQ

clang/include/clang/Driver/SanitizerArgs.h

class SanitizerArgs {
   bool CfiCrossDso = false;
   bool CfiICallGeneralizePointers = false;
   int AsanFieldPadding = 0;
   bool SharedRuntime = false;
   bool AsanUseAfterScope = true;
   bool AsanPoisonCustomArrayCookie = false;
   bool AsanGlobalsDeadStripping = false;
   bool AsanUseOdrIndicator = false;
+  std::string HwasanAbi;
   bool LinkCXXRuntimes = false;
   bool NeedPIE = false;
   bool SafeStackRuntime = false;
   bool Stats = false;
   bool TsanMemoryAccess = true;
   bool TsanFuncEntryExit = true;
   bool TsanAtomics = true;
   bool MinimalRuntime = false;

clang/lib/CodeGen/CGCall.cpp

void CodeGenModule::ConstructDefaultFnAttrList(StringRef Name, bool HasOptnone,
   if (getLangOpts().CUDA && getLangOpts().CUDAIsDevice) {
     // Exceptions aren't supported in CUDA device code.
     FuncAttrs.addAttribute(llvm::Attribute::NoUnwind);
 
     // Respect -fcuda-flush-denormals-to-zero.
     if (CodeGenOpts.FlushDenorm)
       FuncAttrs.addAttribute("nvptx-f32ftz", "true");
   }
+
+  for (StringRef Attr : CodeGenOpts.DefaultFunctionAttrs) {
+    StringRef Var, Value;
+    std::tie(Var, Value) = Attr.split('=');
+    FuncAttrs.addAttribute(Var, Value);
+  }
 }
 
 void CodeGenModule::AddDefaultFnAttrs(llvm::Function &F) {
   llvm::AttrBuilder FuncAttrs;
   ConstructDefaultFnAttrList(F.getName(),
                              F.hasFnAttribute(llvm::Attribute::OptimizeNone),
                              /* AttrOnCallsite = */ false, FuncAttrs);
   F.addAttributes(llvm::AttributeList::FunctionIndex, FuncAttrs);

clang/lib/Driver/SanitizerArgs.cpp

if (AllAddedKinds & Address) {
     AsanUseOdrIndicator =
         Args.hasFlag(options::OPT_fsanitize_address_use_odr_indicator,
                      options::OPT_fno_sanitize_address_use_odr_indicator,
                      AsanUseOdrIndicator);
   } else {
     AsanUseAfterScope = false;
   }
 
+  if (AllAddedKinds & HWAddress) {
+    if (Arg *HwasanAbiArg =
+            Args.getLastArg(options::OPT_fsanitize_hwaddress_abi_EQ)) {
+      HwasanAbi = HwasanAbiArg->getValue();
+      if (HwasanAbi != "platform" && HwasanAbi != "interceptor")
+        D.Diag(clang::diag::err_drv_invalid_value)
+            << HwasanAbiArg->getAsString(Args) << HwasanAbi;
+    } else {
+      HwasanAbi = "interceptor";
+    }
+  }
+
   if (AllAddedKinds & SafeStack) {
     // SafeStack runtime is built into the system on Fuchsia.
     SafeStackRuntime = !TC.getTriple().isOSFuchsia();
   }
 
   // Parse -link-cxx-sanitizer flag.
   LinkCXXRuntimes =
       Args.hasArg(options::OPT_fsanitize_link_cxx_runtime) || D.CCCIsCXX();
if (AsanPoisonCustomArrayCookie)
     CmdArgs.push_back("-fsanitize-address-poison-custom-array-cookie");
 
   if (AsanGlobalsDeadStripping)
     CmdArgs.push_back("-fsanitize-address-globals-dead-stripping");
 
   if (AsanUseOdrIndicator)
     CmdArgs.push_back("-fsanitize-address-use-odr-indicator");
 
+  if (!HwasanAbi.empty()) {
+    CmdArgs.push_back("-default-function-attr");
+    CmdArgs.push_back(Args.MakeArgString("hwasan-abi=" + HwasanAbi));
+  }
+
   // MSan: Workaround for PR16386.
   // ASan: This is mainly to help LSan with cases such as
   // https://github.com/google/sanitizers/issues/373
   // We can't make this conditional on -fsanitize=leak, as that flag shouldn't
   // affect compilation.
   if (Sanitizers.has(Memory) || Sanitizers.has(Address))
     CmdArgs.push_back("-fno-assume-sane-operator-new");
 

clang/lib/Frontend/CompilerInvocation.cpp

static bool ParseCodeGenArgs(CodeGenOptions &Opts, ArgList &Args, InputKind IK,
   }
 
   Opts.BranchTargetEnforcement = Args.hasArg(OPT_mbranch_target_enforce);
 
   Opts.KeepStaticConsts = Args.hasArg(OPT_fkeep_static_consts);
 
   Opts.SpeculativeLoadHardening = Args.hasArg(OPT_mspeculative_load_hardening);
 
+  Opts.DefaultFunctionAttrs = Args.getAllArgValues(OPT_default_function_attr);
+
   return Success;
 }
 
 static void ParseDependencyOutputArgs(DependencyOutputOptions &Opts,
                                       ArgList &Args) {
   Opts.OutputFile = Args.getLastArgValue(OPT_dependency_file);
   Opts.Targets = Args.getAllArgValues(OPT_MT);
   Opts.IncludeSystemHeaders = Args.hasArg(OPT_sys_header_deps);

clang/test/CodeGen/default-function-attr.c

+// RUN: %clang_cc1 -default-function-attr foo=bar -emit-llvm %s -o - | FileCheck %s
+
+// CHECK: define void @foo() #[[X:[0-9]+]]
+void foo() {}
+
+// CHECK: attributes #[[X]] = {{.*}} "foo"="bar"

clang/test/Driver/fsanitize.c

 // CHECK-SCUDO-MSAN: error: invalid argument '-fsanitize=scudo' not allowed with '-fsanitize=memory'
 // RUN: %clang -target x86_64-linux-gnu -fsanitize=scudo,thread  %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-SCUDO-TSAN
 // CHECK-SCUDO-TSAN: error: invalid argument '-fsanitize=scudo' not allowed with '-fsanitize=thread'
 // RUN: %clang -target x86_64-linux-gnu -fsanitize=scudo,hwaddress  %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-SCUDO-HWASAN
 // CHECK-SCUDO-HWASAN: error: invalid argument '-fsanitize=scudo' not allowed with '-fsanitize=hwaddress'
 //
 // RUN: %clang -target x86_64-linux-gnu -fsanitize=scudo,kernel-memory  %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-SCUDO-KMSAN
 // CHECK-SCUDO-KMSAN: error: invalid argument '-fsanitize=kernel-memory' not allowed with '-fsanitize=scudo'
+
+// RUN: %clang -target x86_64-linux-gnu -fsanitize=hwaddress %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-HWASAN-INTERCEPTOR-ABI
+// RUN: %clang -target x86_64-linux-gnu -fsanitize=hwaddress -fsanitize-hwaddress-abi=interceptor %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-HWASAN-INTERCEPTOR-ABI
+// RUN: %clang -target x86_64-linux-gnu -fsanitize=hwaddress -fsanitize-hwaddress-abi=platform %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-HWASAN-PLATFORM-ABI
+// RUN: %clang -target x86_64-linux-gnu -fsanitize=hwaddress -fsanitize-hwaddress-abi=foo %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-HWASAN-FOO-ABI
+// CHECK-HWASAN-INTERCEPTOR-ABI: "-default-function-attr" "hwasan-abi=interceptor"
+// CHECK-HWASAN-PLATFORM-ABI: "-default-function-attr" "hwasan-abi=platform"
+// CHECK-HWASAN-FOO-ABI: error: invalid value 'foo' in '-fsanitize-hwaddress-abi=foo'

compiler-rt/lib/hwasan/hwasan_interceptors.cc

 INTERCEPTOR_ALIAS(void *, pvalloc, SIZE_T size);
 INTERCEPTOR_ALIAS(void, cfree, void *ptr);
 INTERCEPTOR_ALIAS(__sanitizer_struct_mallinfo, mallinfo);
 INTERCEPTOR_ALIAS(int, mallopt, int cmd, int value);
 INTERCEPTOR_ALIAS(void, malloc_stats, void);
 #endif
 #endif // HWASAN_WITH_INTERCEPTORS
 
-
-#if HWASAN_WITH_INTERCEPTORS
-extern "C" int pthread_attr_init(void *attr);
-extern "C" int pthread_attr_destroy(void *attr);
-
-struct ThreadStartArg {
-  thread_callback_t callback;
-  void *param;
-};
-
-static void *HwasanThreadStartFunc(void *arg) {
-  __hwasan_thread_enter();
-  ThreadStartArg A = *reinterpret_cast<ThreadStartArg*>(arg);
-  UnmapOrDie(arg, GetPageSizeCached());
-  return A.callback(A.param);
-}
-
-INTERCEPTOR(int, pthread_create, void *th, void *attr, void *(*callback)(void*),
-            void * param) {
-  ScopedTaggingDisabler disabler;
-  ThreadStartArg *A = reinterpret_cast<ThreadStartArg *> (MmapOrDie(
-      GetPageSizeCached(), "pthread_create"));
-  *A = {callback, param};
-  int res = REAL(pthread_create)(UntagPtr(th), UntagPtr(attr),
-                                 &HwasanThreadStartFunc, A);
-  return res;
-}
-#endif // HWASAN_WITH_INTERCEPTORS
-
 static void BeforeFork() {
   StackDepotLockAll();
 }
 
 static void AfterFork() {
   StackDepotUnlockAll();
 }
 

 
 void InitializeInterceptors() {
   static int inited = 0;
   CHECK_EQ(inited, 0);
 
   INTERCEPT_FUNCTION(fork);
 
 #if HWASAN_WITH_INTERCEPTORS
-  INTERCEPT_FUNCTION(pthread_create);
   INTERCEPT_FUNCTION(realloc);
   INTERCEPT_FUNCTION(free);
 #endif
 
   inited = 1;
 }
 } // namespace __hwasan

compiler-rt/lib/hwasan/hwasan_linux.cc

if (*(uptr *)get_android_tls_ptr() != kMagicValue) {
     Die();
   }
 }
 #else
 void AndroidTestTlsSlot() {}
 #endif
 
 Thread *GetCurrentThread() {
-  auto *R = (StackAllocationsRingBuffer*)GetCurrentThreadLongPtr();
+  uptr *ThreadLong = GetCurrentThreadLongPtr();
+#if HWASAN_WITH_INTERCEPTORS
+  if (!*ThreadLong)
+    __hwasan_thread_enter();
+#endif
+  auto *R = (StackAllocationsRingBuffer *)ThreadLong;
   return hwasanThreadList().GetThreadByBufferAddress((uptr)(R->Next()));
 }
 
 struct AccessInfo {
   uptr addr;
   uptr size;
   bool is_store;
   bool is_load;

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

private:
 
   Function *HwasanCtorFunction;
 
   Function *HwasanMemoryAccessCallback[2][kNumberOfAccessSizes];
   Function *HwasanMemoryAccessCallbackSized[2];
 
   Function *HwasanTagMemoryFunc;
   Function *HwasanGenerateTagFunc;
+  Function *HwasanThreadEnterFunc;
 
   Constant *ShadowGlobal;
 
   Value *LocalDynamicShadow = nullptr;
   GlobalValue *ThreadPtrGlobal = nullptr;
 };
 
 } // end anonymous namespace
HWAsanMemmove = checkSanitizerInterfaceFunction(M.getOrInsertFunction(
       MemIntrinCallbackPrefix + "memmove", IRB.getInt8PtrTy(),
       IRB.getInt8PtrTy(), IRB.getInt8PtrTy(), IntptrTy));
   HWAsanMemcpy = checkSanitizerInterfaceFunction(M.getOrInsertFunction(
       MemIntrinCallbackPrefix + "memcpy", IRB.getInt8PtrTy(),
       IRB.getInt8PtrTy(), IRB.getInt8PtrTy(), IntptrTy));
   HWAsanMemset = checkSanitizerInterfaceFunction(M.getOrInsertFunction(
       MemIntrinCallbackPrefix + "memset", IRB.getInt8PtrTy(),
       IRB.getInt8PtrTy(), IRB.getInt32Ty(), IntptrTy));
+
+  HwasanThreadEnterFunc = checkSanitizerInterfaceFunction(
+      M.getOrInsertFunction("__hwasan_thread_enter", IRB.getVoidTy()));
 }
 
 Value *HWAddressSanitizer::getDynamicShadowNonTls(IRBuilder<> &IRB) {
   // Generate code only when dynamic addressing is needed.
   if (Mapping.Offset != kDynamicShadowSentinel)
     return nullptr;
 
   if (Mapping.InGlobal) {

 Value *HWAddressSanitizer::emitPrologue(IRBuilder<> &IRB,
                                         bool WithFrameRecord) {
   if (!Mapping.InTls)
     return getDynamicShadowNonTls(IRB);
 
   Value *SlotPtr = getHwasanThreadSlotPtr(IRB, IntptrTy);
   assert(SlotPtr);
 
-  Value *ThreadLong = IRB.CreateLoad(SlotPtr);
+  Instruction *ThreadLong = IRB.CreateLoad(SlotPtr);
+
+  Function *F = IRB.GetInsertBlock()->getParent();
+  if (F->getFnAttribute("hwasan-abi").getValueAsString() == "interceptor") {
+    Value *ThreadLongEqZero =
+        IRB.CreateICmpEQ(ThreadLong, ConstantInt::get(IntptrTy, 0));
+    auto *Br = cast<BranchInst>(SplitBlockAndInsertIfThen(
+        ThreadLongEqZero, cast<Instruction>(ThreadLongEqZero)->getNextNode(),
+        false, MDBuilder(*C).createBranchWeights(1, 100000)));
+
+    IRB.SetInsertPoint(Br);
+    // FIXME: This should call a new runtime function with a custom calling
+    // convention to avoid needing to spill all arguments here.
+    IRB.CreateCall(HwasanThreadEnterFunc);
+    LoadInst *ReloadThreadLong = IRB.CreateLoad(SlotPtr);
+
+    IRB.SetInsertPoint(&*Br->getSuccessor(0)->begin());
+    PHINode *ThreadLongPhi = IRB.CreatePHI(IntptrTy, 2);
+    ThreadLongPhi->addIncoming(ThreadLong, ThreadLong->getParent());
+    ThreadLongPhi->addIncoming(ReloadThreadLong, ReloadThreadLong->getParent());
+    ThreadLong = ThreadLongPhi;
+  }
+
   // Extract the address field from ThreadLong. Unnecessary on AArch64 with TBI.
   Value *ThreadLongMaybeUntagged =
       TargetTriple.isAArch64() ? ThreadLong : untagPointer(IRB, ThreadLong);
 
   if (WithFrameRecord) {
     // Prepare ring buffer data.
-    Function *F = IRB.GetInsertBlock()->getParent();
     auto PC = IRB.CreatePtrToInt(F, IntptrTy);
     auto GetStackPointerFn =
         Intrinsic::getDeclaration(F->getParent(), Intrinsic::frameaddress);
     Value *SP = IRB.CreatePtrToInt(
         IRB.CreateCall(GetStackPointerFn,
                        {Constant::getNullValue(IRB.getInt32Ty())}),
         IntptrTy);
     // Mix SP and PC. TODO: also add the tag to the mix.

llvm/test/Instrumentation/HWAddressSanitizer/lazy-thread-init.ll

+; RUN: opt -S -hwasan < %s | FileCheck %s
+
+target triple = "x86_64-unknown-linux-gnu"
+
+declare void @bar([16 x i32]* %p)
+
+define void @foo() sanitize_hwaddress "hwasan-abi"="interceptor" {
+  ; CHECK: [[LOAD:%[^ ]*]] = load i64, i64* @__hwasan_tls
+  ; CHECK: [[ICMP:%[^ ]*]] = icmp eq i64 [[LOAD]], 0
+  ; CHECK: br i1 [[ICMP]], label %[[INIT:[^,]*]], label %[[CONT:[^,]*]], !prof [[PROF:![0-9]+]]
+
+  ; CHECK: [[INIT]]:
+  ; CHECK: call void @__hwasan_thread_enter()
+  ; CHECK: [[RELOAD:%[^ ]*]] = load i64, i64* @__hwasan_tls
+  ; CHECK: br label %[[CONT]]
+
+  ; CHECK: [[CONT]]:
+  ; CHECK: phi i64 [ [[LOAD]], %0 ], [ [[RELOAD]], %[[INIT]] ]
+
+  %p = alloca [16 x i32]
+  call void @bar([16 x i32]* %p)
+  ret void
+}
+
+; CHECK: [[PROF]] = !{!"branch_weights", i32 1, i32 100000}