StackProtector: add unwind cleanup paths for instrumentation.

[t.p.northover]

Summary

This is a mitigation patch for https://bugs.chromium.org/p/llvm/issues/detail?id=30, where existing stack protection is skipped if a function is returned through by an unwinder rather than the normal call/return path. The recent patch D139254 added the ability to instrument a visible unwind path, at least in the IR case (I'm working on the SelectionDAG instrumentation too) but there are still invisible unwinds it can't reach.

So this patch adds logic to DwarfEHPrepare that goes through a function, converting any call that might throw into an invoke to a simple resume cleanup, and adding cleanup clauses to existing landingpads that lack them. Obviously we don't really want to do this if it's wasted effort, so I also exposed requiresStackProtector from the actual StackProtector code to skip the extra paths if they won't be used.

[t.p.northover]

Ping.

[t.p.northover]

Realised I hadn't updated the DomTree correctly first time round.

[xiangzhangllvm]

I took a quick look, Seems we can't reproduce the report for test in https://bugs.chromium.org/p/llvm/issues/detail?id=30 in latest llvm.

$./a.out
*** stack smashing detected ***: terminated
Aborted (core dumped)
$./a.out test
*** stack smashing detected ***: terminated
Aborted (core dumped)

or we mark "noreturn" for all the throw cases.
I'll take a deeper look on this patch. Anyway thanks first.

[xiangzhangllvm]

llvm/lib/CodeGen/DwarfEHPrepare.cpp

297

I just header the func __gxx_personality_v0 in libc++ will handle the exception.
Does all the callees in the caller (with personality attribute) may do invisible unwinds for exception ? So you visibly generate the cleanup resume (to let current stack protect can see it then generate protection code. )

[t.p.northover]

llvm/lib/CodeGen/DwarfEHPrepare.cpp

297

That's pretty much it.

Actually, even without a personality function callees might unwind through this function (that just needs uwtables or similar) but without a personality function we have no way to intercept the exception and check the stack is OK on the way back through.

[xiangzhangllvm]

llvm/lib/CodeGen/DwarfEHPrepare.cpp

176

LGTM, just a question.
Do we will really generate a Landing Pad for try catch pairs without cleanup in real world ?
Should it report error earlier if without cleanup?
Anyway this patch enhance for them.

[xiangzhangllvm]

llvm/lib/CodeGen/DwarfEHPrepare.cpp

176

Sorry, here only for personality function not visible try catch.

[t.p.northover]

Thanks, committed as 203b6f31bb71.

llvm/lib/CodeGen/DwarfEHPrepare.cpp

176

Do we will really generate a Landing Pad for try catch pairs without cleanup in real world ?

Yes, cleanup is only used if there's some code that needs to be run regardless of whether an exception is caught (e.g. C++ destructors). So for example this leads to a landingpad without cleanup:

void func();
void foo() {
  try {
    func();
  } catch (int i) {
  }
}

[thakis]

It looks like this broke building libunwind. Building it now fails with (when building clang and lld and so on with this patch, and then building libunwind with just-built clang and lld and so on):

ld64.lld: error: undefined symbol: __gxx_personality_v0

I suppose this needs some flag to turn it off when building libunwind?

This is in a GN build. I didn't check if it's also broken in a cmake build as it's difficult (for me, at least) to do a build of libunwind that builds with just-build clang and lld. The GN build doesn't build libunwind with any -fstack-protector- or _FORTIFY_SOURCE things, but maybe it does something else silly that causes stack protection to fire when it shouldn't (…any ideas what this might be?). Or maybe stack protection is just enabled by default on darwin? But I don't see a fno-stack-protector flag in the libunwind cmake build either.

…hm, adding "-fno-stack-protector" to libunwind's cflags seems to unbreak the libunwind build. So I figure I'll add that in the GN build. The CMake build doesn't seem to explicitly add that, so I'm guessing libunwind is broken there too. Can you look at that?

