Merge pull request cisco-open#1 from niktesic/ntesic-taint-param

[TA] Analyze out-of-BT func with tainted parameters

llvm-15.0.3/llvm-crash-analyzer/include/Analysis/RegisterEquivalence.h

@@ -90,6 +90,7 @@ class RegisterEquivalence {
   bool applyRegisterCopy(MachineInstr &MI);
   bool applyLoad(MachineInstr &MI);
   bool applyStore(MachineInstr &MI);
+  bool applyArithInstr(MachineInstr &MI);
   bool applyCall(MachineInstr &MI);
   bool applyRegDef(MachineInstr &MI);
 

llvm-15.0.3/llvm-crash-analyzer/include/Analysis/TaintAnalysis.h

@@ -82,6 +82,8 @@ class TaintAnalysis {
   ConcreteReverseExec *CRE = nullptr;
   RegisterEquivalence *REA = nullptr;
 
+  // Used for functions out of the backtrace.
+  SmallVector<TaintInfo, 8> TL_Of_Call;
 public:
   TaintAnalysis(StringRef TaintDotFileName, StringRef MirDotFileName,
                 bool PrintPotentialCrashCauseLocation);
@@ -105,6 +107,15 @@ class TaintAnalysis {
   void startTaint(DestSourcePair &DS, SmallVectorImpl<TaintInfo> &TL,
                   const MachineInstr &MI, TaintDataFlowGraph &TaintDFG,
                   RegisterEquivalence &REAnalysis);
+  bool forwardMFAnalysis(BlameModule &BM, const MachineFunction &MF,
+                         TaintDataFlowGraph &TaintDFG,
+                         unsigned levelOfCalledFn = 0,
+                         SmallVector<TaintInfo, 8> *TL_Of_Caller = nullptr,
+                         const MachineInstr *CallMI = nullptr);
+  bool propagateTaintFwd(DestSourcePair &DS, SmallVectorImpl<TaintInfo> &TL,
+                         const MachineInstr &MI, TaintDataFlowGraph &TaintDFG,
+                         RegisterEquivalence &REAnalysis,
+                         const MachineInstr *CallMI = nullptr);
   void addNewTaint(TaintInfo &Ti, SmallVectorImpl<TaintInfo> &TL,
                    const MachineInstr &MI, TaintDataFlowGraph &TaintDFG,
                    std::shared_ptr<Node> crashNode);
@@ -121,6 +132,24 @@ class TaintAnalysis {
   void printTaintList2(SmallVectorImpl<TaintInfo> &TL);
   void printDestSrcInfo(DestSourcePair &DS, const MachineInstr &MI);
   bool shouldAnalyzeCall(SmallVectorImpl<TaintInfo> &TL);
+  bool areParamsTainted(const MachineInstr *CallMI,
+                        SmallVectorImpl<TaintInfo> &TL,
+                        SmallVectorImpl<TaintInfo> *TL_Of_Caller,
+                        TaintDataFlowGraph &TaintDFG,
+                        RegisterEquivalence &REAnalysis);
+  const MachineInstr *findParamLoadingInstr(TaintInfo &Ti,
+                                            const MachineInstr *CallMI);
+  void transformBPtoSPTaints(const MachineFunction &MF,
+                             TaintDataFlowGraph &TaintDFG,
+                             SmallVectorImpl<TaintInfo> &TL);
+  void transformSPtoBPTaints(const MachineFunction &MF,
+                             TaintDataFlowGraph &TaintDFG,
+                             SmallVectorImpl<TaintInfo> &TL);
+  bool isStackSlotTainted(const MachineInstr *CallMI,
+                          SmallVectorImpl<TaintInfo> &TL,
+                          SmallVectorImpl<TaintInfo> *TL_Of_Caller,
+                          TaintDataFlowGraph &TaintDFG,
+                          RegisterEquivalence *REAnalysis);
   TaintInfo isTainted(TaintInfo &Op, SmallVectorImpl<TaintInfo> &TL,
                       RegisterEquivalence *REAnalysis = nullptr,
                       const MachineInstr *MI = nullptr);

llvm-15.0.3/llvm-crash-analyzer/include/Analysis/TaintDataFlowGraph.h

@@ -19,6 +19,7 @@
 
 #include <memory>
 #include <string>
+#include <algorithm>
 
 using namespace llvm;
 using namespace crash_analyzer;
@@ -32,6 +33,7 @@ struct Node {
   static unsigned NextID;
   bool IsCrashNode;
   bool IsContant;
+  unsigned Depth = 0;
 
   // Call instruction that performed the call to function that is out
   // of bt.
@@ -161,13 +163,13 @@ class TaintDataFlowGraph {
 public:
   // Map operand to the latest taint node.
   // FIXME: This should be private.
-  std::map<TaintInfo, std::shared_ptr<Node>> lastTaintedNode;
+  std::map<const MachineOperand*, std::shared_ptr<Node>> lastTaintedNode;
 
   void addEdge(std::shared_ptr<Node> src, std::shared_ptr<Node> dest,
                EdgeType e_type = EdgeType::Assigment);
   void addNode(std::shared_ptr<Node> n);
 
-  void updateLastTaintedNode(TaintInfo Op, std::shared_ptr<Node> N);
+  void updateLastTaintedNode(const MachineOperand* Op, std::shared_ptr<Node> N);
   unsigned getBlameNodesSize() { return Nodes.size(); }
 
   Node *getCrashNode() { return Nodes[0].get(); }

llvm-15.0.3/llvm-crash-analyzer/include/Target/CATargetInfo.h

@@ -18,6 +18,8 @@
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Host.h"
 
@@ -39,7 +41,8 @@ class CATargetInfo {
   std::unordered_map<unsigned, RegAliasTuple> RegMap;
 
   // Save PC value for each instruction.
-  std::unordered_map<const MachineInstr*, std::pair<uint64_t,uint64_t>> InstAddrs;
+  std::unordered_map<const MachineInstr *, std::pair<uint64_t, uint64_t>>
+      InstAddrs;
 
   // Singleton class for the CATargetInfo instance.
   template <typename T> class Singleton {
@@ -53,34 +56,47 @@ class CATargetInfo {
 
 public:
   CATargetInfo() {}
-  virtual ~CATargetInfo() { RegMap.clear(); InstAddrs.clear(); }
+  virtual ~CATargetInfo() {
+    RegMap.clear();
+    InstAddrs.clear();
+  }
 
   // Get register index in the RegMap.
   virtual Optional<unsigned> getID(std::string RegName) const = 0;
 
+  // Get register unsigned (MCRegister) from the RegMap.
+  virtual Optional<unsigned> getRegister(std::string RegName,
+                                         const MachineInstr *MI) const = 0;
+
   virtual unsigned getRegSize(std::string RegName) const = 0;
 
   // Get RegAliasTuple from the RegMap with selected Id.
   RegAliasTuple &getRegMap(unsigned Id) const {
     return const_cast<RegAliasTuple &>(RegMap.at(Id));
   }
 
+  // Get whole RegMap.
+  const std::unordered_map<unsigned, RegAliasTuple> &getWholeRegMap() const {
+    return RegMap;
+  }
+
   // Get InstAddr from the InstAddrs map for the MI.
-  Optional<uint64_t> getInstAddr(const MachineInstr* MI) {
+  Optional<uint64_t> getInstAddr(const MachineInstr *MI) {
     if (InstAddrs.count(MI) == 0)
       return None;
     return InstAddrs[MI].first;
   }
 
   // Get InstAddr from the InstAddrs map for the MI.
-  Optional<uint64_t> getInstSize(const MachineInstr* MI) {
+  Optional<uint64_t> getInstSize(const MachineInstr *MI) {
     if (InstAddrs.count(MI) == 0)
       return None;
     return InstAddrs[MI].second;
   }
 
   // Set InstAddr in the InstAddrs map for the MI.
-  void setInstAddr(const MachineInstr* MI, uint64_t InstAddr, uint64_t InstSize = 0) {
+  void setInstAddr(const MachineInstr *MI, uint64_t InstAddr,
+                   uint64_t InstSize = 0) {
     InstAddrs[MI] = {InstAddr, InstSize};
   }
 
@@ -93,12 +109,24 @@ class CATargetInfo {
   // Return name of the Program Counter Register.
   virtual Optional<std::string> getPC() const = 0;
 
+  // Return name of the Base Pointer Register.
+  virtual Optional<std::string> getBP() const = 0;
+
+  // Return name of the Stack Pointer Register.
+  virtual Optional<std::string> getSP() const = 0;
+
   // Return true if the register is Stack Pointer Register.
   virtual bool isSPRegister(std::string RegName) const = 0;
 
+  // Return SP adjustment to BP in the callee.
+  virtual int64_t getSPAdjust() const = 0;
+
   // Return true if the register is Base Pointer Register.
   virtual bool isBPRegister(std::string RegName) const = 0;
 
+  // Return true if the register can be used to forward parameters.
+  virtual bool isParamFwdRegister(std::string RegName) const = 0;
+
   // Set target Triple of the CATargetInfo instance.
   static void initializeCATargetInfo(Triple *Triple) {
     if (!TT)
@@ -118,6 +146,9 @@ class X86CATargetInfo : public CATargetInfo {
 
   Optional<unsigned> getID(std::string RegName) const override;
 
+  Optional<unsigned> getRegister(std::string RegName,
+                                 const MachineInstr *MI) const override;
+
   unsigned getRegSize(std::string RegName) const override;
 
   bool isRetValRegister(std::string RegName) const override;
@@ -126,10 +157,19 @@ class X86CATargetInfo : public CATargetInfo {
 
   Optional<std::string> getPC() const override;
 
+  Optional<std::string> getBP() const override;
+
+  Optional<std::string> getSP() const override;
+
   bool isSPRegister(std::string RegName) const override;
 
+  // Return SP adjustment to BP in the callee.
+  int64_t getSPAdjust() const override;
+
   bool isBPRegister(std::string RegName) const override;
 
+  bool isParamFwdRegister(std::string RegName) const override;
+
   // Define static instance getter for each target.
   static X86CATargetInfo *instance() {
     return CATargetInfo::Singleton<X86CATargetInfo>::get();

llvm-15.0.3/llvm-crash-analyzer/lib/Analysis/RegisterEquivalence.cpp

@@ -277,6 +277,57 @@ bool RegisterEquivalence::applyStore(MachineInstr &MI) {
   return true;
 }
 
+bool RegisterEquivalence::applyArithInstr(MachineInstr &MI) {
+  // Currently only LEA is concerned, but this should work for ADD, SUB...
+  if (!TII->isLEAInstr(MI))
+    return false;
+
+  auto srcDest = TII->getDestAndSrc(MI);
+  if (!srcDest)
+    return false;
+
+  auto SrcReg = srcDest->Source->getReg();
+  auto DestReg = srcDest->Destination->getReg();
+
+  int64_t SrcOffset = 0;
+  // Currently cannot handle if dest and src use the same reg.
+  if (SrcReg == DestReg)
+    return false;
+
+  // Take the offset into account.
+  if (srcDest->SrcOffset)
+    SrcOffset = *srcDest->SrcOffset;
+
+  // Transform deref->$rip+(off) to deref->$noreg+(rip+off).
+  auto CATI = getCATargetInfoInstance();
+  std::string RegName = TRI->getRegAsmName(SrcReg).lower();
+  if (CATI->isPCRegister(RegName) && CATI->getInstAddr(&MI)) {
+    SrcReg = 0;
+    SrcOffset += *CATI->getInstAddr(&MI) + *CATI->getInstSize(&MI);
+  }
+
+  RegisterOffsetPair Src{SrcReg, SrcOffset};
+  // Same as Load, but Src Memory address is not dereferenced.
+  Src.IsDeref = false;
+  RegisterOffsetPair Dest{DestReg};
+
+  // First invalidate dest reg, since it is being rewritten.
+  invalidateAllRegUses(MI, Dest);
+
+  // If SrcReg is redefined (same as DestReg), set only identity equivalence.
+  if (Src.RegNum == Dest.RegNum) {
+    if (RegInfo[&MI][Dest].find(Src) == RegInfo[&MI][Dest].end())
+      RegInfo[&MI][Src].insert(Src);
+    return true;
+  }
+
+  // Set (transitive) equivalence.
+  setRegEq(MI, Src, Dest);
+  // dumpRegTableAfterMI(&MI);
+
+  return true;
+}
+
 bool RegisterEquivalence::applyCall(MachineInstr &MI) {
   // TODO: Implement this by invalidating registers
   // that will be clobbered by the call.
@@ -309,6 +360,8 @@ void RegisterEquivalence::processMI(MachineInstr &MI) {
     return;
   if (applyStore(MI))
     return;
+  if (applyArithInstr(MI))
+    return;
   if (applyCall(MI))
     return;
   if (applyRegDef(MI))