1 files changed, 150 insertions, 22 deletions

diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
index 9b5b0c301f1..6f035ae27ab 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -26,6 +26,7 @@
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/FileOutputBuffer.h"
 #include "llvm/Support/Parallel.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/RandomNumberGenerator.h"
 #include <algorithm>
 #include <cstdio>
@@ -123,6 +124,12 @@ private:
   void openFile(StringRef OutputPath);
   template <typename PEHeaderTy> void writeHeader();
   void createSEHTable(OutputSection *RData);
+  void createGFIDTable(OutputSection *RData);
+  void markSymbolsForRVATable(ObjFile *File,
+                              ArrayRef<SectionChunk *> SymIdxChunks,
+                              SymbolRVASet &TableSymbols);
+  void maybeAddRVATable(OutputSection *RData, SymbolRVASet TableSymbols,
+                        StringRef TableSym, StringRef CountSym);
   void setSectionPermissions();
   void writeSections();
   void writeBuildId();
@@ -146,7 +153,8 @@ private:
   IdataContents Idata;
   DelayLoadContents DelayIdata;
   EdataContents Edata;
-  SEHTableChunk *SEHTable = nullptr;
+  RVATableChunk *GuardFidsTable = nullptr;
+  RVATableChunk *SEHTable = nullptr;
 
   Chunk *DebugDirectory = nullptr;
   std::vector<Chunk *> DebugRecords;
@@ -428,7 +436,13 @@ void Writer::createMiscChunks() {
       RData->addChunk(C);
   }
 
-  createSEHTable(RData);
+  // Create SEH table. x86-only.
+  if (Config->Machine == I386)
+    createSEHTable(RData);
+
+  // Create the guard function id table if requested.
+  if (Config->GuardCF)
+    createGFIDTable(RData);
 }
 
 // Create .idata section for the DLL-imported symbol table.
