[WinCOFF] Add support for the .safeseh directive

.safeseh adds an entry to the .sxdata section to register all the
appropriate functions which may handle an exception.  This entry is not
a relocation to the symbol but instead the symbol table index of the
function.

llvm-svn: 238641

7 files changed, 89 insertions, 17 deletions

diff --git a/llvm/lib/MC/MCAsmStreamer.cpp b/llvm/lib/MC/MCAsmStreamer.cpp
index 572634a99ec..8a5624bc333 100644
--- a/llvm/lib/MC/MCAsmStreamer.cpp
+++ b/llvm/lib/MC/MCAsmStreamer.cpp
@@ -136,6 +136,7 @@ public:
   void EmitCOFFSymbolStorageClass(int StorageClass) override;
   void EmitCOFFSymbolType(int Type) override;
   void EndCOFFSymbolDef() override;
+  void EmitCOFFSafeSEH(MCSymbol const *Symbol) override;
   void EmitCOFFSectionIndex(MCSymbol const *Symbol) override;
   void EmitCOFFSecRel32(MCSymbol const *Symbol) override;
   void EmitELFSize(MCSymbol *Symbol, const MCExpr *Value) override;
@@ -486,6 +487,11 @@ void MCAsmStreamer::EndCOFFSymbolDef() {
   EmitEOL();
 }
 
