2 files changed, 39 insertions, 14 deletions

diff --git a/llvm/tools/llvm-cfi-verify/lib/FileAnalysis.cpp b/llvm/tools/llvm-cfi-verify/lib/FileAnalysis.cpp
index 75482544718..15bf20f091d 100644
--- a/llvm/tools/llvm-cfi-verify/lib/FileAnalysis.cpp
+++ b/llvm/tools/llvm-cfi-verify/lib/FileAnalysis.cpp
@@ -154,7 +154,8 @@ const Instr &FileAnalysis::getInstructionOrDie(uint64_t Address) const {
 }
 
 bool FileAnalysis::isCFITrap(const Instr &InstrMeta) const {
-  return MII->getName(InstrMeta.Instruction.getOpcode()) == "TRAP";
+  const auto &InstrDesc = MII->get(InstrMeta.Instruction.getOpcode());
+  return InstrDesc.isTrap();
 }
 
 bool FileAnalysis::canFallThrough(const Instr &InstrMeta) const {
@@ -296,20 +297,38 @@ uint64_t FileAnalysis::indirectCFOperandClobber(const GraphResult &Graph) const
     else
       Node = Branch.Fallthrough;
 
-    while (Node != Graph.BaseAddress) {
+    // Some architectures (e.g., AArch64) cannot load in an indirect branch, so
+    // we allow them one load.
+    bool canLoad = !MII->get(IndirectCF.Instruction.getOpcode()).mayLoad();
+
+    // We walk backwards from the indirect CF.  It is the last node returned by
+    // Graph.flattenAddress, so we skip it since we already handled it.
+    DenseSet<unsigned> CurRegisterNumbers = RegisterNumbers;
+    std::vector<uint64_t> Nodes = Graph.flattenAddress(Node);
+    for (auto I = Nodes.rbegin() + 1, E = Nodes.rend(); I != E; ++I) {
+      Node = *I;
       const Instr &NodeInstr = getInstructionOrDie(Node);
       const auto &InstrDesc = MII->get(NodeInstr.Instruction.getOpcode());
 
-      for (unsigned RegNum : RegisterNumbers) {
+      for (auto RI = CurRegisterNumbers.begin(), RE = CurRegisterNumbers.end();
+           RI != RE; ++RI) {
+        unsigned RegNum = *RI;
         if (InstrDesc.hasDefOfPhysReg(NodeInstr.Instruction, RegNum,
-                                      *RegisterInfo))
-          return Node;
+                                      *RegisterInfo)) {
+          if (!canLoad || !InstrDesc.mayLoad())
+            return Node;
+          canLoad = false;
+          CurRegisterNumbers.erase(RI);
+          // Add the registers this load reads to those we check for clobbers.
+          for (unsigned i = InstrDesc.getNumDefs(),
+                        e = InstrDesc.getNumOperands(); i != e; i++) {
+            const auto Operand = NodeInstr.Instruction.getOperand(i);
+            if (Operand.isReg())
+              CurRegisterNumbers.insert(Operand.getReg());
+          }
+          break;
+        }
       }
-
-      const auto &KV = Graph.IntermediateNodes.find(Node);
-      assert((KV != Graph.IntermediateNodes.end()) &&
-             "Could not get next node.");
-      Node = KV->second;
     }
   }
 
@@ -456,6 +475,9 @@ void FileAnalysis::parseSectionContents(ArrayRef<uint8_t> SectionBytes,
     if (!usesRegisterOperand(InstrMeta))
       continue;
 
+    if (InstrDesc.isReturn())
+      continue;
+
     // Check if this instruction exists in the range of the DWARF metadata.
     if (!IgnoreDWARFFlag) {
       auto LineInfo =
diff --git a/llvm/tools/llvm-cfi-verify/lib/FileAnalysis.h b/llvm/tools/llvm-cfi-verify/lib/FileAnalysis.h
index ce81f8bfbe3..3f0a7048788 100644
--- a/llvm/tools/llvm-cfi-verify/lib/FileAnalysis.h
+++ b/llvm/tools/llvm-cfi-verify/lib/FileAnalysis.h
@@ -149,10 +149,13 @@ public:
   CFIProtectionStatus validateCFIProtection(const GraphResult &Graph) const;
 
   // Returns the first place the operand register is clobbered between the CFI-
-  // check and the indirect CF instruction execution. If the register is not
-  // modified, returns the address of the indirect CF instruction. The result is
-  // undefined if the provided graph does not fall under either the
-  // FAIL_REGISTER_CLOBBERED or PROTECTED status (see CFIProtectionStatus).
+  // check and the indirect CF instruction execution. We do this by walking
+  // backwards from the indirect CF and ensuring there is at most one load
+  // involving the operand register (which is the indirect CF itself on x86).
+  // If the register is not modified, returns the address of the indirect CF
+  // instruction. The result is undefined if the provided graph does not fall
+  // under either the FAIL_REGISTER_CLOBBERED or PROTECTED status (see
+  // CFIProtectionStatus).
   uint64_t indirectCFOperandClobber(const GraphResult& Graph) const;
 
   // Prints an instruction to the provided stream using this object's pretty-