[safestack] Fix handling of array allocas.

[eugenis]

Summary

The current code does not take alloca array size into account and,
as a result, considers any access past the first array element to be
unsafe.

[pcc]

lib/Transforms/Instrumentation/SafeStack.cpp

187

AI can be a dynamic alloca at this point, I believe (see findInsts -> IsSafeStackAlloca -> IsAccessSafe -> getStaticAllocaAllocationSize), so this may not be a constant.

[eugenis]

lib/Transforms/Instrumentation/SafeStack.cpp

187

Right. Returning 0 for unknown size allocas.

[eugenis]

lib/Transforms/Instrumentation/SafeStack.cpp

187

This would check such allocas for safety as if they are size 0, i.e. some operations are still allowed, but not any loads or stores.
We could actually do better by passing down the size as a Value and applying the SCEV magic to it. But (a) one thing at a time and (b) I doubt it actually matters.

[pcc]

Lgtm

[eugenis]

r254350

lib/Transforms/Instrumentation/SafeStack.cpp

class SafeStack : public FunctionPass {
   /// \brief Find all static allocas, dynamic allocas, return instructions and
   /// stack restore points (exception unwind blocks and setjmp calls) in the
   /// given function and append them to the respective vectors.
   void findInsts(Function &F, SmallVectorImpl<AllocaInst *> &StaticAllocas,
                  SmallVectorImpl<AllocaInst *> &DynamicAllocas,
                  SmallVectorImpl<ReturnInst *> &Returns,
                  SmallVectorImpl<Instruction *> &StackRestorePoints);
 
+  /// \brief Calculate the allocation size of a given alloca. Returns 0 if the
+  /// size can not be statically determined.
+  uint64_t getStaticAllocaAllocationSize(const AllocaInst* AI);
+
   /// \brief Allocate space for all static allocas in \p StaticAllocas,
   /// replace allocas with pointers into the unsafe stack and generate code to
   /// restore the stack pointer before all return instructions in \p Returns.
   ///
   /// \returns A pointer to the top of the unsafe stack after all unsafe static
   /// allocas are allocated.
   Value *moveStaticAllocasToUnsafeStack(IRBuilder<> &IRB, Function &F,
                                         ArrayRef<AllocaInst *> StaticAllocas,
bool doInitialization(Module &M) override {
     Int8Ty = Type::getInt8Ty(M.getContext());
 
     return false;
   }
 
   bool runOnFunction(Function &F) override;
 }; // class SafeStack
 
+uint64_t SafeStack::getStaticAllocaAllocationSize(const AllocaInst* AI) {
+  uint64_t Size = DL->getTypeAllocSize(AI->getAllocatedType());
+  if (AI->isArrayAllocation()) {
+    auto C = dyn_cast<ConstantInt>(AI->getArraySize());
+    if (!C)
+      return 0;
+    Size *= C->getZExtValue();
+  }
+  return Size;
+}
+
 bool SafeStack::IsAccessSafe(Value *Addr, uint64_t Size, const AllocaInst *AI) {
   AllocaOffsetRewriter Rewriter(*SE, AI);
   const SCEV *Expr = Rewriter.visit(SE->getSCEV(Addr));
 
   uint64_t BitWidth = SE->getTypeSizeInBits(Expr->getType());
   ConstantRange AccessStartRange = SE->getUnsignedRange(Expr);
   ConstantRange SizeRange =
       ConstantRange(APInt(BitWidth, 0), APInt(BitWidth, Size));
   ConstantRange AccessRange = AccessStartRange.add(SizeRange);
   ConstantRange AllocaRange = ConstantRange(
-      APInt(BitWidth, 0),
-      APInt(BitWidth, DL->getTypeStoreSize(AI->getAllocatedType())));
+      APInt(BitWidth, 0), APInt(BitWidth, getStaticAllocaAllocationSize(AI)));
   bool Safe = AllocaRange.contains(AccessRange);
 