+void MCAsmStreamer::EmitCOFFSafeSEH(MCSymbol const *Symbol) {
+  OS << "\t.safeseh\t" << *Symbol;
+  EmitEOL();
+}
+
 void MCAsmStreamer::EmitCOFFSectionIndex(MCSymbol const *Symbol) {
   OS << "\t.secidx\t" << *Symbol;
   EmitEOL();
diff --git a/llvm/lib/MC/MCAssembler.cpp b/llvm/lib/MC/MCAssembler.cpp
index c8e3371c8f3..cd14e33a7aa 100644
--- a/llvm/lib/MC/MCAssembler.cpp
+++ b/llvm/lib/MC/MCAssembler.cpp
@@ -473,6 +473,9 @@ uint64_t MCAssembler::computeFragmentSize(const MCAsmLayout &Layout,
   case MCFragment::FT_LEB:
     return cast<MCLEBFragment>(F).getContents().size();
 
+  case MCFragment::FT_SafeSEH:
+    return 4;
+
   case MCFragment::FT_Align: {
     const MCAlignFragment &AF = cast<MCAlignFragment>(F);
     unsigned Offset = Layout.getFragmentOffset(&AF);
@@ -705,6 +708,12 @@ static void writeFragment(const MCAssembler &Asm, const MCAsmLayout &Layout,
     break;
   }
 
+  case MCFragment::FT_SafeSEH: {
+    const MCSafeSEHFragment &SF = cast<MCSafeSEHFragment>(F);
+    OW->Write32(SF.getSymbol()->getIndex());
+    break;
+  }
+
   case MCFragment::FT_Org: {
     ++stats::EmittedOrgFragments;
     const MCOrgFragment &OF = cast<MCOrgFragment>(F);
@@ -1086,6 +1095,7 @@ void MCFragment::dump() {
   case MCFragment::FT_Dwarf: OS << "MCDwarfFragment"; break;
   case MCFragment::FT_DwarfFrame: OS << "MCDwarfCallFrameFragment"; break;
   case MCFragment::FT_LEB:   OS << "MCLEBFragment"; break;
+  case MCFragment::FT_SafeSEH:    OS << "MCSafeSEHFragment"; break;
   }
 
   OS << "<MCFragment " << (void*) this << " LayoutOrder:" << LayoutOrder
@@ -1178,6 +1188,13 @@ void MCFragment::dump() {
     OS << " Value:" << LF->getValue() << " Signed:" << LF->isSigned();
     break;
   }
+  case MCFragment::FT_SafeSEH: {
+    const MCSafeSEHFragment *F = cast<MCSafeSEHFragment>(this);
+    OS << "\n       ";
+    OS << " Sym:";
+    F->getSymbol()->print(OS);
+    break;
+  }
   }
   OS << ">";
 }
@@ -1215,5 +1232,6 @@ void MCAlignFragment::anchor() { }
 void MCFillFragment::anchor() { }
 void MCOrgFragment::anchor() { }
 void MCLEBFragment::anchor() { }
+void MCSafeSEHFragment::anchor() { }
 void MCDwarfLineAddrFragment::anchor() { }
 void MCDwarfCallFrameFragment::anchor() { }
diff --git a/llvm/lib/MC/MCObjectFileInfo.cpp b/llvm/lib/MC/MCObjectFileInfo.cpp
index e99f036af16..b6b72e6a798 100644
--- a/llvm/lib/MC/MCObjectFileInfo.cpp
+++ b/llvm/lib/MC/MCObjectFileInfo.cpp
@@ -714,6 +714,9 @@ void MCObjectFileInfo::InitCOFFMCObjectFileInfo(Triple T) {
       ".xdata", COFF::IMAGE_SCN_CNT_INITIALIZED_DATA | COFF::IMAGE_SCN_MEM_READ,
       SectionKind::getDataRel());
 
+  SXDataSection = Ctx->getCOFFSection(".sxdata", COFF::IMAGE_SCN_LNK_INFO,
+                                      SectionKind::getMetadata());
+
   TLSDataSection = Ctx->getCOFFSection(
       ".tls$", COFF::IMAGE_SCN_CNT_INITIALIZED_DATA | COFF::IMAGE_SCN_MEM_READ |
                    COFF::IMAGE_SCN_MEM_WRITE,
diff --git a/llvm/lib/MC/MCParser/COFFAsmParser.cpp b/llvm/lib/MC/MCParser/COFFAsmParser.cpp
index 82f7f228052..53b527fb9ef 100644
--- a/llvm/lib/MC/MCParser/COFFAsmParser.cpp
+++ b/llvm/lib/MC/MCParser/COFFAsmParser.cpp
@@ -57,6 +57,7 @@ class COFFAsmParser : public MCAsmParserExtension {
     addDirectiveHandler<&COFFAsmParser::ParseDirectiveEndef>(".endef");
     addDirectiveHandler<&COFFAsmParser::ParseDirectiveSecRel32>(".secrel32");
     addDirectiveHandler<&COFFAsmParser::ParseDirectiveSecIdx>(".secidx");
+    addDirectiveHandler<&COFFAsmParser::ParseDirectiveSafeSEH>(".safeseh");
     addDirectiveHandler<&COFFAsmParser::ParseDirectiveLinkOnce>(".linkonce");
 
     // Win64 EH directives.
@@ -118,6 +119,7 @@ class COFFAsmParser : public MCAsmParserExtension {
   bool ParseDirectiveEndef(StringRef, SMLoc);
   bool ParseDirectiveSecRel32(StringRef, SMLoc);
   bool ParseDirectiveSecIdx(StringRef, SMLoc);
+  bool ParseDirectiveSafeSEH(StringRef, SMLoc);
   bool parseCOMDATType(COFF::COMDATType &Type);
   bool ParseDirectiveLinkOnce(StringRef, SMLoc);
 
@@ -453,6 +455,21 @@ bool COFFAsmParser::ParseDirectiveSecRel32(StringRef, SMLoc) {
   return false;
 }
 
+bool COFFAsmParser::ParseDirectiveSafeSEH(StringRef, SMLoc) {
+  StringRef SymbolID;
+  if (getParser().parseIdentifier(SymbolID))
+    return TokError("expected identifier in directive");
+
+  if (getLexer().isNot(AsmToken::EndOfStatement))
+    return TokError("unexpected token in directive");
+
+  MCSymbol *Symbol = getContext().getOrCreateSymbol(SymbolID);
+
+  Lex();
+  getStreamer().EmitCOFFSafeSEH(Symbol);
+  return false;
+}
+
 bool COFFAsmParser::ParseDirectiveSecIdx(StringRef, SMLoc) {
   StringRef SymbolID;
   if (getParser().parseIdentifier(SymbolID))
diff --git a/llvm/lib/MC/MCStreamer.cpp b/llvm/lib/MC/MCStreamer.cpp
index 280dbe28816..426bf90afcf 100644
--- a/llvm/lib/MC/MCStreamer.cpp
+++ b/llvm/lib/MC/MCStreamer.cpp
@@ -555,6 +555,9 @@ void MCStreamer::EmitWinCFIEndProlog() {
   CurrentWinFrameInfo->PrologEnd = Label;
 }
 
+void MCStreamer::EmitCOFFSafeSEH(MCSymbol const *Symbol) {
+}
+
 void MCStreamer::EmitCOFFSectionIndex(MCSymbol const *Symbol) {
 }
 
diff --git a/llvm/lib/MC/WinCOFFObjectWriter.cpp b/llvm/lib/MC/WinCOFFObjectWriter.cpp
index ddec16cd42e..2f6d1805605 100644
--- a/llvm/lib/MC/WinCOFFObjectWriter.cpp
+++ b/llvm/lib/MC/WinCOFFObjectWriter.cpp
@@ -21,6 +21,7 @@
 #include "llvm/MC/MCAssembler.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCExpr.h"
+#include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCSection.h"
 #include "llvm/MC/MCSectionCOFF.h"
@@ -76,6 +77,13 @@ public:
   void set_name_offset(uint32_t Offset);
 
   bool should_keep() const;
+
+  int64_t getIndex() const { return Index; }
+  void setIndex(int Value) {
+    Index = Value;
+    if (MC)
+      MC->setIndex(static_cast<uint32_t>(Value));
+  }
 };
 
 // This class contains staging data for a COFF relocation entry.
@@ -219,6 +227,10 @@ bool COFFSymbol::should_keep() const {
     return true;
   }
 
+  // if this is a safeseh handler, keep it
+  if (MC && (MC->getFlags() & COFF::SF_SafeSEH))
+    return true;
+
   // if the section its in is being droped, drop it
   if (Section->Number == -1)
     return false;
@@ -416,11 +428,13 @@ void WinCOFFObjectWriter::DefineSymbol(const MCSymbol &Symbol,
     const MCSymbol *Base = Layout.getBaseSymbol(Symbol);
     coff_symbol->Data.Value = getSymbolValue(Symbol, Layout);
 
-    coff_symbol->Data.Type = (Symbol.getFlags() & 0x0000FFFF) >> 0;
-    coff_symbol->Data.StorageClass = (Symbol.getFlags() & 0x00FF0000) >> 16;
+    coff_symbol->Data.Type =
+        (Symbol.getFlags() & COFF::SF_TypeMask) >> COFF::SF_TypeShift;
+    coff_symbol->Data.StorageClass =
+        (Symbol.getFlags() & COFF::SF_ClassMask) >> COFF::SF_ClassShift;
 
     // If no storage class was specified in the streamer, define it here.
-    if (coff_symbol->Data.StorageClass == 0) {
+    if (coff_symbol->Data.StorageClass == COFF::IMAGE_SYM_CLASS_NULL) {
       bool IsExternal = Symbol.isExternal() ||
                         (!Symbol.getFragment() && !Symbol.isVariable());
 
@@ -828,13 +842,9 @@ void WinCOFFObjectWriter::WriteObject(MCAssembler &Asm,
 
   UseBigObj = NumberOfSections > COFF::MaxNumberOfSections16;
 
-  DenseMap<COFFSection *, int32_t> SectionIndices(
-      NextPowerOf2(NumberOfSections));
-
   // Assign section numbers.
   size_t Number = 1;
   for (const auto &Section : Sections) {
-    SectionIndices[Section.get()] = Number;
     Section->Number = Number;
     Section->Symbol->Data.SectionNumber = Number;
     Section->Symbol->Aux[0].Aux.SectionDefinition.Number = Number;
@@ -877,12 +887,13 @@ void WinCOFFObjectWriter::WriteObject(MCAssembler &Asm,
     if (Symbol->Section)
       Symbol->Data.SectionNumber = Symbol->Section->Number;
     if (Symbol->should_keep()) {
-      Symbol->Index = Header.NumberOfSymbols++;
+      Symbol->setIndex(Header.NumberOfSymbols++);
       // Update auxiliary symbol info.
       Symbol->Data.NumberOfAuxSymbols = Symbol->Aux.size();
       Header.NumberOfSymbols += Symbol->Data.NumberOfAuxSymbols;
-    } else
-      Symbol->Index = -1;
+    } else {
+      Symbol->setIndex(-1);
+    }
   }
 
   // Build string table.
@@ -904,11 +915,11 @@ void WinCOFFObjectWriter::WriteObject(MCAssembler &Asm,
   // Fixup weak external references.
   for (auto &Symbol : Symbols) {
     if (Symbol->Other) {
-      assert(Symbol->Index != -1);
+      assert(Symbol->getIndex() != -1);
       assert(Symbol->Aux.size() == 1 && "Symbol must contain one aux symbol!");
       assert(Symbol->Aux[0].AuxType == ATWeakExternal &&
              "Symbol's aux symbol must be a Weak External!");
-      Symbol->Aux[0].Aux.WeakExternal.TagIndex = Symbol->Other->Index;
+      Symbol->Aux[0].Aux.WeakExternal.TagIndex = Symbol->Other->getIndex();
     }
   }
 
@@ -934,8 +945,7 @@ void WinCOFFObjectWriter::WriteObject(MCAssembler &Asm,
     if (Assoc->Number == -1)
       continue;
 
-    Section->Symbol->Aux[0].Aux.SectionDefinition.Number =
-        SectionIndices[Assoc];
+    Section->Symbol->Aux[0].Aux.SectionDefinition.Number = Assoc->Number;
   }
 
   // Assign file offsets to COFF object file structures.
@@ -984,8 +994,8 @@ void WinCOFFObjectWriter::WriteObject(MCAssembler &Asm,
       offset += COFF::RelocationSize * Sec->Relocations.size();
 
       for (auto &Relocation : Sec->Relocations) {
-        assert(Relocation.Symb->Index != -1);
-        Relocation.Data.SymbolTableIndex = Relocation.Symb->Index;
+        assert(Relocation.Symb->getIndex() != -1);
+        Relocation.Data.SymbolTableIndex = Relocation.Symb->getIndex();
       }
     }
 
@@ -1067,7 +1077,7 @@ void WinCOFFObjectWriter::WriteObject(MCAssembler &Asm,
          "Header::PointerToSymbolTable is insane!");
 
   for (auto &Symbol : Symbols)
-    if (Symbol->Index != -1)
+    if (Symbol->getIndex() != -1)
       WriteSymbol(*Symbol);
 
   OS.write(Strings.data().data(), Strings.data().size());
diff --git a/llvm/lib/MC/WinCOFFStreamer.cpp b/llvm/lib/MC/WinCOFFStreamer.cpp
index 72c18788068..abbe4164958 100644
--- a/llvm/lib/MC/WinCOFFStreamer.cpp
+++ b/llvm/lib/MC/WinCOFFStreamer.cpp
@@ -158,6 +158,21 @@ void MCWinCOFFStreamer::EndCOFFSymbolDef() {
   CurSymbol = nullptr;
 }
 
+void MCWinCOFFStreamer::EmitCOFFSafeSEH(MCSymbol const *Symbol) {
+  if (Symbol->getFlags() & COFF::SF_SafeSEH)
+    return;
+
+  MCSection *SXData = getContext().getObjectFileInfo()->getSXDataSection();
+  getAssembler().registerSection(*SXData);
+  if (SXData->getAlignment() < 4)
+    SXData->setAlignment(4);
+
+  new MCSafeSEHFragment(Symbol, SXData);
+
+  getAssembler().registerSymbol(*Symbol);
+  Symbol->modifyFlags(COFF::SF_SafeSEH, COFF::SF_SafeSEH);
+}
+
 void MCWinCOFFStreamer::EmitCOFFSectionIndex(MCSymbol const *Symbol) {
   MCDataFragment *DF = getOrCreateDataFragment();
   const MCSymbolRefExpr *SRE = MCSymbolRefExpr::create(Symbol, getContext());