(clang/cmake/caches/Apple-stage2.cmake globally adds -fno-stack-protector as far as I can tell, so the Apple bots probably don't show this if they use that cache file. But if you don't use that cache file, it's probably broken now?)

I filed https://github.com/llvm/llvm-project/issues/61501 for this.

[hans]

Heads up that we bisected test failures in Chromium (targeting iOS/x86_64 simulator) to this revision: https://bugs.chromium.org/p/chromium/issues/detail?id=1428624 Hopefully we'll be able to figure out what's going on next week, but if anyone else has seen problems too it would be interesting to know.

[hans]

In D143637#4237166, @hans wrote:

Heads up that we bisected test failures in Chromium (targeting iOS/x86_64 simulator) to this revision: https://bugs.chromium.org/p/chromium/issues/detail?id=1428624 Hopefully we'll be able to figure out what's going on next week, but if anyone else has seen problems too it would be interesting to know.

I still don't have a root cause or stand-alone repro, but I did notice that this kicks in even with clang's -fno-exceptions, which seems like it might be wrong?

[hans]

Finally, here's a stand-alone reproducer for desktop: https://bugs.chromium.org/p/chromium/issues/detail?id=1428624#c11

Can we revert this until that's addressed?

[t.p.northover]

Oh dear, thanks for investigating this, I didn't notice until I got an e-mail about the revert.

ObjC never fails to disappoint it seems; I'll see what I can do about accommodating it.

[t.p.northover]

I think I'd like to reapply this one.

The __gxx_personality_v0 error only reproduces for me if I compile libunwind with -fexceptions (or perhaps more likely remove the LLVM default -fno-exceptions), and I think that's not something that should be possible. It conceptually introduces exactly that dependence from libunwind -> libc++abi lld is complaining about, which is in the wrong direction (and libunwind doesn't actually do exceptions anyway). I've been unable to get this to happen in an actual CMake build though, LLVM_ENABLE_EH doesn't get forwarded to the LLVM_EANBLE_RUNTIMES builds in my tests.

So if necessary I think we should add logic to libunwind to ensure -fno-exceptions, but I haven't updated the patch for now because I can't find a scenario where it triggers.

And, having read up a bit on ObjC lifetime semantics, I think the test-case there (and the Chromium one it's derived from, which is basically the same) is relying on lifetime optimizations that aren't actually guaranteed. In main:

int main() {
  NSObject *p = create();
  __weak NSObject *q = p;
  @autoreleasepool {
    p = nil;
  }
  printf("p = %p\n", p);
  printf("q = %p\n", q); // Both p and q should be nil.
  return 0;
}

create returns an object that's notionally got a retain count of 1 and a queued autorelease.

Before this change, the retain that was part of assigning to p cancelled out the autorelease (the call to _objc_retainAutoreleasedReturnValue was right after create) so it went into the autoreleasepool block with a count of 1 and nothing pending. Assigning nil decremented the refcount and freed the object (zeroing weak ref q).

After this change the retain actually happened (the invoke separated the call to create from _objc_retainAutoreleasedReturnValue). So instead so it went into the block with a refcount of 2 and an autorelease pending. Since the object wasn't deleted, the weak reference never got zeroed out (I assume the real runtime will drain all autoreleasing objects at some point, just not here). But it was only ever an opportunistic optimization.

It's notable that

1. NSObject *p; p = create(); always failed to be optimized and q remains valid even now.
2. This is only at -O0: we actually make a stronger effort to ensure an _objc_retainAutoreleasedReturnValue stays with its call at higher levels so the optimization would still happen there.

So I think the Chromium test needs updating in this case (though I don't really know what it was trying to do so I'm not sure how). @hans, what do you think? Is that feasible? Do you disagree with the reasoning?

[hans]

So I think the Chromium test needs updating in this case (though I don't really know what it was trying to do so I'm not sure how). @hans, what do you think? Is that feasible? Do you disagree with the reasoning?

I'm certainly not an expert on ObjC lifetime management. Based on the original test (quoted in https://bugs.chromium.org/p/chromium/issues/detail?id=1428624#c4) it sounds like it's already sensitive to when exactly things get released. So if we're confident that the new behavior is also correct, I suppose the test should be updated.

But, I also notice that my reproducer works with -fno-exceptions, which is surprising to me. Why is this change also affecting -fno-exceptions code?

[smeenai]

There's actually a valid reason to build libunwind with -fexceptions, which is getting functions like _Unwind_Resume to not be marked nounwind, which would otherwise cause invalid LSDA to be generated when you LTO libunwind together with other libraries. See https://github.com/llvm/llvm-project/issues/56825 and @jyknight's comment on it.

When we reland this, would there be any way to account for the -disable-check-noreturn-call flag added in D141556, to avoid the size increase for people who don't want to pay it? CC @nickdesaulniers, who also had size concerns that motivated D147975.

llvm/include/llvm/CodeGen/StackProtector.h

private:
   /// StackProtector analysis will update this map when determining if an
   /// AllocaInst triggers a stack protector.
   SSPLayoutMap Layout;
 
   /// The minimum size of buffers that will receive stack smashing
   /// protection when -fstack-protection is used.
   unsigned SSPBufferSize = DefaultSSPBufferSize;
 
-  /// VisitedPHIs - The set of PHI nodes visited when determining
-  /// if a variable's reference has been taken.  This set
-  /// is maintained to ensure we don't visit the same PHI node multiple
-  /// times.
-  SmallPtrSet<const PHINode *, 16> VisitedPHIs;
-
   // A prologue is generated.
   bool HasPrologue = false;
 
   // IR checking code is generated.
   bool HasIRCheck = false;
 
   /// InsertStackProtectors - Insert code into the prologue and epilogue of
   /// the function.
   ///
   ///  - The prologue code loads and stores the stack guard onto the stack.
   ///  - The epilogue checks the value stored in the prologue against the
   ///    original value. It calls __stack_chk_fail if they differ.
   bool InsertStackProtectors();
 
   /// CreateFailBB - Create a basic block to jump to when the stack protector
   /// check fails.
   BasicBlock *CreateFailBB();
 
-  /// ContainsProtectableArray - Check whether the type either is an array or
-  /// contains an array of sufficient size so that we need stack protectors
-  /// for it.
-  /// \param [out] IsLarge is set to true if a protectable array is found and
-  /// it is "large" ( >= ssp-buffer-size).  In the case of a structure with
-  /// multiple arrays, this gets set if any of them is large.
-  bool ContainsProtectableArray(Type *Ty, bool &IsLarge, bool Strong = false,
-                                bool InStruct = false) const;
-
-  /// Check whether a stack allocation has its address taken.
-  bool HasAddressTaken(const Instruction *AI, TypeSize AllocSize);
-
-  /// RequiresStackProtector - Check whether or not this function needs a
-  /// stack protector based upon the stack protector level.
-  bool RequiresStackProtector();
-
 public:
   static char ID; // Pass identification, replacement for typeid.
 
   StackProtector();
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   // Return true if StackProtector is supposed to be handled by SelectionDAG.
   bool shouldEmitSDCheck(const BasicBlock &BB) const;
 
   bool runOnFunction(Function &Fn) override;
 
   void copyToMachineFrameInfo(MachineFrameInfo &MFI) const;
+
+  /// Check whether or not \p F needs a stack protector based upon the stack
+  /// protector level.
+  static bool requiresStackProtector(Function *F, SSPLayoutMap *Layout = nullptr);
+
 };
 
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_STACKPROTECTOR_H

