[StackProtector] Catch direct out-of-bounds when checking address-tak…

…enness

With -fstack-protector-strong we check if a non-array variable has its address
taken in a way that could cause a potential out-of-bounds access. However what
we don't catch is when the address is directly used to create an out-of-bounds
memory access.

Fix this by examining the offsets of GEPs that are ultimately derived from
allocas and checking if the resulting address is out-of-bounds, and by checking
that any memory operations using such addresses are not over-large.

Fixes PR43478.

Differential revision: https://reviews.llvm.org/D75695

llvm/include/llvm/CodeGen/StackProtector.h

@@ -95,7 +95,7 @@ class StackProtector : public FunctionPass {
                                 bool InStruct = false) const;
 
   /// Check whether a stack allocation has its address taken.
-  bool HasAddressTaken(const Instruction *AI);
+  bool HasAddressTaken(const Instruction *AI, uint64_t AllocSize);
 
   /// RequiresStackProtector - Check whether or not this function needs a
   /// stack protector based upon the stack protector level.

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,19 +197,34 @@ 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: {
       // Keep track of what PHI nodes we have already visited to ensure
       // 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();
       }
     }
   }