Index: lib/Transforms/Instrumentation/SafeStack.cpp =================================================================== --- lib/Transforms/Instrumentation/SafeStack.cpp +++ lib/Transforms/Instrumentation/SafeStack.cpp @@ -118,6 +118,10 @@ SmallVectorImpl &Returns, SmallVectorImpl &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. @@ -177,6 +181,17 @@ 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(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)); @@ -187,8 +202,7 @@ 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" @@ -463,10 +477,8 @@ for (AllocaInst *AI : StaticAllocas) { IRB.SetInsertPoint(AI); - auto CArraySize = cast(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. Index: test/Transforms/SafeStack/array.ll =================================================================== --- test/Transforms/SafeStack/array.ll +++ test/Transforms/SafeStack/array.ll @@ -35,4 +35,52 @@ 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*)