4 files changed, 55 insertions, 5 deletions

diff --git a/clang/lib/CodeGen/CGClass.cpp b/clang/lib/CodeGen/CGClass.cpp
index 0853d6aa1fb..2adb74aa0fc 100644
--- a/clang/lib/CodeGen/CGClass.cpp
+++ b/clang/lib/CodeGen/CGClass.cpp
@@ -2607,10 +2607,22 @@ void CodeGenFunction::EmitVTablePtrCheck(const CXXRecordDecl *RD,
   auto TypeId = CGM.CreateCfiIdForTypeMetadata(MD);
   if (CGM.getCodeGenOpts().SanitizeCfiCrossDso && TypeId) {
     EmitCfiSlowPathCheck(M, BitSetTest, TypeId, CastedVTable, StaticData);
-  } else {
-    EmitCheck(std::make_pair(BitSetTest, M), "cfi_check_fail", StaticData,
-              CastedVTable);
+    return;
   }
+
+  if (CGM.getCodeGenOpts().SanitizeTrap.has(M)) {
+    EmitTrapCheck(BitSetTest);
+    return;
+  }
+
+  llvm::Value *AllVtables = llvm::MetadataAsValue::get(
+      CGM.getLLVMContext(),
+      llvm::MDString::get(CGM.getLLVMContext(), "all-vtables"));
+  llvm::Value *ValidVtable =
+      Builder.CreateCall(CGM.getIntrinsic(llvm::Intrinsic::bitset_test),
+                         {CastedVTable, AllVtables});
+  EmitCheck(std::make_pair(BitSetTest, M), "cfi_check_fail", StaticData,
+            {CastedVTable, ValidVtable});
 }
 
 // FIXME: Ideally Expr::IgnoreParenNoopCasts should do this, but it doesn't do
diff --git a/clang/lib/CodeGen/CGExpr.cpp b/clang/lib/CodeGen/CGExpr.cpp
index 10547a4ecf0..0855b450797 100644
--- a/clang/lib/CodeGen/CGExpr.cpp
+++ b/clang/lib/CodeGen/CGExpr.cpp
@@ -2636,6 +2636,14 @@ void CodeGenFunction::EmitCfiCheckFail() {
   Address CheckKindAddr(V, getIntAlign());
   llvm::Value *CheckKind = Builder.CreateLoad(CheckKindAddr);
 
+  llvm::Value *AllVtables = llvm::MetadataAsValue::get(
+      CGM.getLLVMContext(),
+      llvm::MDString::get(CGM.getLLVMContext(), "all-vtables"));
+  llvm::Value *ValidVtable = Builder.CreateZExt(
+      Builder.CreateCall(CGM.getIntrinsic(llvm::Intrinsic::bitset_test),
+                         {Addr, AllVtables}),
+      IntPtrTy);
+
   const std::pair<int, SanitizerMask> CheckKinds[] = {
       {CFITCK_VCall, SanitizerKind::CFIVCall},
       {CFITCK_NVCall, SanitizerKind::CFINVCall},
@@ -2649,7 +2657,8 @@ void CodeGenFunction::EmitCfiCheckFail() {
     SanitizerMask Mask = CheckKindMaskPair.second;
     llvm::Value *Cond =
         Builder.CreateICmpNE(CheckKind, llvm::ConstantInt::get(Int8Ty, Kind));
-    EmitCheck(std::make_pair(Cond, Mask), "cfi_check_fail", {}, {Data, Addr});
+    EmitCheck(std::make_pair(Cond, Mask), "cfi_check_fail", {},
+              {Data, Addr, ValidVtable});
   }
 
   FinishFunction();
@@ -3970,7 +3979,8 @@ RValue CodeGenFunction::EmitCall(QualType CalleeType, llvm::Value *Callee,
                            CastedCallee, StaticData);
     } else {
       EmitCheck(std::make_pair(BitSetTest, SanitizerKind::CFIICall),
-                "cfi_check_fail", StaticData, CastedCallee);
+                "cfi_check_fail", StaticData,
+                {CastedCallee, llvm::UndefValue::get(IntPtrTy)});
     }
   }
 
diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
index 4d7f62750a1..a30e62499fa 100644
--- a/clang/lib/CodeGen/CodeGenModule.cpp
+++ b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -4021,6 +4021,20 @@ llvm::Metadata *CodeGenModule::CreateMetadataIdentifierForType(QualType T) {
   return InternalId;
 }
 
+/// Returns whether this module needs the "all-vtables" bitset.
+bool CodeGenModule::NeedAllVtablesBitSet() const {
+  // Returns true if at least one of vtable-based CFI checkers is enabled and
+  // is not in the trapping mode.
+  return ((LangOpts.Sanitize.has(SanitizerKind::CFIVCall) &&
+           !CodeGenOpts.SanitizeTrap.has(SanitizerKind::CFIVCall)) ||
+          (LangOpts.Sanitize.has(SanitizerKind::CFINVCall) &&
+           !CodeGenOpts.SanitizeTrap.has(SanitizerKind::CFINVCall)) ||
+          (LangOpts.Sanitize.has(SanitizerKind::CFIDerivedCast) &&
+           !CodeGenOpts.SanitizeTrap.has(SanitizerKind::CFIDerivedCast)) ||
+          (LangOpts.Sanitize.has(SanitizerKind::CFIUnrelatedCast) &&
+           !CodeGenOpts.SanitizeTrap.has(SanitizerKind::CFIUnrelatedCast)));
+}
+
 void CodeGenModule::CreateVTableBitSetEntry(llvm::NamedMDNode *BitsetsMD,
                                             llvm::GlobalVariable *VTable,
                                             CharUnits Offset,
@@ -4043,6 +4057,17 @@ void CodeGenModule::CreateVTableBitSetEntry(llvm::NamedMDNode *BitsetsMD,
       BitsetsMD->addOperand(llvm::MDTuple::get(getLLVMContext(), BitsetOps2));
     }
   }
+
+  if (NeedAllVtablesBitSet()) {
+    llvm::Metadata *MD = llvm::MDString::get(getLLVMContext(), "all-vtables");
+    llvm::Metadata *BitsetOps[] = {
+        MD, llvm::ConstantAsMetadata::get(VTable),
+        llvm::ConstantAsMetadata::get(
+            llvm::ConstantInt::get(Int64Ty, Offset.getQuantity()))};
+    // Avoid adding a node to BitsetsMD twice.
+    if (!llvm::MDTuple::getIfExists(getLLVMContext(), BitsetOps))
+      BitsetsMD->addOperand(llvm::MDTuple::get(getLLVMContext(), BitsetOps));
+  }
 }
 
 // Fills in the supplied string map with the set of target features for the
diff --git a/clang/lib/CodeGen/CodeGenModule.h b/clang/lib/CodeGen/CodeGenModule.h
index 00119c3ae4e..f8084ef43d9 100644
--- a/clang/lib/CodeGen/CodeGenModule.h
+++ b/clang/lib/CodeGen/CodeGenModule.h
@@ -1128,6 +1128,9 @@ public:
   /// Create a bitset entry for the given function and add it to BitsetsMD.
   void CreateFunctionBitSetEntry(const FunctionDecl *FD, llvm::Function *F);
 
+  /// Returns whether this module needs the "all-vtables" bitset.
+  bool NeedAllVtablesBitSet() const;
+
   /// Create a bitset entry for the given vtable and add it to BitsetsMD.
   void CreateVTableBitSetEntry(llvm::NamedMDNode *BitsetsMD,
                                llvm::GlobalVariable *VTable, CharUnits Offset,