Recommit DwarfEHPrepare: insert extra unwind paths for stack protector to instrument

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.

Changes:
  * Move test to AArch64 directory as it relies on target presence.
  * Re-add Dominator-tree maintenance. Accidentally cherry-picked wrong patch.
  * Skip adding paths on Windows EH functions.

https://reviews.llvm.org/D143637

Author: TNorthover

llvm/lib/CodeGen/DwarfEHPrepare.cpp

@@ -18,6 +18,7 @@
 #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"
@@ -28,6 +29,7 @@
 #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"
@@ -37,6 +39,7 @@
 #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;
@@ -166,7 +169,136 @@ size_t DwarfEHPrepare::pruneUnreachableResumes(
   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() &&
+      !isScopedEHPersonality(classifyEHPersonality(F.getPersonalityFn())) &&
+      StackProtector::requiresStackProtector(&F, nullptr))
+    addCleanupPathsForStackProtector(F, DTU);
+
   SmallVector<ResumeInst *, 16> Resumes;
   SmallVector<LandingPadInst *, 16> CleanupLPads;
   if (F.doesNotThrow())

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

@@ -0,0 +1,140 @@
+; 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)