[SanitizerBinaryMetadata] Declare callbacks extern weak

[melver]

Summary

Declare callbacks extern weak (if no existing declaration exists), and
only call if the function address is non-null.

This allows to attach semantic metadata to binaries where no user of
that metadata exists, avoiding to have to link empty stub callbacks.

Once the binary is linked (statically or dynamically) against a tool
runtime that implements the callbacks, the respective callbacks will be
called. This vastly simplifies gradual deployment of tools using the
metadata, esp. avoiding having to recompile large codebases with
different compiler flags (which negatively impacts compiler caches).

[vitalybuka]

llvm/lib/Transforms/Utils/ModuleUtils.cpp

181

maybe we want skip only if it's defined?
What is the reason to respect declaration?

[melver]

Only skip setting weak linkage if function is defined in module.

llvm/lib/Transforms/Utils/ModuleUtils.cpp

181

Right, the AsmPrinter complained on definition only:

Should never emit this
UNREACHABLE executed at /usr/local/google/home/elver/third_party/llvm/llvm-project/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp:633!

Changing it to skip if definition works.

[dvyukov]

Nice!

compiler-rt/test/metadata/nocallback.cpp

+// RUN: %clangxx %s -o %t -fexperimental-sanitize-metadata=all && %t | FileCheck %s
+
+// Test that the compiler emits weak declarations to the callbacks, which are
+// not called if they do not exist.
+
+#include <stdio.h>
+
+int main() {
+  printf("main\n");
+  return 0;
+}
+
+// CHECK: main

llvm/include/llvm/Transforms/Utils/ModuleUtils.h

 void appendToGlobalDtors(Module &M, Function *F, int Priority,
                          Constant *Data = nullptr);
 
 /// Sets the KCFI type for the function. Used for compiler-generated functions
 /// that are indirectly called in instrumented code.
 void setKCFIType(Module &M, Function &F, StringRef MangledType);
 
 FunctionCallee declareSanitizerInitFunction(Module &M, StringRef InitName,
-                                            ArrayRef<Type *> InitArgTypes);
+                                            ArrayRef<Type *> InitArgTypes,
+                                            bool Weak = false);
 
 /// Creates sanitizer constructor function.
 /// \return Returns pointer to constructor.
 Function *createSanitizerCtor(Module &M, StringRef CtorName);
 
 /// Creates sanitizer constructor function, and calls sanitizer's init
 /// function from it.
 /// \return Returns pair of pointers to constructor, and init functions
 /// respectively.
 std::pair<Function *, FunctionCallee> createSanitizerCtorAndInitFunctions(
     Module &M, StringRef CtorName, StringRef InitName,
     ArrayRef<Type *> InitArgTypes, ArrayRef<Value *> InitArgs,
-    StringRef VersionCheckName = StringRef());
+    StringRef VersionCheckName = StringRef(), bool Weak = false);
 
 /// Creates sanitizer constructor function lazily. If a constructor and init
 /// function already exist, this function returns it. Otherwise it calls \c
 /// createSanitizerCtorAndInitFunctions. The FunctionsCreatedCallback is invoked
 /// in that case, passing the new Ctor and Init function.
 ///
 /// \return Returns pair of pointers to constructor, and init functions
 /// respectively.
 std::pair<Function *, FunctionCallee> getOrCreateSanitizerCtorAndInitFunctions(
     Module &M, StringRef CtorName, StringRef InitName,
     ArrayRef<Type *> InitArgTypes, ArrayRef<Value *> InitArgs,
     function_ref<void(Function *, FunctionCallee)> FunctionsCreatedCallback,
-    StringRef VersionCheckName = StringRef());
+    StringRef VersionCheckName = StringRef(), bool Weak = false);
 
 /// Rename all the anon globals in the module using a hash computed from
 /// the list of public globals in the module.
 bool nameUnamedGlobals(Module &M);
 
 /// Adds global values to the llvm.used list.
 void appendToUsed(Module &M, ArrayRef<GlobalValue *> Values);
 

llvm/lib/Transforms/Instrumentation/SanitizerBinaryMetadata.cpp

 
 // The only instances of MetadataInfo are the constants above, so a set of
 // them may simply store pointers to them. To deterministically generate code,
 // we need to use a set with stable iteration order, such as SetVector.
 using MetadataInfoSet = SetVector<const MetadataInfo *>;
 
 //===--- Command-line options ---------------------------------------------===//
 