   DEBUG(dbgs() << "[SafeStack] Alloca " << *AI << "\n"
                << "            Access " << *Addr << "\n"
                << "            SCEV " << *Expr
                << " U: " << SE->getUnsignedRange(Expr)
                << ", S: " << SE->getSignedRange(Expr) << "\n"
                << "            Range " << AccessRange << "\n"
BasePointer = cast<Instruction>(IRB.CreateIntToPtr(
         StackPtrTy));
   }
 
   // Allocate space for every unsafe static AllocaInst on the unsafe stack.
   int64_t StaticOffset = 0; // Current stack top.
   for (AllocaInst *AI : StaticAllocas) {
     IRB.SetInsertPoint(AI);
 
-    auto CArraySize = cast<ConstantInt>(AI->getArraySize());
     Type *Ty = AI->getAllocatedType();
-
-    uint64_t Size = DL->getTypeAllocSize(Ty) * CArraySize->getZExtValue();
+    uint64_t Size = getStaticAllocaAllocationSize(AI);
     if (Size == 0)
       Size = 1; // Don't create zero-sized stack objects.
 
     // Ensure the object is properly aligned.
     unsigned Align =
         std::max((unsigned)DL->getPrefTypeAlignment(Ty), AI->getAlignment());
 
     // Add alignment.

test/Transforms/SafeStack/array.ll

entry:
 
   ; CHECK: call i8* @strcpy(i8* %[[GEP]], i8* %[[A2]])
   %call = call i8* @strcpy(i8* %gep, i8* %a2)
 
   ; CHECK: store i8* %[[USP]], i8** @__safestack_unsafe_stack_ptr
   ret void
 }
 
+; Load from an array at a fixed offset, no overflow.
+define i8 @StaticArrayFixedSafe() nounwind uwtable safestack {
+entry:
+  ; CHECK-LABEL: define i8 @StaticArrayFixedSafe(
+  ; CHECK-NOT: __safestack_unsafe_stack_ptr
+  ; CHECK: ret i8
+  %buf = alloca i8, i32 4, align 1
+  %gep = getelementptr inbounds i8, i8* %buf, i32 2
+  %x = load i8, i8* %gep, align 1
+  ret i8 %x
+}
+
+; Load from an array at a fixed offset with overflow.
+define i8 @StaticArrayFixedUnsafe() nounwind uwtable safestack {
+entry:
+  ; CHECK-LABEL: define i8 @StaticArrayFixedUnsafe(
+  ; CHECK: __safestack_unsafe_stack_ptr
+  ; CHECK: ret i8
+  %buf = alloca i8, i32 4, align 1
+  %gep = getelementptr inbounds i8, i8* %buf, i32 5
+  %x = load i8, i8* %gep, align 1
+  ret i8 %x
+}
+
+; Load from an array at an unknown offset.
+define i8 @StaticArrayVariableUnsafe(i32 %ofs) nounwind uwtable safestack {
+entry:
+  ; CHECK-LABEL: define i8 @StaticArrayVariableUnsafe(
+  ; CHECK: __safestack_unsafe_stack_ptr
+  ; CHECK: ret i8
+  %buf = alloca i8, i32 4, align 1
+  %gep = getelementptr inbounds i8, i8* %buf, i32 %ofs
+  %x = load i8, i8* %gep, align 1
+  ret i8 %x
+}
+
+; Load from an array at an unknown offset.
+define i8 @DynamicArrayUnsafe(i32 %sz) nounwind uwtable safestack {
+entry:
+  ; CHECK-LABEL: define i8 @DynamicArrayUnsafe(
+  ; CHECK: __safestack_unsafe_stack_ptr
+  ; CHECK: ret i8
+  %buf = alloca i8, i32 %sz, align 1
+  %gep = getelementptr inbounds i8, i8* %buf, i32 2
+  %x = load i8, i8* %gep, align 1
+  ret i8 %x
+}
+
 declare i8* @strcpy(i8*, i8*)