Vulnerability.h

/*
 * Vulnerability.h
 *
 * Vulnerability is the base class of different
 * vulnerabilities.
 *
 * Each Vulnerability provides the same interfaces
 * whose implementations are customized for one kind of
 * vulnerability, e.g. NPD.
 *
 * The instance of the class should be passed to
 * PPPathSensitiveChecker as an input, so that the checker
 * can use it to detect specific bugs.
 *
 * Author: Qingkai
 */
 
#ifndef CHECKER_PSA_VULNERABILITY_H
#define CHECKER_PSA_VULNERABILITY_H
 
#include <llvm/Pass.h>
 
#include "IR/SEG/SEGCallSite.h"
#include "IR/SEG/SEGReturnSite.h"
#include "IR/SEG/SymbolicExprGraph.h"
#include "IR/SEG/SymbolicExprGraphSolver.h"
#include "Report/BugReport/BugReportMgr.h"
#include "VulnerabilityTrace.h"
#include "VulnerabilityTraceBuilder.h"
 
#include "Checker/PSA/PSAReportDecoratorDefault.h"
#include "Checker/PSA/PSATaintReportDecorator.h"
#include "IR/SEG/SEGCallSiteOutputNode.h"
#include "IR/SEG/SEGSimpleSite.h"
 
using namespace llvm;
 
class SailFishVulnerability;
class SrcMustNotReachSinkVulnerability;
class SrcMustReachSinkVulnerability;
class Vulnerability {
public:
  enum SiteType {
    ST_Return = 1 << 0,
    ST_Call = 1 << 1,
    ST_Sink = 1 << 2,
    ST_Others = 1 << 3,
  };
 
  enum VulnerabilityCategoryType {
    VCT_Begin,
    VCT_SinkMustReach,
    VCT_SinkMustNotReach,
    VCT_Taint,
    VCT_End
  };
 
  typedef typename std::pair<const SEGOperandNode *, const SEGSiteBase *>
      ValueSitePairType;
 
private:
  VulnerabilityCategoryType CategoryType;
 
  bool Parasitical = true;
 
  char *VulnName;
  const BugDescription::BugImportance Importance;
  const BugDescription::BugClassification Classification;
  const char *VulnDescription;
 
  // The constructor is made private and, thus it cannot be extended.
  // Users must extend its sub-classes to implement different types of
  // vulnerabilities.
  Vulnerability(VulnerabilityCategoryType CTy, const char *Name,
                const BugDescription::BugImportance Importance,
                const BugDescription::BugClassification Classification,
                const char *Desc)
      : CategoryType(CTy), Importance(Importance),
        Classification(Classification), VulnDescription(Desc) {
    VulnName = (char *)malloc(300);
    std::strcpy(VulnName, (char *)Name);
  }
 
  friend class SinkMustReachVulnerability;
  friend class SinkMustNotReachVulnerability;
  friend class TaintStyleVulnerability;
  friend class SrcMustNotReachSinkVulnerability;
  friend class SrcMustReachSinkVulnerability;
  friend class SailFishVulnerability;
 
public:
  virtual ~Vulnerability() {}
 
  virtual void setSources(const SymbolicExprGraph *SEG,
                          std::vector<ValueSitePairType> &Sources) = 0;
 
  virtual void setPrerequisites(SymbolicExprGraphSolver *Solver,
                                const SEGSiteBase *CurrSite,
                                const VulnerabilityTraceBuilder &TraceHistory,
                                SMTExprVec &Prerequisites) = 0;
 
  virtual bool checkNode(const SEGNodeBase *CurrNode,
                         const VulnerabilityTraceBuilder &TraceHistory) {
    return false;
  }
 
  virtual bool finalCheck(const VulnerabilityTraceBuilder &TraceHistory) {
    return true;
  }
 
  virtual void
  ConstantCheck(const SymbolicExprGraph *SEG,
                std::list<shared_ptr<VulnerabilityTrace>> &AllTraces) {
    return;
  }
 
