7 files changed, 102 insertions, 1 deletions

diff --git a/test/tools/llvm-cfi-verify/AArch64/Inputs/function-only-check.o b/test/tools/llvm-cfi-verify/AArch64/Inputs/function-only-check.o
new file mode 100644
index 00000000000..49142ff8170
--- /dev/null
+++ b/test/tools/llvm-cfi-verify/AArch64/Inputs/function-only-check.o
diff --git a/test/tools/llvm-cfi-verify/AArch64/function-only-check.s b/test/tools/llvm-cfi-verify/AArch64/function-only-check.s
new file mode 100644
index 00000000000..30460c8cefc
--- /dev/null
+++ b/test/tools/llvm-cfi-verify/AArch64/function-only-check.s
@@ -0,0 +1,9 @@
+# RUN: llvm-cfi-verify -search-length-undef=6 %S/Inputs/function-only-check.o | FileCheck %s
+
+# CHECK-LABEL: {{^Instruction: .* \(PROTECTED\)}}
+# CHECK-NEXT: tiny.cc:9
+
+# CHECK: Expected Protected: 1 (100.00%)
+# CHECK: Unexpected Protected: 0 (0.00%)
+# CHECK: Expected Unprotected: 0 (0.00%)
+# CHECK: Unexpected Unprotected (BAD): 0 (0.00%)
diff --git a/test/tools/llvm-cfi-verify/X86/Inputs/function-only-check.o b/test/tools/llvm-cfi-verify/X86/Inputs/function-only-check.o
new file mode 100644
index 00000000000..099076a69ef
--- /dev/null
+++ b/test/tools/llvm-cfi-verify/X86/Inputs/function-only-check.o
diff --git a/test/tools/llvm-cfi-verify/X86/function-only-check.s b/test/tools/llvm-cfi-verify/X86/function-only-check.s
new file mode 100644
index 00000000000..250f5dc5ab7
--- /dev/null
+++ b/test/tools/llvm-cfi-verify/X86/function-only-check.s
@@ -0,0 +1,8 @@
+# RUN: llvm-cfi-verify %S/Inputs/function-only-check.o | FileCheck %s
+
+# CHECK-LABEL: {{^Instruction: .* \(PROTECTED\)}}
+
+# CHECK: Expected Protected: 1 (100.00%)
+# CHECK: Unexpected Protected: 0 (0.00%)
+# CHECK: Expected Unprotected: 0 (0.00%)
+# CHECK: Unexpected Unprotected (BAD): 0 (0.00%)
diff --git a/tools/llvm-cfi-verify/lib/FileAnalysis.cpp b/tools/llvm-cfi-verify/lib/FileAnalysis.cpp
index 29819d8d28e..8dc0802c465 100644
--- a/tools/llvm-cfi-verify/lib/FileAnalysis.cpp
+++ b/tools/llvm-cfi-verify/lib/FileAnalysis.cpp
@@ -105,6 +105,9 @@ Expected<FileAnalysis> FileAnalysis::Create(StringRef Filename) {
   if (auto SectionParseResponse = Analysis.parseCodeSections())
     return std::move(SectionParseResponse);
 
+  if (auto SymbolTableParseResponse = Analysis.parseSymbolTable())
+    return std::move(SymbolTableParseResponse);
+
   return std::move(Analysis);
 }
 
