diff --git a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h --- a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h +++ b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h @@ -72,6 +72,38 @@ Instruction *getUntagLocationIfFunctionExit(Instruction &Inst); +struct AllocaInfo { + AllocaInst *AI; + SmallVector LifetimeStart; + SmallVector LifetimeEnd; +}; + +struct StackInfo { + MapVector AllocasToInstrument; + SmallVector UnrecognizedLifetimes; + DenseMap> AllocaDbgMap; + SmallVector RetVec; + bool CallsReturnTwice = false; +}; + +class StackInfoBuilderDelegate { +public: + virtual bool isInterestingAlloca(const AllocaInst &) = 0; + virtual ~StackInfoBuilderDelegate() = default; +}; + +class StackInfoBuilder { +public: + StackInfoBuilder(StackInfoBuilderDelegate *Delegate) : Delegate(Delegate) {} + + void visit(Instruction &Inst); + StackInfo &get() { return Info; }; + +private: + StackInfo Info; + StackInfoBuilderDelegate *Delegate; +}; + } // namespace llvm #endif diff --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp --- a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp +++ b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp @@ -245,14 +245,7 @@ /// An instrumentation pass implementing detection of addressability bugs /// using tagged pointers. -class HWAddressSanitizer { -private: - struct AllocaInfo { - AllocaInst *AI; - SmallVector LifetimeStart; - SmallVector LifetimeEnd; - }; - +class HWAddressSanitizer : StackInfoBuilderDelegate { public: HWAddressSanitizer(Module &M, bool CompileKernel, bool Recover, const StackSafetyGlobalInfo *SSI) @@ -299,18 +292,14 @@ void getInterestingMemoryOperands( Instruction *I, SmallVectorImpl &Interesting); - bool isInterestingAlloca(const AllocaInst &AI); + bool isInterestingAlloca(const AllocaInst &AI) override; void tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag, size_t Size); Value *tagPointer(IRBuilder<> &IRB, Type *Ty, Value *PtrLong, Value *Tag); Value *untagPointer(IRBuilder<> &IRB, Value *PtrLong); - bool instrumentStack( - bool ShouldDetectUseAfterScope, - MapVector &AllocasToInstrument, - SmallVector &UnrecognizedLifetimes, - DenseMap> &AllocaDbgMap, - SmallVectorImpl &RetVec, Value *StackTag, - llvm::function_ref GetDT, - llvm::function_ref GetPDT); + bool instrumentStack(bool ShouldDetectUseAfterScope, StackInfo &Info, + Value *StackTag, + llvm::function_ref GetDT, + llvm::function_ref GetPDT); Value *readRegister(IRBuilder<> &IRB, StringRef Name); bool instrumentLandingPads(SmallVectorImpl &RetVec); Value *getNextTagWithCall(IRBuilder<> &IRB); @@ -1331,11 +1320,7 @@ } bool HWAddressSanitizer::instrumentStack( - bool ShouldDetectUseAfterScope, - MapVector &AllocasToInstrument, - SmallVector &UnrecognizedLifetimes, - DenseMap> &AllocaDbgMap, - SmallVectorImpl &RetVec, Value *StackTag, + bool ShouldDetectUseAfterScope, StackInfo &SInfo, Value *StackTag, llvm::function_ref GetDT, llvm::function_ref GetPDT) { // Ideally, we want to calculate tagged stack base pointer, and rewrite all @@ -1345,7 +1330,7 @@ // This generates one extra instruction per alloca use. unsigned int I = 0; - for (auto &KV : AllocasToInstrument) { + for (auto &KV : SInfo.AllocasToInstrument) { auto N = I++; auto *AI = KV.first; AllocaInfo &Info = KV.second; @@ -1362,7 +1347,7 @@ AI->replaceUsesWithIf(Replacement, [AILong](Use &U) { return U.getUser() != AILong; }); - for (auto *DDI : AllocaDbgMap.lookup(AI)) { + for (auto *DDI : SInfo.AllocaDbgMap.lookup(AI)) { // Prepend "tag_offset, N" to the dwarf expression. // Tag offset logically applies to the alloca pointer, and it makes sense // to put it at the beginning of the expression. @@ -1382,7 +1367,7 @@ tagAlloca(IRB, AI, UARTag, AlignedSize); }; bool StandardLifetime = - UnrecognizedLifetimes.empty() && + SInfo.UnrecognizedLifetimes.empty() && isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, &GetDT(), ClMaxLifetimes); if (ShouldDetectUseAfterScope && StandardLifetime) { @@ -1390,13 +1375,13 @@ IRB.SetInsertPoint(Start->getNextNode()); tagAlloca(IRB, AI, Tag, Size); if (!forAllReachableExits(GetDT(), GetPDT(), Start, Info.LifetimeEnd, - RetVec, TagEnd)) { + SInfo.RetVec, TagEnd)) { for (auto *End : Info.LifetimeEnd) End->eraseFromParent(); } } else { tagAlloca(IRB, AI, Tag, Size); - for (auto *RI : RetVec) + for (auto *RI : SInfo.RetVec) TagEnd(RI); if (!StandardLifetime) { for (auto &II : Info.LifetimeStart) @@ -1406,7 +1391,7 @@ } } } - for (auto &I : UnrecognizedLifetimes) + for (auto &I : SInfo.UnrecognizedLifetimes) I->eraseFromParent(); return true; } @@ -1475,52 +1460,12 @@ SmallVector OperandsToInstrument; SmallVector IntrinToInstrument; - MapVector AllocasToInstrument; - SmallVector RetVec; SmallVector LandingPadVec; - SmallVector UnrecognizedLifetimes; - DenseMap> AllocaDbgMap; - bool CallsReturnTwice = false; + + StackInfoBuilder SIB(this); for (auto &Inst : instructions(F)) { - if (CallInst *CI = dyn_cast(&Inst)) { - if (CI->canReturnTwice()) { - CallsReturnTwice = true; - } - } if (InstrumentStack) { - if (AllocaInst *AI = dyn_cast(&Inst)) { - if (isInterestingAlloca(*AI)) - AllocasToInstrument.insert({AI, {}}); - continue; - } - auto *II = dyn_cast(&Inst); - if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start || - II->getIntrinsicID() == Intrinsic::lifetime_end)) { - AllocaInst *AI = findAllocaForValue(II->getArgOperand(1)); - if (!AI) { - UnrecognizedLifetimes.push_back(&Inst); - continue; - } - if (!isInterestingAlloca(*AI)) - continue; - if (II->getIntrinsicID() == Intrinsic::lifetime_start) - AllocasToInstrument[AI].LifetimeStart.push_back(II); - else - AllocasToInstrument[AI].LifetimeEnd.push_back(II); - continue; - } - } - - Instruction *ExitUntag = getUntagLocationIfFunctionExit(Inst); - if (ExitUntag) - RetVec.push_back(ExitUntag); - - if (auto *DVI = dyn_cast(&Inst)) { - for (Value *V : DVI->location_ops()) { - if (auto *Alloca = dyn_cast_or_null(V)) - if (!AllocaDbgMap.count(Alloca) || AllocaDbgMap[Alloca].back() != DVI) - AllocaDbgMap[Alloca].push_back(DVI); - } + SIB.visit(Inst); } if (InstrumentLandingPads && isa(Inst)) @@ -1533,6 +1478,8 @@ IntrinToInstrument.push_back(MI); } + StackInfo &SInfo = SIB.get(); + initializeCallbacks(*F.getParent()); bool Changed = false; @@ -1540,7 +1487,7 @@ if (!LandingPadVec.empty()) Changed |= instrumentLandingPads(LandingPadVec); - if (AllocasToInstrument.empty() && F.hasPersonalityFn() && + if (SInfo.AllocasToInstrument.empty() && F.hasPersonalityFn() && F.getPersonalityFn()->getName() == kHwasanPersonalityThunkName) { // __hwasan_personality_thunk is a no-op for functions without an // instrumented stack, so we can drop it. @@ -1548,7 +1495,7 @@ Changed = true; } - if (AllocasToInstrument.empty() && OperandsToInstrument.empty() && + if (SInfo.AllocasToInstrument.empty() && OperandsToInstrument.empty() && IntrinToInstrument.empty()) return Changed; @@ -1558,24 +1505,24 @@ IRBuilder<> EntryIRB(InsertPt); emitPrologue(EntryIRB, /*WithFrameRecord*/ ClRecordStackHistory && - Mapping.WithFrameRecord && !AllocasToInstrument.empty()); + Mapping.WithFrameRecord && + !SInfo.AllocasToInstrument.empty()); - if (!AllocasToInstrument.empty()) { + if (!SInfo.AllocasToInstrument.empty()) { Value *StackTag = ClGenerateTagsWithCalls ? nullptr : getStackBaseTag(EntryIRB); // Calls to functions that may return twice (e.g. setjmp) confuse the // postdominator analysis, and will leave us to keep memory tagged after // function return. Work around this by always untagging at every return // statement if return_twice functions are called. - instrumentStack(DetectUseAfterScope && !CallsReturnTwice, - AllocasToInstrument, UnrecognizedLifetimes, AllocaDbgMap, - RetVec, StackTag, GetDT, GetPDT); + instrumentStack(DetectUseAfterScope && !SInfo.CallsReturnTwice, SIB.get(), + StackTag, GetDT, GetPDT); } // Pad and align each of the allocas that we instrumented to stop small // uninteresting allocas from hiding in instrumented alloca's padding and so // that we have enough space to store real tags for short granules. DenseMap AllocaToPaddedAllocaMap = - padInterestingAllocas(AllocasToInstrument); + padInterestingAllocas(SInfo.AllocasToInstrument); if (!AllocaToPaddedAllocaMap.empty()) { for (auto &Inst : instructions(F)) { diff --git a/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp b/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp --- a/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp +++ b/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp @@ -12,6 +12,8 @@ #include "llvm/Transforms/Utils/MemoryTaggingSupport.h" +#include "llvm/Analysis/ValueTracking.h" + namespace llvm { namespace { bool maybeReachableFromEachOther(const SmallVectorImpl &Insts, @@ -54,4 +56,44 @@ } return nullptr; } + +void StackInfoBuilder::visit(Instruction &Inst) { + if (CallInst *CI = dyn_cast(&Inst)) { + if (CI->canReturnTwice()) { + Info.CallsReturnTwice = true; + } + } + if (AllocaInst *AI = dyn_cast(&Inst)) { + if (Delegate->isInterestingAlloca(*AI)) + Info.AllocasToInstrument.insert({AI, {}}); + return; + } + auto *II = dyn_cast(&Inst); + if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start || + II->getIntrinsicID() == Intrinsic::lifetime_end)) { + AllocaInst *AI = findAllocaForValue(II->getArgOperand(1)); + if (!AI) { + Info.UnrecognizedLifetimes.push_back(&Inst); + return; + } + if (!Delegate->isInterestingAlloca(*AI)) + return; + if (II->getIntrinsicID() == Intrinsic::lifetime_start) + Info.AllocasToInstrument[AI].LifetimeStart.push_back(II); + else + Info.AllocasToInstrument[AI].LifetimeEnd.push_back(II); + return; + } + if (auto *DVI = dyn_cast(&Inst)) { + for (Value *V : DVI->location_ops()) { + if (auto *Alloca = dyn_cast_or_null(V)) + if (!Info.AllocaDbgMap.count(Alloca) || + Info.AllocaDbgMap[Alloca].back() != DVI) + Info.AllocaDbgMap[Alloca].push_back(DVI); + } + } + Instruction *ExitUntag = getUntagLocationIfFunctionExit(Inst); + if (ExitUntag) + Info.RetVec.push_back(ExitUntag); +} } // namespace llvm