Index: lib/Transforms/Instrumentation/MemorySanitizer.cpp =================================================================== --- lib/Transforms/Instrumentation/MemorySanitizer.cpp +++ lib/Transforms/Instrumentation/MemorySanitizer.cpp @@ -92,6 +92,7 @@ //===----------------------------------------------------------------------===// #include "llvm/Transforms/Instrumentation.h" +#include "llvm/IR/DIBuilder.h" #include "llvm/ADT/DepthFirstIterator.h" #include "llvm/ADT/SmallString.h" #include "llvm/ADT/SmallVector.h" @@ -191,6 +192,10 @@ cl::desc("Insert checks for constant shadow values"), cl::Hidden, cl::init(false)); +static cl::opt ClOneBigAlloca("msan-one-big-alloca", + cl::desc("Lump all local variables into one alloca"), + cl::Hidden, cl::init(false)); + namespace { // Memory map parameters used in application-to-shadow address calculation. @@ -336,6 +341,7 @@ InlineAsm *EmptyAsm; friend struct MemorySanitizerVisitor; + friend struct FunctionStackPoisoner; friend struct VarArgAMD64Helper; }; } // namespace @@ -512,6 +518,115 @@ namespace { +struct MemorySanitizerVisitor; + +struct FunctionStackPoisoner : public InstVisitor { + Function &F; + MemorySanitizer &MS; + MemorySanitizerVisitor &MSV; + bool PoisonStack; + DIBuilder DIB; + LLVMContext *C; + Type *IntptrTy; + Type *IntptrPtrTy; + + SmallVector AllocaVec; + SmallVector DynamicAllocaVec; + + bool HasNonEmptyInlineAsm; + std::unique_ptr EmptyInlineAsm; + + FunctionStackPoisoner(Function &F, MemorySanitizer &MS, + MemorySanitizerVisitor &MSV, bool PoisonStack) + : F(F), + MS(MS), + MSV(MSV), + PoisonStack(PoisonStack), + DIB(*F.getParent(), /*AllowUnresolved*/ false), + C(MS.C), + IntptrTy(MS.IntptrTy), + IntptrPtrTy(PointerType::get(IntptrTy, 0)), + HasNonEmptyInlineAsm(false), + EmptyInlineAsm(CallInst::Create(MS.EmptyAsm)) {} + + bool runOnFunction() { + // Collect alloca instructions. + for (BasicBlock *BB : depth_first(&F.getEntryBlock())) visit(*BB); + + if (AllocaVec.empty() && DynamicAllocaVec.empty()) return false; + + poisonStack(); + return true; + } + + // Poison all found Alloca instructions. + void poisonStack(); + + // ----------------------- Visitors. + + // Poison dynamic alloca. + void handleDynamicAllocaCall(AllocaInst &I); + + struct MSanStackVariableDescription { + StringRef Name; // Name of the variable that will be displayed by MSan + // if a stack-related bug is reported. + uint64_t Size; // Size of the variable in bytes. + size_t Alignment; // Alignment of the variable (power of 2). + AllocaInst *AI; // The actual AllocaInst. + size_t Offset; // Offset from the beginning of the frame; + // set by ComputeMSanStackFrameLayout. + Value *NewAllocaPtr; // New address of the variable (points inside the big + // alloca). + }; + + struct MSanStackFrameLayout { + size_t FrameAlignment; // Alignment for the entire frame. + size_t FrameSize; // Size of the frame in bytes. + }; + + void ComputeStackFrameLayout( + SmallVectorImpl &Vars, + MSanStackFrameLayout *Layout); + + /// \brief Collect Alloca instructions we want (and can) handle. + void visitAllocaInst(AllocaInst &AI) { + if (!isInterestingAlloca(AI)) return; + + if (isDynamicAlloca(AI)) + DynamicAllocaVec.push_back(&AI); + else + AllocaVec.push_back(&AI); + } + + void visitCallInst(CallInst &CI) { + HasNonEmptyInlineAsm |= + CI.isInlineAsm() && !CI.isIdenticalTo(EmptyInlineAsm.get()); + } + + // ---------------------- Helpers. + void initializeCallbacks(Module &M); + + bool isDynamicAlloca(AllocaInst &AI) const { + return AI.isArrayAllocation() || !AI.isStaticAlloca(); + } + + // Check if we want (and can) handle this alloca. + bool isInterestingAlloca(AllocaInst &AI) const { + return (AI.getAllocatedType()->isSized() && + // alloca() may be called with 0 size, ignore it. + getAllocaSizeInBytes(&AI) > 0); + } + + uint64_t getAllocaSizeInBytes(AllocaInst *AI) const { + Type *Ty = AI->getAllocatedType(); + uint64_t SizeInBytes = MS.DL->getTypeAllocSize(Ty); + return SizeInBytes; + } + + Value *createAllocaForLayout(IRBuilder<> &IRB, const MSanStackFrameLayout &L, + bool Dynamic); +}; + /// \brief A helper class that handles instrumentation of VarArg /// functions on a particular platform. /// @@ -785,7 +900,6 @@ for (BasicBlock *BB : depth_first(&F.getEntryBlock())) visit(*BB); - // Finalize PHI nodes. for (PHINode *PN : ShadowPHINodes) { PHINode *PNS = cast(getShadow(PN)); @@ -810,6 +924,11 @@ // Insert shadow value checks. materializeChecks(InstrumentWithCalls); + if (ClOneBigAlloca) { + FunctionStackPoisoner FSP(F, MS, *this, PoisonStack); + FSP.runOnFunction(); + } + return true; } @@ -2540,39 +2659,45 @@ "_msphi_o")); } + void SetAllocaOrigin(IRBuilder<> &IRB, StringRef Name, Value *Ptr, + Value *Size) { + SmallString<2048> StackDescriptionStorage; + raw_svector_ostream StackDescription(StackDescriptionStorage); + // We create a string with a description of the stack allocation and + // pass it into __msan_set_alloca_origin. + // It will be printed by the run-time if stack-originated UMR is found. + // The first 4 bytes of the string are set to '----' and will be replaced + // by __msan_va_arg_overflow_size_tls at the first call. + StackDescription << "----" << Name << "@" << F.getName(); + Value *Descr = createPrivateNonConstGlobalForString(*F.getParent(), + StackDescription.str()); + + IRB.CreateCall4(MS.MsanSetAllocaOrigin4Fn, + IRB.CreatePointerCast(Ptr, IRB.getInt8PtrTy()), Size, + IRB.CreatePointerCast(Descr, IRB.getInt8PtrTy()), + IRB.CreatePointerCast(&F, MS.IntptrTy)); + } + void visitAllocaInst(AllocaInst &I) { setShadow(&I, getCleanShadow(&I)); setOrigin(&I, getCleanOrigin()); - IRBuilder<> IRB(I.getNextNode()); - uint64_t Size = MS.DL->getTypeAllocSize(I.getAllocatedType()); - if (PoisonStack && ClPoisonStackWithCall) { - IRB.CreateCall2(MS.MsanPoisonStackFn, - IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), - ConstantInt::get(MS.IntptrTy, Size)); - } else { - Value *ShadowBase = getShadowPtr(&I, Type::getInt8PtrTy(*MS.C), IRB); - Value *PoisonValue = IRB.getInt8(PoisonStack ? ClPoisonStackPattern : 0); - IRB.CreateMemSet(ShadowBase, PoisonValue, Size, I.getAlignment()); - } - - if (PoisonStack && MS.TrackOrigins) { - SmallString<2048> StackDescriptionStorage; - raw_svector_ostream StackDescription(StackDescriptionStorage); - // We create a string with a description of the stack allocation and - // pass it into __msan_set_alloca_origin. - // It will be printed by the run-time if stack-originated UMR is found. - // The first 4 bytes of the string are set to '----' and will be replaced - // by __msan_va_arg_overflow_size_tls at the first call. - StackDescription << "----" << I.getName() << "@" << F.getName(); - Value *Descr = - createPrivateNonConstGlobalForString(*F.getParent(), - StackDescription.str()); - - IRB.CreateCall4(MS.MsanSetAllocaOrigin4Fn, - IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), - ConstantInt::get(MS.IntptrTy, Size), - IRB.CreatePointerCast(Descr, IRB.getInt8PtrTy()), - IRB.CreatePointerCast(&F, MS.IntptrTy)); + if (!ClOneBigAlloca) { + IRBuilder<> IRB(I.getNextNode()); + uint64_t Size = MS.DL->getTypeAllocSize(I.getAllocatedType()); + if (PoisonStack && ClPoisonStackWithCall) { + IRB.CreateCall2(MS.MsanPoisonStackFn, + IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), + ConstantInt::get(MS.IntptrTy, Size)); + } else { + Value *ShadowBase = getShadowPtr(&I, Type::getInt8PtrTy(*MS.C), IRB); + Value *PoisonValue = + IRB.getInt8(PoisonStack ? ClPoisonStackPattern : 0); + IRB.CreateMemSet(ShadowBase, PoisonValue, Size, I.getAlignment()); + } + + if (PoisonStack && MS.TrackOrigins) + SetAllocaOrigin(IRB, I.getName(), &I, + ConstantInt::get(MS.IntptrTy, Size)); } } @@ -2883,6 +3008,142 @@ return new VarArgNoOpHelper(Func, Msan, Visitor); } +static DebugLoc getFunctionEntryDebugLocation(Function &F) { + for (const auto &Inst : F.getEntryBlock()) + if (!isa(Inst)) return Inst.getDebugLoc(); + return DebugLoc(); +} + +Value *FunctionStackPoisoner::createAllocaForLayout( + IRBuilder<> &IRB, const MSanStackFrameLayout &L, bool Dynamic) { + AllocaInst *Alloca; + if (Dynamic) { + Alloca = IRB.CreateAlloca(IRB.getInt8Ty(), + ConstantInt::get(IRB.getInt64Ty(), L.FrameSize), + "MSanAlloca"); + } else { + Alloca = IRB.CreateAlloca(ArrayType::get(IRB.getInt8Ty(), L.FrameSize), + nullptr, "MSanAlloca"); + assert(Alloca->isStaticAlloca()); + } + Alloca->setAlignment(L.FrameAlignment); + return IRB.CreatePointerCast(Alloca, IntptrTy); +} + +void FunctionStackPoisoner::handleDynamicAllocaCall(AllocaInst &I) { + IRBuilder<> IRB(I.getNextNode()); + uint64_t Size = MS.DL->getTypeAllocSize(I.getAllocatedType()); + if (PoisonStack && ClPoisonStackWithCall) { + IRB.CreateCall2(MS.MsanPoisonStackFn, + IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), + ConstantInt::get(MS.IntptrTy, Size)); + } else { + Value *ShadowBase = MSV.getShadowPtr(&I, Type::getInt8PtrTy(*MS.C), IRB); + Value *PoisonValue = IRB.getInt8(PoisonStack ? ClPoisonStackPattern : 0); + IRB.CreateMemSet(ShadowBase, PoisonValue, Size, I.getAlignment()); + } + + if (PoisonStack && MS.TrackOrigins) + MSV.SetAllocaOrigin(IRB, I.getName(), &I, + ConstantInt::get(MS.IntptrTy, Size)); +} + +void FunctionStackPoisoner::ComputeStackFrameLayout( + SmallVectorImpl &Vars, + MSanStackFrameLayout *Layout) { + size_t NumVars = Vars.size(); + assert(NumVars > 0); + + std::stable_sort(Vars.begin(), Vars.end(), + [](const MSanStackVariableDescription &a, + const MSanStackVariableDescription &b) { + return a.Alignment > b.Alignment; + }); + Layout->FrameAlignment = Vars[0].Alignment; + size_t Offset = 0; + for (size_t i = 0; i < NumVars; i++) { + size_t Alignment = Vars[i].Alignment; + size_t Size = Vars[i].Size; + assert((Alignment & (Alignment - 1)) == 0); + assert(Layout->FrameAlignment >= Alignment); + assert(Size > 0); + + Offset = RoundUpToAlignment(Offset, Alignment); + Vars[i].Offset = Offset; + Offset += Size; + } + Layout->FrameSize = Offset; +} + +void FunctionStackPoisoner::poisonStack() { + assert(AllocaVec.size() > 0 || DynamicAllocaVec.size() > 0); + + // Handle dynamic allocas. + for (auto &AllocaCall : DynamicAllocaVec) + handleDynamicAllocaCall(*AllocaCall); + + if (AllocaVec.size() == 0) return; + + DebugLoc EntryDebugLocation = getFunctionEntryDebugLocation(F); + + Instruction *InsBefore = AllocaVec[0]; + IRBuilder<> IRB(InsBefore); + IRB.SetCurrentDebugLocation(EntryDebugLocation); + + SmallVector SVD; + SVD.reserve(AllocaVec.size()); + for (AllocaInst *AI : AllocaVec) { + unsigned Alignment = AI->getAlignment(); + if (Alignment == 0) { + Type *Ty = AI->getType(); // FIXME??? + Alignment = std::max(Alignment, MS.DL->getABITypeAlignment(Ty)); + } + // We want each local to get its own origin id. + // Pad everything to 4 bytes. + Alignment = std::max(Alignment, kMinOriginAlignment); + + MSanStackVariableDescription D = {AI->getName(), getAllocaSizeInBytes(AI), + Alignment, AI, 0, 0}; + SVD.push_back(D); + } + MSanStackFrameLayout L; + ComputeStackFrameLayout(SVD, &L); + // Don't do dynamic alloca in presence of inline asm: too often it + // makes assumptions on which registers are available. + bool DoDynamicAlloca = !HasNonEmptyInlineAsm; + + Value *LocalStackBase = createAllocaForLayout(IRB, L, DoDynamicAlloca); + + // Replace Alloca instructions with base+offset. + for (auto &Desc : SVD) { + AllocaInst *AI = Desc.AI; + Desc.NewAllocaPtr = IRB.CreateIntToPtr( + IRB.CreateAdd(LocalStackBase, ConstantInt::get(IntptrTy, Desc.Offset)), + AI->getType()); + replaceDbgDeclareForAlloca(AI, Desc.NewAllocaPtr, DIB, /*Deref=*/true); + AI->replaceAllUsesWith(Desc.NewAllocaPtr); + } + + // Poison the stack frame at the entry. + if (ClPoisonStackWithCall) { + IRB.CreateCall2(MS.MsanPoisonStackFn, LocalStackBase, + ConstantInt::get(MS.IntptrTy, L.FrameSize)); + } else { + Value *ShadowBase = + MSV.getShadowPtr(LocalStackBase, Type::getInt8PtrTy(*MS.C), IRB); + Value *PoisonValue = IRB.getInt8(PoisonStack ? ClPoisonStackPattern : 0); + IRB.CreateMemSet(ShadowBase, PoisonValue, L.FrameSize, L.FrameAlignment); + } + + if (PoisonStack && MS.TrackOrigins) + for (const auto &Desc : SVD) + MSV.SetAllocaOrigin(IRB, Desc.Name, Desc.NewAllocaPtr, + ConstantInt::get(MS.IntptrTy, Desc.Size)); + + // We are done. Remove the old unused alloca instructions. + for (auto AI : AllocaVec) AI->eraseFromParent(); +} + } // namespace bool MemorySanitizer::runOnFunction(Function &F) { @@ -2898,3 +3159,4 @@ return Visitor.runOnFunction(); } + Index: test/Instrumentation/MemorySanitizer/alloca.ll =================================================================== --- test/Instrumentation/MemorySanitizer/alloca.ll +++ test/Instrumentation/MemorySanitizer/alloca.ll @@ -0,0 +1,23 @@ +; RUN: opt < %s -msan -msan-check-access-address=0 -msan-one-big-alloca=1 -S | FileCheck %s +; RUN: opt < %s -msan -msan-check-access-address=0 -msan-one-big-alloca=1 -msan-track-origins=1 -S | FileCheck -check-prefix=CHECK -check-prefix=CHECK-ORIGINS %s + +target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128" +target triple = "x86_64-unknown-linux-gnu" + +; Test that stack allocations are unpoisoned in one memset call + +define i32 @ManyAlloca() sanitize_memory { +entry: + %p = alloca i8, align 1 + %q = alloca i32, align 8 + %r = alloca i64, align 8 + %x = call i32 @ManyAllocaHelper(i8* %p, i32* %q, i64* %r) + ret i32 %x +} + +declare i32 @ManyAllocaHelper(i8* %p, i32* %q, i64* %r) + +; CHECK: @ManyAlloca +; CHECK: call void @llvm.memset.p0i8.i64(i8* {{.*}}, i8 -1, i64 17, i32 8, i1 false) +; CHECK: call i32 @ManyAllocaHelper( +; CHECK: ret i32