r253694 - Thread Safety Analysis: Fix DenseMap iterator invalidation UAF

[Reid Kleckner via cfe-commits]

Author: rnk
Date: Fri Nov 20 13:08:30 2015
New Revision: 253694

URL: http://llvm.org/viewvc/llvm-project?rev=253694&view=rev
Log:
Thread Safety Analysis: Fix DenseMap iterator invalidation UAF

Rather than storing BeforeInfo in the DenseMap by value, this stores a
unique_ptr to it, so that we can keep a pointer to it live across
subsequent DenseMap insertions.

This change also removes the unique_ptr wrapper around BeforeVect
because now we're indirecting at a higher level.

Modified:
    cfe/trunk/lib/Analysis/ThreadSafety.cpp

Modified: cfe/trunk/lib/Analysis/ThreadSafety.cpp
URL: http://llvm.org/viewvc/llvm-project/cfe/trunk/lib/Analysis/ThreadSafety.cpp?rev=253694&r1=253693&r2=253694&view=diff
==============================================================================
--- cfe/trunk/lib/Analysis/ThreadSafety.cpp (original)
+++ cfe/trunk/lib/Analysis/ThreadSafety.cpp Fri Nov 20 13:08:30 2015
@@ -258,16 +258,15 @@ private:
   typedef SmallVector<const ValueDecl*, 4>  BeforeVect;
 
   struct BeforeInfo {
-    BeforeInfo() : Vect(nullptr), Visited(false) { }
-    BeforeInfo(BeforeInfo &&O)
-        : Vect(std::move(O.Vect)), Visited(O.Visited)
-    {}
+    BeforeInfo() : Visited(0) {}
+    BeforeInfo(BeforeInfo &&O) : Vect(std::move(O.Vect)), Visited(O.Visited) {}
 
-    std::unique_ptr<BeforeVect> Vect;
-    int                         Visited;
+    BeforeVect Vect;
+    int Visited;
   };
 
-  typedef llvm::DenseMap<const ValueDecl*, BeforeInfo>  BeforeMap;
+  typedef llvm::DenseMap<const ValueDecl *, std::unique_ptr<BeforeInfo>>
+      BeforeMap;
   typedef llvm::DenseMap<const ValueDecl*, bool>        CycleMap;
 
 public:
@@ -276,6 +275,9 @@ public:
   BeforeInfo* insertAttrExprs(const ValueDecl* Vd,
                               ThreadSafetyAnalyzer& Analyzer);
 
+  BeforeInfo *getBeforeInfoForDecl(const ValueDecl *Vd,
+                                   ThreadSafetyAnalyzer &Analyzer);
+
   void checkBeforeAfter(const ValueDecl* Vd,
                         const FactSet& FSet,
                         ThreadSafetyAnalyzer& Analyzer,
@@ -965,26 +967,27 @@ public:
 BeforeSet::BeforeInfo* BeforeSet::insertAttrExprs(const ValueDecl* Vd,
     ThreadSafetyAnalyzer& Analyzer) {
   // Create a new entry for Vd.
-  auto& Entry = BMap.FindAndConstruct(Vd);
-  BeforeInfo* Info = &Entry.second;
-  BeforeVect* Bv = nullptr;
+  BeforeInfo *Info = nullptr;
+  {
+    // Keep InfoPtr in its own scope in case BMap is modified later and the
+    // reference becomes invalid.
+    std::unique_ptr<BeforeInfo> &InfoPtr = BMap[Vd];
+    if (!InfoPtr)
+      InfoPtr.reset(new BeforeInfo());
+    Info = InfoPtr.get();
+  }
 
   for (Attr* At : Vd->attrs()) {
     switch (At->getKind()) {
       case attr::AcquiredBefore: {
         auto *A = cast<AcquiredBeforeAttr>(At);
 
-        // Create a new BeforeVect for Vd if necessary.
-        if (!Bv) {
-          Bv = new BeforeVect;
-          Info->Vect.reset(Bv);
-        }
         // Read exprs from the attribute, and add them to BeforeVect.
         for (const auto *Arg : A->args()) {
           CapabilityExpr Cp =
             Analyzer.SxBuilder.translateAttrExpr(Arg, nullptr);
           if (const ValueDecl *Cpvd = Cp.valueDecl()) {
-            Bv->push_back(Cpvd);
+            Info->Vect.push_back(Cpvd);
             auto It = BMap.find(Cpvd);
             if (It == BMap.end())
               insertAttrExprs(Cpvd, Analyzer);
@@ -1001,20 +1004,8 @@ BeforeSet::BeforeInfo* BeforeSet::insert
             Analyzer.SxBuilder.translateAttrExpr(Arg, nullptr);
           if (const ValueDecl *ArgVd = Cp.valueDecl()) {
             // Get entry for mutex listed in attribute
-            BeforeInfo* ArgInfo;
-            auto It = BMap.find(ArgVd);
-            if (It == BMap.end())
-              ArgInfo = insertAttrExprs(ArgVd, Analyzer);
-            else
-              ArgInfo = &It->second;
-
-            // Create a new BeforeVect if necessary.
-            BeforeVect* ArgBv = ArgInfo->Vect.get();
-            if (!ArgBv) {
-              ArgBv = new BeforeVect;
-              ArgInfo->Vect.reset(ArgBv);
-            }
-            ArgBv->push_back(Vd);
+            BeforeInfo *ArgInfo = getBeforeInfoForDecl(ArgVd, Analyzer);
+            ArgInfo->Vect.push_back(Vd);
           }
         }
         break;
@@ -1027,6 +1018,18 @@ BeforeSet::BeforeInfo* BeforeSet::insert
   return Info;
 }
 
+BeforeSet::BeforeInfo *
+BeforeSet::getBeforeInfoForDecl(const ValueDecl *Vd,
+                                ThreadSafetyAnalyzer &Analyzer) {
+  auto It = BMap.find(Vd);
+  BeforeInfo *Info = nullptr;
+  if (It == BMap.end())
+    Info = insertAttrExprs(Vd, Analyzer);
+  else
+    Info = It->second.get();
+  assert(Info && "BMap contained nullptr?");
+  return Info;
+}
 
 /// Return true if any mutexes in FSet are in the acquired_before set of Vd.
 void BeforeSet::checkBeforeAfter(const ValueDecl* StartVd,
@@ -1041,12 +1044,7 @@ void BeforeSet::checkBeforeAfter(const V
     if (!Vd)
       return false;
 
-    BeforeSet::BeforeInfo* Info;
-    auto It = BMap.find(Vd);
-    if (It == BMap.end())
-      Info = insertAttrExprs(Vd, Analyzer);
-    else
-      Info = &It->second;
+    BeforeSet::BeforeInfo *Info = getBeforeInfoForDecl(Vd, Analyzer);
 
     if (Info->Visited == 1)
       return true;
@@ -1054,13 +1052,12 @@ void BeforeSet::checkBeforeAfter(const V
     if (Info->Visited == 2)
       return false;
 
-    BeforeVect* Bv = Info->Vect.get();
-    if (!Bv)
+    if (Info->Vect.empty())
       return false;
 
     InfoVect.push_back(Info);
     Info->Visited = 1;
-    for (auto *Vdb : *Bv) {
+    for (auto *Vdb : Info->Vect) {
       // Exclude mutexes in our immediate before set.
       if (FSet.containsMutexDecl(Analyzer.FactMan, Vdb)) {
         StringRef L1 = StartVd->getName();