@@ -720,6 +734,8 @@ template <typename PEHeaderTy> void Writer::writeHeader() {
     PE->DLLCharacteristics |= IMAGE_DLL_CHARACTERISTICS_NX_COMPAT;
   if (!Config->AllowIsolation)
     PE->DLLCharacteristics |= IMAGE_DLL_CHARACTERISTICS_NO_ISOLATION;
+  if (Config->GuardCF)
+    PE->DLLCharacteristics |= IMAGE_DLL_CHARACTERISTICS_GUARD_CF;
   if (Config->Machine == I386 && !SEHTable &&
       !Symtab->findUnderscore("_load_config_used"))
     PE->DLLCharacteristics |= IMAGE_DLL_CHARACTERISTICS_NO_SEH;
@@ -825,34 +841,146 @@ void Writer::openFile(StringRef Path) {
 }
 
 void Writer::createSEHTable(OutputSection *RData) {
-  // Create SEH table. x86-only.
-  if (Config->Machine != I386)
-    return;
+  SymbolRVASet Handlers;
+  for (ObjFile *File : ObjFile::Instances) {
+    // FIXME: We should error here instead of earlier unless /safeseh:no was
+    // passed.
+    if (!File->hasSafeSEH())
+      return;
+
+    markSymbolsForRVATable(File, File->getSXDataChunks(), Handlers);
+  }
 
-  std::set<Defined *> Handlers;
+  maybeAddRVATable(RData, std::move(Handlers), "__safe_se_handler_table",
+                   "__safe_se_handler_count");
+}
+
+// Add a symbol to an RVA set. Two symbols may have the same RVA, but an RVA set
+// cannot contain duplicates. Therefore, the set is uniqued by Chunk and the
+// symbol's offset into that Chunk.
+static void addSymbolToRVASet(SymbolRVASet &RVASet, Defined *S) {
+  Chunk *C = S->getChunk();
+  if (auto *SC = dyn_cast<SectionChunk>(C))
+    C = SC->Repl; // Look through ICF replacement.
+  uint32_t Off = S->getRVA() - (C ? C->getRVA() : 0);
+  RVASet.insert({C, Off});
+}
+
+// Visit all relocations from all section contributions of this object file and
+// mark the relocation target as address-taken.
+static void markSymbolsWithRelocations(ObjFile *File,
+                                       SymbolRVASet &UsedSymbols) {
+  for (Chunk *C : File->getChunks()) {
+    // We only care about live section chunks. Common chunks and other chunks
+    // don't generally contain relocations.
+    SectionChunk *SC = dyn_cast<SectionChunk>(C);
+    if (!SC || !SC->isLive())
+      continue;
+
+    // Look for relocations in this section against symbols in executable output
+    // sections.
+    for (Symbol *Ref : SC->symbols()) {
+      // FIXME: Do further testing to see if the relocation type matters,
+      // especially for 32-bit where taking the address of something usually
+      // uses an absolute relocation instead of a relative one.
+      if (auto *D = dyn_cast_or_null<Defined>(Ref)) {
+        Chunk *RefChunk = D->getChunk();
+        OutputSection *OS = RefChunk ? RefChunk->getOutputSection() : nullptr;
+        if (OS && OS->getPermissions() & IMAGE_SCN_MEM_EXECUTE)
+          addSymbolToRVASet(UsedSymbols, D);
+      }
+    }
+  }
+}
 
+// Create the guard function id table. This is a table of RVAs of all
+// address-taken functions. It is sorted and uniqued, just like the safe SEH
+// table.
+void Writer::createGFIDTable(OutputSection *RData) {
+  SymbolRVASet AddressTakenSyms;
   for (ObjFile *File : ObjFile::Instances) {
-    if (!File->SEHCompat)
-      return;
-    for (uint32_t I : File->SXData)
-      if (Symbol *B = File->getSymbol(I))
-        if (B->isLive())
-          Handlers.insert(cast<Defined>(B));
+    // If the object was compiled with /guard:cf, the address taken symbols are
+    // in the .gfids$y sections. Otherwise, we approximate the set of address
+    // taken symbols by checking which symbols were used by relocations in live
+    // sections.
+    if (File->hasGuardCF())
+      markSymbolsForRVATable(File, File->getGuardFidChunks(), AddressTakenSyms);
+    else
+      markSymbolsWithRelocations(File, AddressTakenSyms);
   }
 
-  if (Handlers.empty())
+  // Mark the image entry as address-taken.
+  if (Config->Entry)
+    addSymbolToRVASet(AddressTakenSyms, cast<Defined>(Config->Entry));
+
+  maybeAddRVATable(RData, std::move(AddressTakenSyms), "__guard_fids_table",
+                   "__guard_fids_count");
+
+  // Set __guard_flags, which will be used in the load config to indicate that
+  // /guard:cf was enabled.
+  uint32_t GuardFlags = uint32_t(coff_guard_flags::CFInstrumented) |
+                        uint32_t(coff_guard_flags::HasFidTable);
+  Symbol *FlagSym = Symtab->findUnderscore("__guard_flags");
+  cast<DefinedAbsolute>(FlagSym)->setVA(GuardFlags);
+}
+
+// Take a list of input sections containing symbol table indices and add those
+// symbols to an RVA table. The challenge is that symbol RVAs are not known and
+// depend on the table size, so we can't directly build a set of integers.
+void Writer::markSymbolsForRVATable(ObjFile *File,
+                                    ArrayRef<SectionChunk *> SymIdxChunks,
+                                    SymbolRVASet &TableSymbols) {
+  for (SectionChunk *C : SymIdxChunks) {
+    // Skip sections discarded by linker GC. This comes up when a .gfids section
+    // is associated with something like a vtable and the vtable is discarded.
+    // In this case, the associated gfids section is discarded, and we don't
+    // mark the virtual member functions as address-taken by the vtable.
+    if (!C->isLive())
+      continue;
+
+    // Validate that the contents look like symbol table indices.
+    ArrayRef<uint8_t> Data = C->getContents();
+    if (Data.size() % 4 != 0) {
+      warn("ignoring " + C->getSectionName() +
+           " symbol table index section in object " + toString(File));
+      continue;
+    }
+
+    // Read each symbol table index and check if that symbol was included in the
+    // final link. If so, add it to the table symbol set.
+    ArrayRef<ulittle32_t> SymIndices(
+        reinterpret_cast<const ulittle32_t *>(Data.data()), Data.size() / 4);
+    ArrayRef<Symbol *> ObjSymbols = File->getSymbols();
+    for (uint32_t SymIndex : SymIndices) {
+      if (SymIndex >= ObjSymbols.size()) {
+        warn("ignoring invalid symbol table index in section " +
+             C->getSectionName() + " in object " + toString(File));
+        continue;
+      }
+      if (Symbol *S = ObjSymbols[SymIndex]) {
+        if (S->isLive())
+          addSymbolToRVASet(TableSymbols, cast<Defined>(S));
+      }
+    }
+  }
+}
+
+// Replace the absolute table symbol with a synthetic symbol pointing to
+// TableChunk so that we can emit base relocations for it and resolve section
+// relative relocations.
+void Writer::maybeAddRVATable(OutputSection *RData,
+                              SymbolRVASet TableSymbols,
+                              StringRef TableSym, StringRef CountSym) {
+  if (TableSymbols.empty())
     return;
 
-  SEHTable = make<SEHTableChunk>(Handlers);
-  RData->addChunk(SEHTable);
+  RVATableChunk *TableChunk = make<RVATableChunk>(std::move(TableSymbols));
+  RData->addChunk(TableChunk);
 
-  // Replace the absolute table symbol with a synthetic symbol pointing to the
-  // SEHTable chunk so that we can emit base relocations for it and resolve
-  // section relative relocations.
-  Symbol *T = Symtab->find("___safe_se_handler_table");
-  Symbol *C = Symtab->find("___safe_se_handler_count");
-  replaceSymbol<DefinedSynthetic>(T, T->getName(), SEHTable);
-  cast<DefinedAbsolute>(C)->setVA(SEHTable->getSize() / 4);
+  Symbol *T = Symtab->findUnderscore(TableSym);
+  Symbol *C = Symtab->findUnderscore(CountSym);
+  replaceSymbol<DefinedSynthetic>(T, T->getName(), TableChunk);
+  cast<DefinedAbsolute>(C)->setVA(TableChunk->getSize() / 4);
 }
 
 // Handles /section options to allow users to overwrite