  virtual bool isFlowInsensitive() { return false; }
 
  virtual SiteType checkSite(const SEGSiteBase *CurrSite,
                             const VulnerabilityTraceBuilder &TraceHistory) = 0;
 
  virtual bool checkTrace(std::shared_ptr<VulnerabilityTrace> &Trace) {
    return true;
  }
 
  virtual void getAnalysisUsage(AnalysisUsage &AU) {}
 
  virtual void initializeAnalysis(Pass *P) {}
 
public:
  VulnerabilityCategoryType getCategoryType() const { return CategoryType; }
 
  bool isParasitical() const { return Parasitical; }
 
  void setParasitical(bool B) { Parasitical = B; }
 
  const char *getDescription() { return VulnDescription; }
 
  BugDescription::BugImportance getImportance() const { return Importance; }
 
  BugDescription::BugClassification getClassification() const {
    return Classification;
  }
 
  const char *getName() { return VulnName; }
 
  virtual PSAReportDecoratorDefault *allocNewDecorator() {
    return new PSAReportDecoratorDefault;
  }
 
  virtual void destroyDecorator(PSAReportDecoratorDefault *decorator) {
    delete decorator;
  }
};
 
class SinkMustReachVulnerability : public Vulnerability {
public:
  SinkMustReachVulnerability(
      const char *Name, const BugDescription::BugImportance importance,
      const BugDescription::BugClassification classification, const char *Desc)
      : Vulnerability(VCT_SinkMustReach, Name, importance, classification,
                      Desc) {}
};
 
class SinkMustNotReachVulnerability : public Vulnerability {
public:
  SinkMustNotReachVulnerability(
      const char *Name, const BugDescription::BugImportance Importance,
      const BugDescription::BugClassification Classification, const char *Desc)
      : Vulnerability(VCT_SinkMustNotReach, Name, Importance, Classification,
                      Desc) {}
};
 
class SailFishVulnerability : public Vulnerability {
public:
  SailFishVulnerability(VulnerabilityCategoryType CTy, const char *Name,
                        const BugDescription::BugImportance Importance,
                        const BugDescription::BugClassification Classification,
                        const char *Desc)
      : Vulnerability(CTy, Name, Importance, Classification, Desc) {}
 
  virtual void setSources(const SymbolicExprGraph *SEG,
                          std::vector<ValueSitePairType> &Sources) {
    for (auto It = SEG->value_node_begin(), E = SEG->value_node_end(); It != E;
         It++) {
      SEGOperandNode *Node = It->second;
 
      if (Node->use_site_size()) {
        for (auto SiteIt = Node->use_site_begin(), SiteE = Node->use_site_end();
             SiteIt != SiteE; SiteIt++) {
          if (isSource(Node, *SiteIt)) {
            //            outs() << "[Find Source]: " << *Node << " in func: "
            //                   <<
            //                   Node->getParentGraph()->getBaseFunc()->getName()
            //                   << "\n";
            Sources.emplace_back(Node, *SiteIt);
          }
        }
      } else if (isSource(Node, nullptr)) {
        //        outs() << "[Find Source]: " << *Node << " in func: "
        //               << Node->getParentGraph()->getBaseFunc()->getName() <<
        //               "\n";
        Sources.emplace_back(Node, nullptr);
      }
    }
 
    // Add modeling common output value node
    for (auto it = SEG->seg_callsite_begin(), E = SEG->seg_callsite_end();
         it != E; it++) {
      if (*it == nullptr) {
        continue;
      }
 
      if (!(*it)->hasCommonOutput()) {
        continue;
      }
 
      auto CO = (*it)->getCommonOutput();
      if (auto *CSO = const_cast<SEGCallSiteCommonOutputNode *>(CO)) {
        if (isSource(CSO, *it)) {
          //          outs() << "[Find Source]: " << *CSO << " in func: "
          //                 << CSO->getParentGraph()->getBaseFunc()->getName()
          //                 << "\n";
          Sources.emplace_back(CSO, *it);
        }
      }
    }
  }
 
