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 @@ -21,6 +21,7 @@ #include "llvm/IR/Module.h" namespace llvm { +namespace memtag { // For an alloca valid between lifetime markers Start and Ends, call the // Callback for all possible exits out of the lifetime in the containing // function, which can return from the instructions in RetVec. @@ -71,6 +72,34 @@ 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 StackInfoBuilder { +public: + StackInfoBuilder(std::function IsInterestingAlloca) + : IsInterestingAlloca(IsInterestingAlloca) {} + + void visit(Instruction &Inst); + StackInfo &get() { return Info; }; + +private: + StackInfo Info; + std::function IsInterestingAlloca; +}; + +} // namespace memtag } // namespace llvm #endif diff --git a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp --- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp +++ b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp @@ -583,7 +583,7 @@ continue; } - Instruction *ExitUntag = getUntagLocationIfFunctionExit(I); + Instruction *ExitUntag = memtag::getUntagLocationIfFunctionExit(I); if (ExitUntag) RetVec.push_back(ExitUntag); } @@ -643,8 +643,8 @@ bool StandardLifetime = UnrecognizedLifetimes.empty() && - isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, DT, - ClMaxLifetimes); + memtag::isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, DT, + ClMaxLifetimes); // 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 @@ -659,8 +659,8 @@ auto TagEnd = [&](Instruction *Node) { untagAlloca(AI, Node, Size); }; if (!DT || !PDT || - !forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd, RetVec, - TagEnd)) { + !memtag::forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd, + RetVec, TagEnd)) { for (auto *End : Info.LifetimeEnd) End->eraseFromParent(); } 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 @@ -247,13 +247,6 @@ /// An instrumentation pass implementing detection of addressability bugs /// using tagged pointers. class HWAddressSanitizer { -private: - struct AllocaInfo { - AllocaInst *AI; - SmallVector LifetimeStart; - SmallVector LifetimeEnd; - }; - public: HWAddressSanitizer(Module &M, bool CompileKernel, bool Recover, const StackSafetyGlobalInfo *SSI) @@ -269,7 +262,7 @@ void setSSI(const StackSafetyGlobalInfo *S) { SSI = S; } DenseMap padInterestingAllocas( - const MapVector &AllocasToInstrument); + const MapVector &AllocasToInstrument); bool sanitizeFunction(Function &F, llvm::function_ref GetDT, llvm::function_ref GetPDT); @@ -304,14 +297,10 @@ 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, memtag::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); @@ -1325,11 +1314,7 @@ } bool HWAddressSanitizer::instrumentStack( - bool ShouldDetectUseAfterScope, - MapVector &AllocasToInstrument, - SmallVector &UnrecognizedLifetimes, - DenseMap> &AllocaDbgMap, - SmallVectorImpl &RetVec, Value *StackTag, + bool ShouldDetectUseAfterScope, memtag::StackInfo &SInfo, Value *StackTag, llvm::function_ref GetDT, llvm::function_ref GetPDT) { // Ideally, we want to calculate tagged stack base pointer, and rewrite all @@ -1339,10 +1324,10 @@ // 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; + memtag::AllocaInfo &Info = KV.second; IRBuilder<> IRB(AI->getNextNode()); // Replace uses of the alloca with tagged address. @@ -1356,7 +1341,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. @@ -1376,21 +1361,22 @@ tagAlloca(IRB, AI, UARTag, AlignedSize); }; bool StandardLifetime = - UnrecognizedLifetimes.empty() && - isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, &GetDT(), - ClMaxLifetimes); + SInfo.UnrecognizedLifetimes.empty() && + memtag::isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, + &GetDT(), ClMaxLifetimes); if (ShouldDetectUseAfterScope && StandardLifetime) { IntrinsicInst *Start = Info.LifetimeStart[0]; IRB.SetInsertPoint(Start->getNextNode()); tagAlloca(IRB, AI, Tag, Size); - if (!forAllReachableExits(GetDT(), GetPDT(), Start, Info.LifetimeEnd, - RetVec, TagEnd)) { + if (!memtag::forAllReachableExits(GetDT(), GetPDT(), Start, + Info.LifetimeEnd, 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) @@ -1400,7 +1386,7 @@ } } } - for (auto &I : UnrecognizedLifetimes) + for (auto &I : SInfo.UnrecognizedLifetimes) I->eraseFromParent(); return true; } @@ -1424,7 +1410,7 @@ } DenseMap HWAddressSanitizer::padInterestingAllocas( - const MapVector &AllocasToInstrument) { + const MapVector &AllocasToInstrument) { DenseMap AllocaToPaddedAllocaMap; for (auto &KV : AllocasToInstrument) { AllocaInst *AI = KV.first; @@ -1469,52 +1455,13 @@ SmallVector OperandsToInstrument; SmallVector IntrinToInstrument; - MapVector AllocasToInstrument; - SmallVector RetVec; SmallVector LandingPadVec; - SmallVector UnrecognizedLifetimes; - DenseMap> AllocaDbgMap; - bool CallsReturnTwice = false; + + memtag::StackInfoBuilder SIB( + [this](const AllocaInst &AI) { return isInterestingAlloca(AI); }); 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)) @@ -1527,6 +1474,8 @@ IntrinToInstrument.push_back(MI); } + memtag::StackInfo &SInfo = SIB.get(); + initializeCallbacks(*F.getParent()); bool Changed = false; @@ -1534,7 +1483,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. @@ -1542,7 +1491,7 @@ Changed = true; } - if (AllocasToInstrument.empty() && OperandsToInstrument.empty() && + if (SInfo.AllocasToInstrument.empty() && OperandsToInstrument.empty() && IntrinToInstrument.empty()) return Changed; @@ -1552,24 +1501,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,7 +12,10 @@ #include "llvm/Transforms/Utils/MemoryTaggingSupport.h" +#include "llvm/Analysis/ValueTracking.h" + namespace llvm { +namespace memtag { namespace { bool maybeReachableFromEachOther(const SmallVectorImpl &Insts, const DominatorTree *DT, size_t MaxLifetimes) { @@ -54,4 +57,46 @@ } 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 (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 (!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 memtag } // namespace llvm