llvm/lib/CodeGen/DwarfEHPrepare.cpp

 
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/CodeGen/RuntimeLibcalls.h"
+#include "llvm/CodeGen/StackProtector.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/EHPersonalities.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include <cstddef>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "dwarfehprepare"
 
 STATISTIC(NumResumesLowered, "Number of resume calls lowered");
 STATISTIC(NumCleanupLandingPadsUnreachable,
if (ResumeReachable[I]) {
       RI->eraseFromParent();
       simplifyCFG(BB, *TTI, DTU);
     }
   }
   Resumes.resize(ResumesLeft);
   return ResumesLeft;
 }
 
+/// If a landingpad block doesn't already have a cleanup case, add one
+/// that feeds directly into a resume instruction.
+static void addCleanupResumeToLandingPad(BasicBlock &BB, DomTreeUpdater *DTU) {
+  LandingPadInst *LP = BB.getLandingPadInst();
+  if (LP->isCleanup())
+    return;
+
+  // There will usually be code testing for the other kinds of exception
+  // immediately after the landingpad. Working out the far end of that chain is
+  // tricky, so put our test for the new cleanup case (i.e. selector == 0) at
+  // the beginning.
+  BasicBlock *ContBB = SplitBlock(&BB, LP->getNextNode(), DTU);
+  BB.getTerminator()->eraseFromParent();
+
+  LP->setCleanup(true);
+  IRBuilder<> B(&BB);
+  Value *Selector = B.CreateExtractValue(LP, 1);
+  Value *Cmp = B.CreateICmpEQ(Selector, ConstantInt::get(Selector->getType(), 0));
+
+  Function *F = BB.getParent();
+  LLVMContext &Ctx = F->getContext();
+  BasicBlock *ResumeBB = BasicBlock::Create(Ctx, "resume", F);
+  ResumeInst::Create(LP, ResumeBB);
+
+  B.CreateCondBr(Cmp, ResumeBB, ContBB);
+  if (DTU) {
+    SmallVector<DominatorTree::UpdateType> Updates;
+    Updates.push_back({DominatorTree::Insert, &BB, ResumeBB});
+    DTU->applyUpdates(Updates);
+  }
+}
+
+/// Create a basic block that has a `landingpad` instruction feeding
+/// directly into a `resume`. Will be set to the unwind destination of a new
+/// invoke.
+static BasicBlock *createCleanupResumeBB(Function &F,  Type *LandingPadTy) {
+  LLVMContext &Ctx = F.getContext();
+  BasicBlock *BB = BasicBlock::Create(Ctx, "cleanup_resume", &F);
+  IRBuilder<> B(BB);
+
+  // If this is going to be the only landingpad in the function, synthesize the
+  // standard type all ABIs use, which is essentially `{ ptr, i32 }`.
+  if (!LandingPadTy)
+    LandingPadTy =
+        StructType::get(Type::getInt8PtrTy(Ctx), IntegerType::get(Ctx, 32));
+
+  LandingPadInst *Except = B.CreateLandingPad(LandingPadTy, 0);
+  Except->setCleanup(true);
+  B.CreateResume(Except);
+  return BB;
+}
+
+/// Convert a call that might throw into an invoke that unwinds to the specified
+/// simple landingpad/resume block.
+static void changeCallToInvokeResume(CallInst &CI, BasicBlock *CleanupResumeBB,
+                                     DomTreeUpdater *DTU) {
+  BasicBlock *BB = CI.getParent();
+  BasicBlock *ContBB = SplitBlock(BB, &CI, DTU);
+  BB->getTerminator()->eraseFromParent();
+
+  IRBuilder<> B(BB);
+  SmallVector<Value *> Args(CI.args());
+  SmallVector<OperandBundleDef> Bundles;
+  CI.getOperandBundlesAsDefs(Bundles);
+  InvokeInst *NewCall =
+      B.CreateInvoke(CI.getFunctionType(), CI.getCalledOperand(), ContBB,
+                     CleanupResumeBB, Args, Bundles, CI.getName());
+  NewCall->setAttributes(CI.getAttributes());
+  NewCall->setCallingConv(CI.getCallingConv());
+  NewCall->copyMetadata(CI);
+
+  if (DTU) {
+    SmallVector<DominatorTree::UpdateType> Updates;
+    Updates.push_back({DominatorTree::Insert, BB, CleanupResumeBB});
+    DTU->applyUpdates(Updates);
+  }
+  CI.replaceAllUsesWith(NewCall);
+  CI.eraseFromParent();
+}
+
+/// Ensure that any call in this function that might throw has an associated
+/// cleanup/resume that the stack protector can instrument later. Existing
+/// invokes will get an added `cleanup` clause if needed, calls will be
+/// converted to an invoke with trivial unwind followup.
+static void addCleanupPathsForStackProtector(Function &F, DomTreeUpdater *DTU) {
+  // First add cleanup -> resume paths to all existing landingpads, noting what
+  // type landingpads in this function actually have along the way.
+  Type *LandingPadTy = nullptr;
+  for (Function::iterator FI = F.begin(); FI != F.end(); ++FI) {
+    BasicBlock &BB = *FI;
+    if (LandingPadInst *LP = BB.getLandingPadInst()) {
+      // We can assume the type is broadly compatible with { ptr, i32 } since
+      // other parts of this pass already try to extract values from it.
+      LandingPadTy = LP->getType();
+      addCleanupResumeToLandingPad(BB, DTU);
+    }
+  }
+
+  // Next convert any call that might throw into an invoke to a resume
+  // instruction for later instrumentation.
+  BasicBlock *CleanupResumeBB = nullptr;
+  for (Function::iterator FI = F.begin(); FI != F.end(); ++FI) {
+    BasicBlock &BB = *FI;
+    for (Instruction &I : BB) {
+      CallInst *CI = dyn_cast<CallInst>(&I);
+      if (!CI || CI->doesNotThrow())
+        continue;
+
+      // Tail calls cannot use our stack so no need to check whether it was
+      // corrupted.
+      if (CI->isTailCall())
+        continue;
+
+      if (!CleanupResumeBB)
+        CleanupResumeBB = createCleanupResumeBB(F, LandingPadTy);
+
+      changeCallToInvokeResume(*CI, CleanupResumeBB, DTU);
+
+      // This block has been split, start again on its continuation.
+      break;
+    }
+  }
+}
+
 bool DwarfEHPrepare::InsertUnwindResumeCalls() {
+  if (F.hasPersonalityFn() &&
+      StackProtector::requiresStackProtector(&F, nullptr))
+    addCleanupPathsForStackProtector(F, DTU);
+
   SmallVector<ResumeInst *, 16> Resumes;
   SmallVector<LandingPadInst *, 16> CleanupLPads;
   if (F.doesNotThrow())
     NumNoUnwind++;
   else
     NumUnwind++;
   for (BasicBlock &BB : F) {
     if (auto *RI = dyn_cast<ResumeInst>(BB.getTerminator()))

llvm/lib/CodeGen/StackProtector.cpp

bool StackProtector::runOnFunction(Function &Fn) {
   TM = &getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
   Trip = TM->getTargetTriple();
   TLI = TM->getSubtargetImpl(Fn)->getTargetLowering();
   HasPrologue = false;
   HasIRCheck = false;
 
   SSPBufferSize = Fn.getFnAttributeAsParsedInteger(
       "stack-protector-buffer-size", DefaultSSPBufferSize);
-  if (!RequiresStackProtector())
+  if (!requiresStackProtector(F, &Layout))
     return false;
 
   // TODO(etienneb): Functions with funclets are not correctly supported now.
   // Do nothing if this is funclet-based personality.
   if (Fn.hasPersonalityFn()) {
     EHPersonality Personality = classifyEHPersonality(Fn.getPersonalityFn());
     if (isFuncletEHPersonality(Personality))
       return false;
   }
 
   ++NumFunProtected;
   bool Changed = InsertStackProtectors();
 #ifdef EXPENSIVE_CHECKS
   assert((!DTU ||
           DTU->getDomTree().verify(DominatorTree::VerificationLevel::Full)) &&
          "Failed to maintain validity of domtree!");
 #endif
   DTU.reset();
   return Changed;
 }
 
 /// \param [out] IsLarge is set to true if a protectable array is found and
 /// it is "large" ( >= ssp-buffer-size).  In the case of a structure with
 /// multiple arrays, this gets set if any of them is large.
-bool StackProtector::ContainsProtectableArray(Type *Ty, bool &IsLarge,
-                                              bool Strong,
-                                              bool InStruct) const {
+static bool ContainsProtectableArray(Type *Ty, Module *M, unsigned SSPBufferSize,
+                                     bool &IsLarge, bool Strong,
+                                     bool InStruct) {
   if (!Ty)
     return false;
   if (ArrayType *AT = dyn_cast<ArrayType>(Ty)) {
     if (!AT->getElementType()->isIntegerTy(8)) {
       // If we're on a non-Darwin platform or we're inside of a structure, don't
       // add stack protectors unless the array is a character array.
       // However, in strong mode any array, regardless of type and size,
       // triggers a protector.
-      if (!Strong && (InStruct || !Trip.isOSDarwin()))
+      if (!Strong && (InStruct || !Triple(M->getTargetTriple()).isOSDarwin()))
         return false;
     }
 
     // If an array has more than SSPBufferSize bytes of allocated space, then we
     // emit stack protectors.
     if (SSPBufferSize <= M->getDataLayout().getTypeAllocSize(AT)) {
       IsLarge = true;
       return true;
     }
 
     if (Strong)
       // Require a protector for all arrays in strong mode
       return true;
   }
 
   const StructType *ST = dyn_cast<StructType>(Ty);
   if (!ST)
     return false;
 
   bool NeedsProtector = false;
   for (Type *ET : ST->elements())
-    if (ContainsProtectableArray(ET, IsLarge, Strong, true)) {
+    if (ContainsProtectableArray(ET, M, SSPBufferSize, IsLarge, Strong, true)) {
       // If the element is a protectable array and is large (>= SSPBufferSize)
       // then we are done.  If the protectable array is not large, then
       // keep looking in case a subsequent element is a large array.
       if (IsLarge)
         return true;
       NeedsProtector = true;
     }
 
   return NeedsProtector;
 }
 
-bool StackProtector::HasAddressTaken(const Instruction *AI,
-                                     TypeSize AllocSize) {
+/// Check whether a stack allocation has its address taken.
+static bool HasAddressTaken(const Instruction *AI, TypeSize AllocSize,
+                            Module *M,
+                            SmallPtrSet<const PHINode *, 16> &VisitedPHIs) {
   const DataLayout &DL = M->getDataLayout();
   for (const User *U : AI->users()) {
     const auto *I = cast<Instruction>(U);
     // If this instruction accesses memory make sure it doesn't access beyond
     // the bounds of the allocated object.
     std::optional<MemoryLocation> MemLoc = MemoryLocation::getOrNone(I);
     if (MemLoc && MemLoc->Size.hasValue() &&
         !TypeSize::isKnownGE(AllocSize,
case Instruction::GetElementPtr: {
       TypeSize OffsetSize = TypeSize::Fixed(Offset.getLimitedValue());
       if (!TypeSize::isKnownGT(AllocSize, OffsetSize))
         return true;
       // Adjust AllocSize to be the space remaining after this offset.
       // We can't subtract a fixed size from a scalable one, so in that case
       // assume the scalable value is of minimum size.
       TypeSize NewAllocSize =
           TypeSize::Fixed(AllocSize.getKnownMinValue()) - OffsetSize;
-      if (HasAddressTaken(I, NewAllocSize))
+      if (HasAddressTaken(I, NewAllocSize, M, VisitedPHIs))
         return true;
       break;
     }
     case Instruction::BitCast:
     case Instruction::Select:
     case Instruction::AddrSpaceCast:
-      if (HasAddressTaken(I, AllocSize))
+      if (HasAddressTaken(I, AllocSize, M, VisitedPHIs))
         return true;
       break;
     case Instruction::PHI: {
       // Keep track of what PHI nodes we have already visited to ensure
       // they are only visited once.
       const auto *PN = cast<PHINode>(I);
       if (VisitedPHIs.insert(PN).second)
-        if (HasAddressTaken(PN, AllocSize))
+        if (HasAddressTaken(PN, AllocSize, M, VisitedPHIs))
           return true;
       break;
     }
     case Instruction::Load:
     case Instruction::AtomicRMW:
     case Instruction::Ret:
       // These instructions take an address operand, but have load-like or
       // other innocuous behavior that should not trigger a stack protector.

 /// call alloca with a either a variable size or a size >= SSPBufferSize,
 /// functions with character buffers larger than SSPBufferSize, and functions
 /// with aggregates containing character buffers larger than SSPBufferSize. The
 /// strong heuristic will add a guard variables to functions that call alloca
 /// regardless of size, functions with any buffer regardless of type and size,
 /// functions with aggregates that contain any buffer regardless of type and
 /// size, and functions that contain stack-based variables that have had their
 /// address taken.
-bool StackProtector::RequiresStackProtector() {
+bool StackProtector::requiresStackProtector(Function *F, SSPLayoutMap *Layout) {
+  Module *M = F->getParent();
   bool Strong = false;
   bool NeedsProtector = false;
 
+  // The set of PHI nodes visited when determining if a variable's reference has
+  // been taken.  This set is maintained to ensure we don't visit the same PHI
+  // node multiple times.
+  SmallPtrSet<const PHINode *, 16> VisitedPHIs;
+
+  unsigned SSPBufferSize = F->getFnAttributeAsParsedInteger(
+      "stack-protector-buffer-size", DefaultSSPBufferSize);
+
   if (F->hasFnAttribute(Attribute::SafeStack))
     return false;
 
   // We are constructing the OptimizationRemarkEmitter on the fly rather than
   // using the analysis pass to avoid building DominatorTree and LoopInfo which
   // are not available this late in the IR pipeline.
   OptimizationRemarkEmitter ORE(F);
 
   if (F->hasFnAttribute(Attribute::StackProtectReq)) {
+    if (!Layout)
+      return true;
     ORE.emit([&]() {
       return OptimizationRemark(DEBUG_TYPE, "StackProtectorRequested", F)
              << "Stack protection applied to function "
              << ore::NV("Function", F)
              << " due to a function attribute or command-line switch";
     });
     NeedsProtector = true;
     Strong = true; // Use the same heuristic as strong to determine SSPLayout
for (const Instruction &I : BB) {
                    << ore::NV("Function", F)
                    << " due to a call to alloca or use of a variable length "
                       "array";
           };
           if (const auto *CI = dyn_cast<ConstantInt>(AI->getArraySize())) {
             if (CI->getLimitedValue(SSPBufferSize) >= SSPBufferSize) {
               // A call to alloca with size >= SSPBufferSize requires
               // stack protectors.
-              Layout.insert(std::make_pair(AI,
-                                           MachineFrameInfo::SSPLK_LargeArray));
+              if (!Layout)
+                return true;
+              Layout->insert(
+                  std::make_pair(AI, MachineFrameInfo::SSPLK_LargeArray));
               ORE.emit(RemarkBuilder);
               NeedsProtector = true;
             } else if (Strong) {
               // Require protectors for all alloca calls in strong mode.
-              Layout.insert(std::make_pair(AI,
-                                           MachineFrameInfo::SSPLK_SmallArray));
+              if (!Layout)
+                return true;
+              Layout->insert(
+                  std::make_pair(AI, MachineFrameInfo::SSPLK_SmallArray));
               ORE.emit(RemarkBuilder);
               NeedsProtector = true;
             }
           } else {
             // A call to alloca with a variable size requires protectors.
-            Layout.insert(std::make_pair(AI,
-                                         MachineFrameInfo::SSPLK_LargeArray));
+            if (!Layout)
+              return true;
+            Layout->insert(
+                std::make_pair(AI, MachineFrameInfo::SSPLK_LargeArray));
             ORE.emit(RemarkBuilder);
             NeedsProtector = true;
           }
           continue;
         }
 
         bool IsLarge = false;
-        if (ContainsProtectableArray(AI->getAllocatedType(), IsLarge, Strong)) {
-          Layout.insert(std::make_pair(AI, IsLarge
-                                       ? MachineFrameInfo::SSPLK_LargeArray
-                                       : MachineFrameInfo::SSPLK_SmallArray));
+        if (ContainsProtectableArray(AI->getAllocatedType(), M, SSPBufferSize,
+                                     IsLarge, Strong, false)) {
+          if (!Layout)
+            return true;
+          Layout->insert(std::make_pair(
+              AI, IsLarge ? MachineFrameInfo::SSPLK_LargeArray
           ORE.emit([&]() {
             return OptimizationRemark(DEBUG_TYPE, "StackProtectorBuffer", &I)
                    << "Stack protection applied to function "
                    << ore::NV("Function", F)
                    << " due to a stack allocated buffer or struct containing a "
                       "buffer";
           });
           NeedsProtector = true;
           continue;
         }
 
-        if (Strong && HasAddressTaken(AI, M->getDataLayout().getTypeAllocSize(
-                                              AI->getAllocatedType()))) {
+        if (Strong &&
+            HasAddressTaken(
+                AI, M->getDataLayout().getTypeAllocSize(AI->getAllocatedType()),
+                M, VisitedPHIs)) {
           ++NumAddrTaken;
-          Layout.insert(std::make_pair(AI, MachineFrameInfo::SSPLK_AddrOf));
+          if (!Layout)
+            return true;
+          Layout->insert(std::make_pair(AI, MachineFrameInfo::SSPLK_AddrOf));
           ORE.emit([&]() {
             return OptimizationRemark(DEBUG_TYPE, "StackProtectorAddressTaken",
                                       &I)
                    << "Stack protection applied to function "
                    << ore::NV("Function", F)
                    << " due to the address of a local variable being taken";
           });
           NeedsProtector = true;

llvm/test/CodeGen/Generic/safestack-unwind.ll

+; RUN: opt --dwarfehprepare %s -S -o - -mtriple=aarch64-linux-gnu | FileCheck %s
+
+define void @call() sspreq personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define void @call()
+; CHECK:     invoke void @bar()
+; CHECK:             to label %[[CONT:.*]] unwind label %[[CLEANUP:.*]]
+; CHECK: [[CONT]]:
+; CHECK:     ret void
+; CHECK: [[CLEANUP]]:
+; CHECK:     [[LP:%.*]] = landingpad { ptr, i32 }
+; CHECK:             cleanup
+; CHECK:     [[EXN:%.*]] = extractvalue { ptr, i32 } [[LP]], 0
+; CHECK:     call void @_Unwind_Resume(ptr [[EXN]])
+; CHECK:     unreachable
+  call void @bar()
+  ret void
+}
+
+define void @invoke_no_cleanup() sspreq personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define void @invoke_no_cleanup
+; CHECK:   invoke void @bar()
+; CHECK:           to label %done unwind label %catch
+
+; CHECK: catch:
+; CHECK:   [[LP:%.*]] = landingpad { ptr, i32 }
+; CHECK:           cleanup
+; CHECK:           catch ptr null
+; CHECK:   [[SEL:%.*]] = extractvalue { ptr, i32 } [[LP]], 1
+; CHECK:   [[CMP:%.*]] = icmp eq i32 [[SEL]], 0
+; CHECK:   br i1 [[CMP]], label %[[RESUME:.*]], label %[[SPLIT:.*]]
+
+; CHECK: [[SPLIT]]:
+; CHECK:   br label %done
+
+; CHECK: done:
+; CHECK:   ret void
+
+; CHECK: [[RESUME]]:
+; CHECK:   [[EXN:%.*]] = extractvalue { ptr, i32 } [[LP]], 0
+; CHECK:   call void @_Unwind_Resume(ptr [[EXN]])
+; CHECK:   unreachable
+  invoke void @bar() to label %done unwind label %catch
+
+catch:
+  %lp = landingpad { ptr, i32 }
+          catch ptr null
+  br label %done
+
+done:
+  ret void
+}
+
+define void @invoke_no_cleanup_catches() sspreq personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define void @invoke_no_cleanup_catches
+; CHECK:   invoke void @bar()
+; CHECK:           to label %done unwind label %catch
+
+; CHECK: catch:
+; CHECK:   [[LP:%.*]] = landingpad { ptr, i32 }
+; CHECK:           cleanup
+; CHECK:           catch ptr null
+; CHECK:   [[SEL:%.*]] = extractvalue { ptr, i32 } [[LP]], 1
+; CEHCK:   [[CMP:%.*]] = icmp eq i32 [[SEL]], 0
+; CEHCK:   br i1 [[CMP]], label %[[RESUME:.*]], label %[[SPLIT:.*]]
+
+; CHECK: [[SPLIT]]:
+; CHECK:   %exn = extractvalue { ptr, i32 } %lp, 0
+; CHECK:   invoke ptr @__cxa_begin_catch(ptr %exn)
+; CHECK:            to label %[[SPLIT2:.*]] unwind label %[[CLEANUP_RESUME:.*]]
+
+; CHECK: [[SPLIT2]]:
+; CHECK:   invoke void @__cxa_end_catch()
+; CHECK:            to label  %[[SPLIT3:.*]] unwind label %[[CLEANUP_RESUME:.*]]
+
+; CHECK: [[SPLIT3]]:
+; CHECK:   br label %done
+
+; CHECK: done:
+; CHECK:   ret void
+
+; CHECK: [[RESUME]]:
+; CHECK:   [[EXN1:%.*]] = extractvalue { ptr, i32 } [[LP]], 0
+; CHECK:   br label %[[RESUME_MERGE:.*]]
+
+; CHECK: [[CLEANUP_RESUME]]:
+; CHECK:   [[LP:%.*]] = landingpad { ptr, i32 }
+; CHECK:           cleanup
+; CHECK:   [[EXN2:%.*]] = extractvalue { ptr, i32 } [[LP]], 0
+; CHECK:   br label %[[RESUME_MERGE]]
+
+; CHECK: [[RESUME_MERGE]]:
+; CHECK:   [[EXN_PHI:%.*]] = phi ptr [ [[EXN1]], %[[RESUME]] ], [ [[EXN2]], %[[CLEANUP_RESUME]] ]
+; CHECK:   call void @_Unwind_Resume(ptr [[EXN_PHI]])
+; CHECK:   unreachable
+  invoke void @bar() to label %done unwind label %catch
+
+catch:
+  %lp = landingpad { ptr, i32 }
+          catch ptr null
+  %exn = extractvalue { ptr, i32 } %lp, 0
+  call ptr @__cxa_begin_catch(ptr %exn)
+  call void @__cxa_end_catch()
+  br label %done
+
+done:
+  ret void
+}
+
+; Don't try to invoke any intrinsics.
+define ptr @call_intrinsic() sspreq personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define ptr @call_intrinsic
+; CHECK: call ptr @llvm.frameaddress.p0(i32 0)
+  %res = call ptr @llvm.frameaddress.p0(i32 0)
+  ret ptr %res
+}
+
+; Check we go along with the existing landingpad type, even if it's a bit
+; outside the normal.
+define void @weird_landingpad() sspreq personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define void @weird_landingpad
+; CHECK: landingpad { ptr, i64 }
+; CHECK: landingpad { ptr, i64 }
+  invoke void @bar() to label %done unwind label %catch
+
+catch:
+  %lp = landingpad { ptr, i64 }
+           catch ptr null
+  resume { ptr, i64 } %lp
+;  br label %done
+
+done:
+  call void @bar()
+  ret void
+}
+
+declare void @bar()
+declare i32 @__gxx_personality_v0(...)
+declare ptr @__cxa_begin_catch(ptr)
+declare void @__cxa_end_catch()
+declare ptr @llvm.frameaddress.p0(i32)