+cl::opt<bool> ClWeakCallbacks(
+    "sanitizer-metadata-weak-callbacks",
+    cl::desc("Declare callbacks extern weak, and only call if non-null."),
+    cl::Hidden, cl::init(true));
+
 cl::opt<bool> ClEmitCovered("sanitizer-metadata-covered",
                             cl::desc("Emit PCs for covered functions."),
                             cl::Hidden, cl::init(false));
 cl::opt<bool> ClEmitAtomics("sanitizer-metadata-atomics",
                             cl::desc("Emit PCs for atomic operations."),
                             cl::Hidden, cl::init(false));
 cl::opt<bool> ClEmitUAR("sanitizer-metadata-uar",
                         cl::desc("Emit PCs for start of functions that are "
bool SanitizerBinaryMetadata::run() {
   auto *Version = ConstantInt::get(Int32Ty, getVersion());
 
   for (const MetadataInfo *MI : MIS) {
     const std::array<Value *, InitTypes.size()> InitArgs = {
         Version,
         getSectionMarker(getSectionStart(MI->SectionSuffix), Int8PtrTy),
         getSectionMarker(getSectionEnd(MI->SectionSuffix), Int8PtrTy),
     };
+    // We declare the _add and _del functions as weak, and only call them if
+    // there is a valid symbol linked. This allows building binaries with
+    // semantic metadata, but without having callbacks. When a tool that wants
+    // the metadata is linked which provides the callbacks, they will be called.
     Function *Ctor =
         createSanitizerCtorAndInitFunctions(
             Mod, (MI->FunctionPrefix + ".module_ctor").str(),
-            (MI->FunctionPrefix + "_add").str(), InitTypes, InitArgs)
+            (MI->FunctionPrefix + "_add").str(), InitTypes, InitArgs,
+            /*VersionCheckName=*/StringRef(), /*Weak=*/ClWeakCallbacks)
             .first;
     Function *Dtor =
         createSanitizerCtorAndInitFunctions(
             Mod, (MI->FunctionPrefix + ".module_dtor").str(),
-            (MI->FunctionPrefix + "_del").str(), InitTypes, InitArgs)
+            (MI->FunctionPrefix + "_del").str(), InitTypes, InitArgs,
+            /*VersionCheckName=*/StringRef(), /*Weak=*/ClWeakCallbacks)
             .first;
     Constant *CtorData = nullptr;
     Constant *DtorData = nullptr;
     if (TargetTriple.supportsCOMDAT()) {
       // Use COMDAT to deduplicate constructor/destructor function.
       Ctor->setComdat(Mod.getOrInsertComdat(Ctor->getName()));
       Dtor->setComdat(Mod.getOrInsertComdat(Dtor->getName()));
       CtorData = Ctor;

llvm/lib/Transforms/Utils/ModuleUtils.cpp

void llvm::setKCFIType(Module &M, Function &F, StringRef MangledType) {
   // we use the same patchable-function-prefix.
   if (auto *MD = mdconst::extract_or_null<ConstantInt>(
           M.getModuleFlag("kcfi-offset"))) {
     if (unsigned Offset = MD->getZExtValue())
       F.addFnAttr("patchable-function-prefix", std::to_string(Offset));
   }
 }
 
-FunctionCallee
-llvm::declareSanitizerInitFunction(Module &M, StringRef InitName,
-                                   ArrayRef<Type *> InitArgTypes) {
+FunctionCallee llvm::declareSanitizerInitFunction(Module &M, StringRef InitName,
+                                                  ArrayRef<Type *> InitArgTypes,
+                                                  bool Weak) {
   assert(!InitName.empty() && "Expected init function name");
-  return M.getOrInsertFunction(
-      InitName,
-      FunctionType::get(Type::getVoidTy(M.getContext()), InitArgTypes, false),
-      AttributeList());
+  auto *VoidTy = Type::getVoidTy(M.getContext());
+  auto *FnTy = FunctionType::get(VoidTy, InitArgTypes, false);
+  auto FnCallee = M.getOrInsertFunction(InitName, FnTy);
+  auto *Fn = cast<Function>(FnCallee.getCallee());
+  if (Weak && Fn->isDeclaration())
+    Fn->setLinkage(Function::ExternalWeakLinkage);
+  return FnCallee;
 }
 
 Function *llvm::createSanitizerCtor(Module &M, StringRef CtorName) {
   Function *Ctor = Function::createWithDefaultAttr(
       FunctionType::get(Type::getVoidTy(M.getContext()), false),
       GlobalValue::InternalLinkage, M.getDataLayout().getProgramAddressSpace(),
       CtorName, &M);
   Ctor->addFnAttr(Attribute::NoUnwind);
   setKCFIType(M, *Ctor, "_ZTSFvvE"); // void (*)(void)
   BasicBlock *CtorBB = BasicBlock::Create(M.getContext(), "", Ctor);
   ReturnInst::Create(M.getContext(), CtorBB);
   // Ensure Ctor cannot be discarded, even if in a comdat.
   appendToUsed(M, {Ctor});
   return Ctor;
 }
 
 std::pair<Function *, FunctionCallee> llvm::createSanitizerCtorAndInitFunctions(
     Module &M, StringRef CtorName, StringRef InitName,
     ArrayRef<Type *> InitArgTypes, ArrayRef<Value *> InitArgs,
-    StringRef VersionCheckName) {
+    StringRef VersionCheckName, bool Weak) {
   assert(!InitName.empty() && "Expected init function name");
   assert(InitArgs.size() == InitArgTypes.size() &&
          "Sanitizer's init function expects different number of arguments");
   FunctionCallee InitFunction =
-      declareSanitizerInitFunction(M, InitName, InitArgTypes);
+      declareSanitizerInitFunction(M, InitName, InitArgTypes, Weak);
   Function *Ctor = createSanitizerCtor(M, CtorName);
-  IRBuilder<> IRB(Ctor->getEntryBlock().getTerminator());
+  IRBuilder<> IRB(M.getContext());
+
+  BasicBlock *RetBB = &Ctor->getEntryBlock();
+  if (Weak) {
+    RetBB->setName("ret");
+    auto *EntryBB = BasicBlock::Create(M.getContext(), "entry", Ctor, RetBB);
+    auto *CallInitBB =
+        BasicBlock::Create(M.getContext(), "callfunc", Ctor, RetBB);
+    auto *InitFn = cast<Function>(InitFunction.getCallee());
+    auto *InitFnPtr =
+        PointerType::get(InitFn->getType(), InitFn->getAddressSpace());
+    IRB.SetInsertPoint(EntryBB);
+    Value *InitNotNull =
+        IRB.CreateICmpNE(InitFn, ConstantPointerNull::get(InitFnPtr));
+    IRB.CreateCondBr(InitNotNull, CallInitBB, RetBB);
+    IRB.SetInsertPoint(CallInitBB);
+  } else {
+    IRB.SetInsertPoint(RetBB->getTerminator());
+  }
+
   IRB.CreateCall(InitFunction, InitArgs);
   if (!VersionCheckName.empty()) {
     FunctionCallee VersionCheckFunction = M.getOrInsertFunction(
         VersionCheckName, FunctionType::get(IRB.getVoidTy(), {}, false),
         AttributeList());
     IRB.CreateCall(VersionCheckFunction, {});
   }
+
+  if (Weak)
+    IRB.CreateBr(RetBB);
+
   return std::make_pair(Ctor, InitFunction);
 }
 
 std::pair<Function *, FunctionCallee>
 llvm::getOrCreateSanitizerCtorAndInitFunctions(
     Module &M, StringRef CtorName, StringRef InitName,
     ArrayRef<Type *> InitArgTypes, ArrayRef<Value *> InitArgs,
     function_ref<void(Function *, FunctionCallee)> FunctionsCreatedCallback,
-    StringRef VersionCheckName) {
+    StringRef VersionCheckName, bool Weak) {
   assert(!CtorName.empty() && "Expected ctor function name");
 
   if (Function *Ctor = M.getFunction(CtorName))
     // FIXME: Sink this logic into the module, similar to the handling of
     // globals. This will make moving to a concurrent model much easier.
     if (Ctor->arg_empty() ||
         Ctor->getReturnType() == Type::getVoidTy(M.getContext()))
-      return {Ctor, declareSanitizerInitFunction(M, InitName, InitArgTypes)};
+      return {Ctor,
+              declareSanitizerInitFunction(M, InitName, InitArgTypes, Weak)};
 
   Function *Ctor;
   FunctionCallee InitFunction;
   std::tie(Ctor, InitFunction) = llvm::createSanitizerCtorAndInitFunctions(
-      M, CtorName, InitName, InitArgTypes, InitArgs, VersionCheckName);
+      M, CtorName, InitName, InitArgTypes, InitArgs, VersionCheckName, Weak);
   FunctionsCreatedCallback(Ctor, InitFunction);
   return std::make_pair(Ctor, InitFunction);
 }
 
 void llvm::filterDeadComdatFunctions(
     SmallVectorImpl<Function *> &DeadComdatFunctions) {
   SmallPtrSet<Function *, 32> MaybeDeadFunctions;
   SmallPtrSet<Comdat *, 32> MaybeDeadComdats;

llvm/test/Instrumentation/SanitizerBinaryMetadata/atomics.ll

entry:
   ret void
 }
 ; CHECK-LABEL: atomic128_cas_seq_cst{{.*}}!pcsections !0
 ; CHECK: !pcsections !2
 
 ; Check that callbacks are emitted.
 
 ; CHECK-LABEL: __sanitizer_metadata_atomics.module_ctor
-; CHECK: call void @__sanitizer_metadata_atomics_add(i32 1, ptr @__start_sanmd_atomics, ptr @__stop_sanmd_atomics)
+; CHECK-DAG: entry:
+; CHECK-NEXT:  br i1 icmp ne (ptr @__sanitizer_metadata_atomics_add, ptr null), label %callfunc, label %ret
+; CHECK-DAG: callfunc:
+; CHECK-NEXT:  call void @__sanitizer_metadata_atomics_add(i32 1, ptr @__start_sanmd_atomics, ptr @__stop_sanmd_atomics)
+; CHECK-NEXT:  br label %ret
+; CHECK-DAG: ret:
+; CHECK-NEXT:  ret void
 
 ; CHECK-LABEL: __sanitizer_metadata_atomics.module_dtor
-; CHECK: call void @__sanitizer_metadata_atomics_del(i32 1, ptr @__start_sanmd_atomics, ptr @__stop_sanmd_atomics)
+; CHECK-DAG: entry:
+; CHECK-NEXT:  br i1 icmp ne (ptr @__sanitizer_metadata_atomics_del, ptr null), label %callfunc, label %ret
+; CHECK-DAG: callfunc:
+; CHECK-NEXT:  call void @__sanitizer_metadata_atomics_del(i32 1, ptr @__start_sanmd_atomics, ptr @__stop_sanmd_atomics)
+; CHECK-NEXT:  br label %ret
+; CHECK-DAG: ret:
+; CHECK-NEXT:  ret void
 
 ; CHECK-LABEL: __sanitizer_metadata_covered.module_ctor
-; CHECK: call void @__sanitizer_metadata_covered_add(i32 1, ptr @__start_sanmd_covered, ptr @__stop_sanmd_covered)
+; CHECK-DAG: entry:
+; CHECK-NEXT:  br i1 icmp ne (ptr @__sanitizer_metadata_covered_add, ptr null), label %callfunc, label %ret
+; CHECK-DAG: callfunc:
+; CHECK-NEXT:  call void @__sanitizer_metadata_covered_add(i32 1, ptr @__start_sanmd_covered, ptr @__stop_sanmd_covered)
+; CHECK-NEXT:  br label %ret
+; CHECK-DAG: ret:
+; CHECK-NEXT:  ret void
 
 ; CHECK-LABEL: __sanitizer_metadata_covered.module_dtor
-; CHECK: call void @__sanitizer_metadata_covered_del(i32 1, ptr @__start_sanmd_covered, ptr @__stop_sanmd_covered)
+; CHECK-DAG: entry:
+; CHECK-NEXT:  br i1 icmp ne (ptr @__sanitizer_metadata_covered_del, ptr null), label %callfunc, label %ret
+; CHECK-DAG: callfunc:
+; CHECK-NEXT:  call void @__sanitizer_metadata_covered_del(i32 1, ptr @__start_sanmd_covered, ptr @__stop_sanmd_covered)
+; CHECK-NEXT:  br label %ret
+; CHECK-DAG: ret:
+; CHECK-NEXT:  ret void
 
 ; CHECK: !0 = !{!"sanmd_covered", !1}
 ; CHECK: !1 = !{i32 1}
 ; CHECK: !2 = !{!"sanmd_atomics"}