3 files changed, 23 insertions, 6 deletions
diff --git a/llvm/lib/CodeGen/AsmPrinter/WinException.cpp b/llvm/lib/CodeGen/AsmPrinter/WinException.cpp
index 7e3a6d5a76f..ffa0ceb1a60 100644
--- a/llvm/lib/CodeGen/AsmPrinter/WinException.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/WinException.cpp
@@ -50,6 +50,14 @@ WinException::~WinException() {}
 /// endModule - Emit all exception information that should come after the
 /// content.
 void WinException::endModule() {
+  auto &OS = *Asm->OutStreamer;
+  const Module *M = MMI->getModule();
+  for (const Function &F : *M) {
+    if (F.hasFnAttribute("safeseh")) {
+      llvm::errs() << ".safeseh " << F.getName() << "\n";
+      OS.EmitCOFFSafeSEH(Asm->getSymbol(&F));
+    }
+  }
 }
 
 void WinException::beginFunction(const MachineFunction *MF) {
diff --git a/llvm/lib/MC/WinCOFFStreamer.cpp b/llvm/lib/MC/WinCOFFStreamer.cpp
index 41fc8e4681e..36dd691f07b 100644
--- a/llvm/lib/MC/WinCOFFStreamer.cpp
+++ b/llvm/lib/MC/WinCOFFStreamer.cpp
@@ -164,7 +164,8 @@ void MCWinCOFFStreamer::EmitCOFFSafeSEH(MCSymbol const *Symbol) {
       Triple::x86)
     return;
 
-  if (cast<MCSymbolCOFF>(Symbol)->isSafeSEH())
+  const MCSymbolCOFF *CSymbol = cast<MCSymbolCOFF>(Symbol);
+  if (CSymbol->isSafeSEH())
     return;
 
   MCSection *SXData = getContext().getObjectFileInfo()->getSXDataSection();
@@ -175,7 +176,12 @@ void MCWinCOFFStreamer::EmitCOFFSafeSEH(MCSymbol const *Symbol) {
   new MCSafeSEHFragment(Symbol, SXData);
 
   getAssembler().registerSymbol(*Symbol);
-  cast<MCSymbolCOFF>(Symbol)->setIsSafeSEH();
+  CSymbol->setIsSafeSEH();
+
+  // The Microsoft linker requires that the symbol type of a handler be
+  // function. Go ahead and oblige it here.
+  CSymbol->setType(COFF::IMAGE_SYM_DTYPE_FUNCTION
+                   << COFF::SCT_COMPLEX_TYPE_SHIFT);
 }
 
 void MCWinCOFFStreamer::EmitCOFFSectionIndex(MCSymbol const *Symbol) {
diff --git a/llvm/lib/Target/X86/X86WinEHState.cpp b/llvm/lib/Target/X86/X86WinEHState.cpp
index 0c4aabab880..16a94daa819 100644
--- a/llvm/lib/Target/X86/X86WinEHState.cpp
+++ b/llvm/lib/Target/X86/X86WinEHState.cpp
@@ -60,7 +60,7 @@ public:
 private:
   void emitExceptionRegistrationRecord(Function *F);
 
-  void linkExceptionRegistration(IRBuilder<> &Builder, Value *Handler);
+  void linkExceptionRegistration(IRBuilder<> &Builder, Function *Handler);
   void unlinkExceptionRegistration(IRBuilder<> &Builder);
   void addCXXStateStores(Function &F, MachineModuleInfo &MMI);
   void addSEHStateStores(Function &F, MachineModuleInfo &MMI);
@@ -365,11 +365,14 @@ Function *WinEHStatePass::generateLSDAInEAXThunk(Function *ParentFunc) {
 }
 
 void WinEHStatePass::linkExceptionRegistration(IRBuilder<> &Builder,
-                                               Value *Handler) {
+                                               Function *Handler) {
+  // Emit the .safeseh directive for this function.
+  Handler->addFnAttr("safeseh");
+
   Type *LinkTy = getEHLinkRegistrationType();
   // Handler = Handler
-  Handler = Builder.CreateBitCast(Handler, Builder.getInt8PtrTy());
-  Builder.CreateStore(Handler, Builder.CreateStructGEP(LinkTy, Link, 1));
+  Value *HandlerI8 = Builder.CreateBitCast(Handler, Builder.getInt8PtrTy());
+  Builder.CreateStore(HandlerI8, Builder.CreateStructGEP(LinkTy, Link, 1));
   // Next = [fs:00]
   Constant *FSZero =
       Constant::getNullValue(LinkTy->getPointerTo()->getPointerTo(257));