  virtual Vulnerability::SiteType
  checkSite(const SEGSiteBase *CurrSite,
            const VulnerabilityTraceBuilder &TraceHistory) override {
    auto *Node = TraceHistory.recentObjAs<SEGOperandNode>();
 
    if (isSink((SEGNodeBase *)Node, const_cast<SEGSiteBase *>(CurrSite))) {
      //      outs() << "[Find Sink]: " << *Node
      //             << " in func: " <<
      //             Node->getParentGraph()->getBaseFunc()->getName()
      //             << "\n";
      if (CurrSite) {
        //        outs() << "[Find Sink Site]: " << *CurrSite << "\n";
      }
      return ST_Sink;
    } else if (isa<SEGReturnSite>(CurrSite)) {
      return ST_Return;
    } else if (isa<SEGCallSite>(CurrSite)) {
      return ST_Call;
    }
    return ST_Others;
  }
 
  virtual bool isSource(SEGNodeBase *Node, SEGSiteBase *Site) { return false; };
 
  virtual bool isSink(SEGNodeBase *Node, SEGSiteBase *Site) { return false; };
 
  virtual void setPrerequisites(SymbolicExprGraphSolver *Solver,
                                const SEGSiteBase *CurrSite,
                                const VulnerabilityTraceBuilder &TraceHistory,
                                SMTExprVec &Prerequisites) override {
    //    Prerequisites.push_back(
    //        Solver->getOrInsertExpr(TraceHistory.recentObjAs<SEGNodeBase>())
    //        == 0);
  }
 
  virtual void transfer(const SEGSiteBase *Site, const SEGNodeBase *Arg,
                        std::vector<const SEGNodeBase *> &TransferDsts) {}
 
public:
  static bool classof(const Vulnerability *vuln) {
    return vuln->getCategoryType() > VCT_Begin &&
           vuln->getCategoryType() < VCT_End;
  }
};
 
class SrcMustReachSinkVulnerability : public SailFishVulnerability {
public:
  SrcMustReachSinkVulnerability(
      const char *Name,
      BugDescription::BugImportance Importance = BugDescription::BI_MEDIUM,
      BugDescription::BugClassification Classification =
          BugDescription::BC_SECURITY,
      const char *Desc = "CUSTOMIZED")
      : SailFishVulnerability(VCT_SinkMustReach, Name, Importance,
                              Classification, Desc) {}
};
 
class SrcMustNotReachSinkVulnerability : public SailFishVulnerability {
public:
  SrcMustNotReachSinkVulnerability(
      const char *Name,
      BugDescription::BugImportance Importance = BugDescription::BI_MEDIUM,
      BugDescription::BugClassification Classification =
          BugDescription::BC_SECURITY,
      const char *Desc = "CUSTOMIZED")
      : SailFishVulnerability(VCT_SinkMustNotReach, Name, Importance,
                              Classification, Desc) {}
};
 
class TaintExtras {
private:
  TaintExtras() {}
  friend class TaintStyleVulnerability;
 
public:
  virtual ~TaintExtras() {}
  virtual bool isSinkFunction(const SymbolicExprGraph *G,
                              const SEGArgumentNode *Arg) const = 0;
  virtual void transfer(const SEGSiteBase *Site, const SEGNodeBase *Arg,
                        std::vector<const SEGNodeBase *> &TransferDsts) {}
};
 
class TaintStyleVulnerability : public Vulnerability, public TaintExtras {
public:
  TaintStyleVulnerability(
      const char *Name, const BugDescription::BugImportance importance,
      const BugDescription::BugClassification classification, const char *Desc)
      : Vulnerability(VCT_Taint, Name, importance, classification, Desc) {}
 
  virtual PSAReportDecoratorDefault *allocNewDecorator() override {
    return new PSATaintReportDecorator;
  }
};
 
