1 files changed, 12 insertions, 0 deletions
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> {