diff --git a/llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h b/llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h --- a/llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h +++ b/llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h @@ -14,6 +14,9 @@ #ifndef LLVM_TRANSFORMS_INSTRUMENTATION_ADDRESSSANITIZERCOMMON_H #define LLVM_TRANSFORMS_INSTRUMENTATION_ADDRESSSANITIZERCOMMON_H +#include "llvm/Analysis/CFG.h" +#include "llvm/Analysis/PostDominators.h" +#include "llvm/IR/Dominators.h" #include "llvm/IR/Instruction.h" #include "llvm/IR/Module.h" @@ -44,6 +47,47 @@ Value *getPtr() { return PtrUse->get(); } }; +// For an alloca valid between lifetime markers Start and End, call the +// Callback for all possible exits out of the lifetime in the containing +// function, which can return from the instructions in RetVec. +// +// Returns whether End was the only possible exit. If it wasn't, the caller +// should remove End to ensure that work done at the other exits does not +// happen outside of the lifetime. +template +bool forAllReachableExits(DominatorTree *DT, PostDominatorTree *PDT, + const Instruction *Start, Instruction *End, + const SmallVectorImpl &RetVec, + F Callback) { + // We need to ensure that if we tag some object, we certainly untag it + // before the function exits. + if (PDT != nullptr && PDT->dominates(End, Start)) { + Callback(End); + } else { + SmallVector ReachableRetVec; + unsigned NumCoveredExits = 0; + for (auto &RI : RetVec) { + if (!isPotentiallyReachable(Start, RI, nullptr, DT)) + continue; + ReachableRetVec.push_back(RI); + if (DT != nullptr && DT->dominates(End, RI)) + ++NumCoveredExits; + } + // If there's a mix of covered and non-covered exits, just put the untag + // on exits, so we avoid the redundancy of untagging twice. + if (NumCoveredExits == ReachableRetVec.size()) { + Callback(End); + } else { + for (auto &RI : ReachableRetVec) + Callback(RI); + // We may have inserted untag outside of the lifetime interval. + // Signal the caller to remove the lifetime end call for this alloca. + return false; + } + } + return true; +} + } // 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 @@ -42,17 +42,18 @@ #include "llvm/IR/Dominators.h" #include "llvm/IR/Function.h" #include "llvm/IR/GetElementPtrTypeIterator.h" +#include "llvm/IR/IRBuilder.h" #include "llvm/IR/Instruction.h" #include "llvm/IR/Instructions.h" #include "llvm/IR/IntrinsicInst.h" #include "llvm/IR/IntrinsicsAArch64.h" -#include "llvm/IR/IRBuilder.h" #include "llvm/IR/Metadata.h" #include "llvm/InitializePasses.h" #include "llvm/Pass.h" #include "llvm/Support/Casting.h" #include "llvm/Support/Debug.h" #include "llvm/Support/raw_ostream.h" +#include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h" #include "llvm/Transforms/Utils/Local.h" #include #include @@ -648,32 +649,10 @@ cast(Start->getArgOperand(0))->getZExtValue(); Size = alignTo(Size, kTagGranuleSize); tagAlloca(AI, Start->getNextNode(), Start->getArgOperand(1), Size); - // We need to ensure that if we tag some object, we certainly untag it - // before the function exits. - if (PDT != nullptr && PDT->dominates(End, Start)) { - untagAlloca(AI, End, Size); - } else { - SmallVector ReachableRetVec; - unsigned NumCoveredExits = 0; - for (auto &RI : RetVec) { - if (!isPotentiallyReachable(Start, RI, nullptr, DT)) - continue; - ReachableRetVec.push_back(RI); - if (DT != nullptr && DT->dominates(End, RI)) - ++NumCoveredExits; - } - // If there's a mix of covered and non-covered exits, just put the untag - // on exits, so we avoid the redundancy of untagging twice. - if (NumCoveredExits == ReachableRetVec.size()) { - untagAlloca(AI, End, Size); - } else { - for (auto &RI : ReachableRetVec) - untagAlloca(AI, RI, Size); - // We may have inserted untag outside of the lifetime interval. - // Remove the lifetime end call for this alloca. - End->eraseFromParent(); - } - } + + auto TagEnd = [&](Instruction *Node) { untagAlloca(AI, Node, Size); }; + if (!forAllReachableExits(DT, PDT, Start, End, RetVec, TagEnd)) + End->eraseFromParent(); } else { uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8; Value *Ptr = IRB.CreatePointerCast(TagPCall, IRB.getInt8PtrTy());