class MultiVulnerability {
protected:
  std::vector<std::shared_ptr<Vulnerability>> Vulns;
 
  std::map<std::shared_ptr<Vulnerability>,
           std::set<std::shared_ptr<Vulnerability>>>
      HeadDeps;
 
  std::map<std::shared_ptr<Vulnerability>,
           std::set<std::shared_ptr<Vulnerability>>>
      TailDeps;
 
  typedef std::pair<std::shared_ptr<Vulnerability>,
                    std::shared_ptr<Vulnerability>>
      VulnPairTy;
  typedef std::function<const SEGObject *(std::shared_ptr<VulnerabilityTrace>)>
      GetTraceFunctionTy;
  std::map<VulnPairTy, GetTraceFunctionTy> HeadDepFctorMap;
  std::map<VulnPairTy, GetTraceFunctionTy> TailDepFctorMap;
 
public:
  virtual ~MultiVulnerability() {}
 
  bool add(std::shared_ptr<Vulnerability> V) {
    Vulns.push_back(V);
    return true;
  }
 
  std::shared_ptr<Vulnerability> get(size_t Index) const {
    return Vulns[Index];
  }
 
  size_t size() const { return Vulns.size(); }
 
  std::vector<std::shared_ptr<Vulnerability>>::const_iterator begin() const {
    return Vulns.begin();
  }
 
  std::vector<std::shared_ptr<Vulnerability>>::const_iterator end() const {
    return Vulns.end();
  }
 
  virtual const char *getFullName() = 0;
 
  virtual const char *getAbbrName() = 0;
 
  virtual void buildDependence() = 0;
 
  void insertHeadDependence(std::shared_ptr<Vulnerability> V1,
                            std::shared_ptr<Vulnerability> V2,
                            GetTraceFunctionTy &Func) {
    HeadDeps[V1].insert(V2);
    HeadDepFctorMap[std::make_pair(V1, V2)] = Func;
  }
 
  void insertTailDependence(std::shared_ptr<Vulnerability> V1,
                            std::shared_ptr<Vulnerability> V2,
                            GetTraceFunctionTy &Func) {
    TailDeps[V1].insert(V2);
    TailDepFctorMap[std::make_pair(V1, V2)] = Func;
  }
 
  std::set<std::shared_ptr<Vulnerability>>::iterator
  head_dep_begin(std::shared_ptr<Vulnerability> V) {
    return HeadDeps[V].begin();
  }
 
  std::set<std::shared_ptr<Vulnerability>>::iterator
  head_dep_end(std::shared_ptr<Vulnerability> V) {
    return HeadDeps[V].end();
  }
 
  std::set<std::shared_ptr<Vulnerability>>::iterator
  tail_dep_begin(std::shared_ptr<Vulnerability> V) {
    return TailDeps[V].begin();
  }
 
  std::set<std::shared_ptr<Vulnerability>>::iterator
  tail_dep_end(std::shared_ptr<Vulnerability> V) {
    return TailDeps[V].end();
  }
 
  const SEGObject *getTraceHead(std::shared_ptr<Vulnerability> V1,
                                std::shared_ptr<VulnerabilityTrace> T1,
                                std::shared_ptr<Vulnerability> V2) {
    auto It = HeadDepFctorMap.find(std::make_pair(V1, V2));
    if (It != HeadDepFctorMap.end()) {
      return It->second(T1);
    }
    return nullptr;
  }
 
  const SEGObject *getTraceTail(std::shared_ptr<Vulnerability> V1,
                                std::shared_ptr<VulnerabilityTrace> T1,
                                std::shared_ptr<Vulnerability> V2) {
    auto It = TailDepFctorMap.find(std::make_pair(V1, V2));
    if (It != TailDepFctorMap.end()) {
      return It->second(T1);
    }
    return nullptr;
  }
 
};
 
#endif /* CHECKER_PSA_VULNERABILITY_H */