1 files changed, 32 insertions, 5 deletions

diff --git a/llvm/lib/CodeGen/StackProtector.cpp b/llvm/lib/CodeGen/StackProtector.cpp
index 4e2189884bb1..a343791807e6 100644
--- a/llvm/lib/CodeGen/StackProtector.cpp
+++ b/llvm/lib/CodeGen/StackProtector.cpp
@@ -18,6 +18,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/EHPersonalities.h"
+#include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetLowering.h"
@@ -161,9 +162,16 @@ bool StackProtector::ContainsProtectableArray(Type *Ty, bool &IsLarge,
   return NeedsProtector;
 }
 
-bool StackProtector::HasAddressTaken(const Instruction *AI) {
+bool StackProtector::HasAddressTaken(const Instruction *AI,
+                                     uint64_t AllocSize) {
+  const DataLayout &DL = M->getDataLayout();
   for (const User *U : AI->users()) {
     const auto *I = cast<Instruction>(U);
+    // If this instruction accesses memory make sure it doesn't access beyond
+    // the bounds of the allocated object.
+    Optional<MemoryLocation> MemLoc = MemoryLocation::getOrNone(I);
+    if (MemLoc.hasValue() && MemLoc->Size.getValue() > AllocSize)
+      return true;
     switch (I->getOpcode()) {
     case Instruction::Store:
       if (AI == cast<StoreInst>(I)->getValueOperand())
@@ -189,11 +197,26 @@ bool StackProtector::HasAddressTaken(const Instruction *AI) {
     }
     case Instruction::Invoke:
       return true;
+    case Instruction::GetElementPtr: {
+      // If the GEP offset is out-of-bounds, or is non-constant and so has to be
+      // assumed to be potentially out-of-bounds, then any memory access that
+      // would use it could also be out-of-bounds meaning stack protection is
+      // required.
+      const GetElementPtrInst *GEP = cast<GetElementPtrInst>(I);
+      unsigned TypeSize = DL.getIndexTypeSizeInBits(I->getType());
+      APInt Offset(TypeSize, 0);
+      APInt MaxOffset(TypeSize, AllocSize);
+      if (!GEP->accumulateConstantOffset(DL, Offset) || Offset.ugt(MaxOffset))
+        return true;
+      // Adjust AllocSize to be the space remaining after this offset.
+      if (HasAddressTaken(I, AllocSize - Offset.getLimitedValue()))
+        return true;
+      break;
+    }
     case Instruction::BitCast:
-    case Instruction::GetElementPtr:
     case Instruction::Select:
     case Instruction::AddrSpaceCast:
-      if (HasAddressTaken(I))
+      if (HasAddressTaken(I, AllocSize))
         return true;
       break;
     case Instruction::PHI: {
@@ -201,7 +224,7 @@ bool StackProtector::HasAddressTaken(const Instruction *AI) {
       // they are only visited once.
       const auto *PN = cast<PHINode>(I);
       if (VisitedPHIs.insert(PN).second)
-        if (HasAddressTaken(PN))
+        if (HasAddressTaken(PN, AllocSize))
           return true;
       break;
     }
@@ -330,7 +353,8 @@ bool StackProtector::RequiresStackProtector() {
           continue;
         }
 
-        if (Strong && HasAddressTaken(AI)) {
+        if (Strong && HasAddressTaken(AI, M->getDataLayout().getTypeAllocSize(
+                                              AI->getAllocatedType()))) {
           ++NumAddrTaken;
           Layout.insert(std::make_pair(AI, MachineFrameInfo::SSPLK_AddrOf));
           ORE.emit([&]() {
@@ -342,6 +366,9 @@ bool StackProtector::RequiresStackProtector() {
           });
           NeedsProtector = true;
         }
+        // Clear any PHIs that we visited, to make sure we examine all uses of
+        // any subsequent allocas that we look at.
+        VisitedPHIs.clear();
       }
     }
   }