@@ -165,7 +168,18 @@ const Instr &FileAnalysis::getInstructionOrDie(uint64_t Address) const {
 
 bool FileAnalysis::isCFITrap(const Instr &InstrMeta) const {
   const auto &InstrDesc = MII->get(InstrMeta.Instruction.getOpcode());
-  return InstrDesc.isTrap();
+  return InstrDesc.isTrap() || willTrapOnCFIViolation(InstrMeta);
+}
+
+bool FileAnalysis::willTrapOnCFIViolation(const Instr &InstrMeta) const {
+  const auto &InstrDesc = MII->get(InstrMeta.Instruction.getOpcode());
+  if (!InstrDesc.isCall())
+    return false;
+  uint64_t Target;
+  if (!MIA->evaluateBranch(InstrMeta.Instruction, InstrMeta.VMAddress,
+                           InstrMeta.InstructionSize, Target))
+    return false;
+  return TrapOnFailFunctionAddresses.count(Target) > 0;
 }
 
 bool FileAnalysis::canFallThrough(const Instr &InstrMeta) const {
@@ -431,6 +445,12 @@ Error FileAnalysis::parseCodeSections() {
     if (!(object::ELFSectionRef(Section).getFlags() & ELF::SHF_EXECINSTR))
       continue;
 
+    // Avoid checking the PLT since it produces spurious failures on AArch64
+    // when ignoring DWARF data.
+    StringRef SectionName;
+    if (!Section.getName(SectionName) && SectionName == ".plt")
+      continue;
+
     StringRef SectionContents;
     if (Section.getContents(SectionContents))
       return make_error<StringError>("Failed to retrieve section contents",
@@ -518,6 +538,40 @@ void FileAnalysis::addInstruction(const Instr &Instruction) {
   }
 }
 
+Error FileAnalysis::parseSymbolTable() {
+  // Functions that will trap on CFI violations.
+  SmallSet<StringRef, 4> TrapOnFailFunctions;
+  TrapOnFailFunctions.insert("__cfi_slowpath");
+  TrapOnFailFunctions.insert("__cfi_slowpath_diag");
+  TrapOnFailFunctions.insert("abort");
+
+  // Look through the list of symbols for functions that will trap on CFI
+  // violations.
+  for (auto &Sym : Object->symbols()) {
+    auto SymNameOrErr = Sym.getName();
+    if (!SymNameOrErr)
+      consumeError(SymNameOrErr.takeError());
+    else if (TrapOnFailFunctions.count(*SymNameOrErr) > 0) {
+      auto AddrOrErr = Sym.getAddress();
+      if (!AddrOrErr)
+        consumeError(AddrOrErr.takeError());
+      else
+        TrapOnFailFunctionAddresses.insert(*AddrOrErr);
+    }
+  }
+  if (auto *ElfObject = dyn_cast<object::ELFObjectFileBase>(Object)) {
+    for (const auto &Addr : ElfObject->getPltAddresses()) {
+      object::SymbolRef Sym(Addr.first, Object);
+      auto SymNameOrErr = Sym.getName();
+      if (!SymNameOrErr)
+        consumeError(SymNameOrErr.takeError());
+      else if (TrapOnFailFunctions.count(*SymNameOrErr) > 0)
+        TrapOnFailFunctionAddresses.insert(Addr.second);
+    }
+  }
+  return Error::success();
+}
+
 UnsupportedDisassembly::UnsupportedDisassembly(StringRef Text) : Text(Text) {}
 
 char UnsupportedDisassembly::ID;
diff --git a/tools/llvm-cfi-verify/lib/FileAnalysis.h b/tools/llvm-cfi-verify/lib/FileAnalysis.h
index 3f0a7048788..67421e80da3 100644
--- a/tools/llvm-cfi-verify/lib/FileAnalysis.h
+++ b/tools/llvm-cfi-verify/lib/FileAnalysis.h
@@ -11,6 +11,7 @@
 #define LLVM_CFI_VERIFY_FILE_ANALYSIS_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/DebugInfo/Symbolize/Symbolize.h"
 #include "llvm/MC/MCAsmInfo.h"
@@ -110,6 +111,10 @@ public:
   // Returns whether this instruction is used by CFI to trap the program.
   bool isCFITrap(const Instr &InstrMeta) const;
 
+  // Returns whether this instruction is a call to a function that will trap on
+  // CFI violations (i.e., it serves as a trap in this instance).
+  bool willTrapOnCFIViolation(const Instr &InstrMeta) const;
+
   // Returns whether this function can fall through to the next instruction.
   // Undefined (and bad) instructions cannot fall through, and instruction that
   // modify the control flow can only fall through if they are conditional
@@ -183,6 +188,10 @@ protected:
   // internal members. Should only be called once by Create().
   Error parseCodeSections();
 
+  // Parses the symbol table to look for the addresses of functions that will
+  // trap on CFI violations.
+  Error parseSymbolTable();
+
 private:
   // Members that describe the input file.
   object::OwningBinary<object::Binary> Binary;
@@ -218,6 +227,9 @@ private:
 
   // A list of addresses of indirect control flow instructions.
   std::set<uint64_t> IndirectInstructions;
+
+  // The addresses of functions that will trap on CFI violations.
+  SmallSet<uint64_t, 4> TrapOnFailFunctionAddresses;
 };
 
 class UnsupportedDisassembly : public ErrorInfo<UnsupportedDisassembly> {
diff --git a/tools/llvm-cfi-verify/lib/GraphBuilder.cpp b/tools/llvm-cfi-verify/lib/GraphBuilder.cpp
index 4153b5f6844..5b2bc6f0c3b 100644
--- a/tools/llvm-cfi-verify/lib/GraphBuilder.cpp
+++ b/tools/llvm-cfi-verify/lib/GraphBuilder.cpp
@@ -311,6 +311,24 @@ void GraphBuilder::buildFlowGraphImpl(const FileAnalysis &Analysis,
     Result.ConditionalBranchNodes.push_back(BranchNode);
   }
 
+  // When using cross-DSO, some indirect calls are not guarded by a branch to a
+  // trap but instead follow a call to __cfi_slowpath.  For example:
+  // if (!InlinedFastCheck(f))
+  //    call *f
+  //  else {
+  //    __cfi_slowpath(CallSiteTypeId, f);
+  //    call *f
+  //  }
+  // To mark the second call as protected, we recognize indirect calls that
+  // directly follow calls to functions that will trap on CFI violations.
+  if (CFCrossRefs.empty()) {
+    const Instr *PrevInstr = Analysis.getPrevInstructionSequential(ChildMeta);
+    if (PrevInstr && Analysis.willTrapOnCFIViolation(*PrevInstr)) {
+      Result.IntermediateNodes[PrevInstr->VMAddress] = Address;
+      HasValidCrossRef = true;
+    }
+  }
+
   if (!HasValidCrossRef)
     Result.OrphanedNodes.push_back(Address);