[LLD][COFF] Store and validate load config in SymbolTable (#120324)

cjacek · web-flow · commit ff29f38c02eb · 2024-12-29T11:43:45.000+01:00
Improve diagnostics for invalid load configurations.
diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
@@ -2824,6 +2824,7 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
 
   if (ctx.symtabEC)
     ctx.symtabEC->initializeECThunks();
+  ctx.forEachSymtab([](SymbolTable &symtab) { symtab.initializeLoadConfig(); });
 
   // Identify unreferenced COMDAT sections.
   if (config->doGC) {
diff --git a/lld/COFF/SymbolTable.cpp b/lld/COFF/SymbolTable.cpp
@@ -27,6 +27,7 @@
 #include <utility>
 
 using namespace llvm;
+using namespace llvm::support;
 
 namespace lld::coff {
 
@@ -596,6 +597,52 @@ std::pair<Symbol *, bool> SymbolTable::insert(StringRef name, InputFile *file) {
   return result;
 }
 
+void SymbolTable::initializeLoadConfig() {
+  auto sym =
+      dyn_cast_or_null<DefinedRegular>(findUnderscore("_load_config_used"));
+  if (!sym) {
+    if (ctx.config.guardCF != GuardCFLevel::Off)
+      Warn(ctx)
+          << "Control Flow Guard is enabled but '_load_config_used' is missing";
+    if (ctx.config.dependentLoadFlags)
+      Warn(ctx) << "_load_config_used not found, /dependentloadflag will have "
+                   "no effect";
+    return;
+  }
+
+  SectionChunk *sc = sym->getChunk();
+  if (!sc->hasData) {
+    Err(ctx) << "_load_config_used points to uninitialized data";
+    return;
+  }
+  uint64_t offsetInChunk = sym->getValue();
+  if (offsetInChunk + 4 > sc->getSize()) {
+    Err(ctx) << "_load_config_used section chunk is too small";
+    return;
+  }
+
+  ArrayRef<uint8_t> secContents = sc->getContents();
+  loadConfigSize =
+      *reinterpret_cast<const ulittle32_t *>(&secContents[offsetInChunk]);
+  if (offsetInChunk + loadConfigSize > sc->getSize()) {
+    Err(ctx) << "_load_config_used specifies a size larger than its containing "
+                "section chunk";
+    return;
+  }
+
+  uint32_t expectedAlign = ctx.config.is64() ? 8 : 4;
+  if (sc->getAlignment() < expectedAlign)
+    Warn(ctx) << "'_load_config_used' is misaligned (expected alignment to be "
+              << expectedAlign << " bytes, got " << sc->getAlignment()
+              << " instead)";
+  else if (!isAligned(Align(expectedAlign), offsetInChunk))
+    Warn(ctx) << "'_load_config_used' is misaligned (section offset is 0x"
+              << Twine::utohexstr(sym->getValue()) << " not aligned to "
+              << expectedAlign << " bytes)";
+
+  loadConfigSym = sym;
+}
+
 void SymbolTable::addEntryThunk(Symbol *from, Symbol *to) {
   entryThunks.push_back({from, to});
 }
diff --git a/lld/COFF/SymbolTable.h b/lld/COFF/SymbolTable.h
@@ -138,6 +138,10 @@ class SymbolTable {
       callback(pair.second);
   }
 
+  DefinedRegular *loadConfigSym = nullptr;
+  uint32_t loadConfigSize = 0;
+  void initializeLoadConfig();
+
 private:
   /// Given a name without "__imp_" prefix, returns a defined symbol
   /// with the "__imp_" prefix, if it exists.
diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
@@ -1837,22 +1837,10 @@ template <typename PEHeaderTy> void Writer::writeHeader() {
     dir[DEBUG_DIRECTORY].RelativeVirtualAddress = debugDirectory->getRVA();
     dir[DEBUG_DIRECTORY].Size = debugDirectory->getSize();
   }
-  if (Symbol *sym = ctx.symtab.findUnderscore("_load_config_used")) {
-    if (auto *b = dyn_cast<DefinedRegular>(sym)) {
-      SectionChunk *sc = b->getChunk();
-      assert(b->getRVA() >= sc->getRVA());
-      uint64_t offsetInChunk = b->getRVA() - sc->getRVA();
-      if (!sc->hasData || offsetInChunk + 4 > sc->getSize())
-        Fatal(ctx) << "_load_config_used is malformed";
-
-      ArrayRef<uint8_t> secContents = sc->getContents();
-      uint32_t loadConfigSize =
-          *reinterpret_cast<const ulittle32_t *>(&secContents[offsetInChunk]);
-      if (offsetInChunk + loadConfigSize > sc->getSize())
-        Fatal(ctx) << "_load_config_used is too large";
-      dir[LOAD_CONFIG_TABLE].RelativeVirtualAddress = b->getRVA();
-      dir[LOAD_CONFIG_TABLE].Size = loadConfigSize;
-    }
+  if (ctx.symtab.loadConfigSym) {
+    dir[LOAD_CONFIG_TABLE].RelativeVirtualAddress =
+        ctx.symtab.loadConfigSym->getRVA();
+    dir[LOAD_CONFIG_TABLE].Size = ctx.symtab.loadConfigSize;
   }
   if (!delayIdata.empty()) {
     dir[DELAY_IMPORT_DESCRIPTOR].RelativeVirtualAddress =
@@ -2649,31 +2637,14 @@ void Writer::fixTlsAlignment() {
 }
 
 void Writer::prepareLoadConfig() {
-  Symbol *sym = ctx.symtab.findUnderscore("_load_config_used");
-  auto *b = cast_if_present<DefinedRegular>(sym);
-  if (!b) {
-    if (ctx.config.guardCF != GuardCFLevel::Off)
-      Warn(ctx)
-          << "Control Flow Guard is enabled but '_load_config_used' is missing";
-    if (ctx.config.dependentLoadFlags)
-      Warn(ctx) << "_load_config_used not found, /dependentloadflag will have "
-                   "no effect";
+  if (!ctx.symtab.loadConfigSym)
     return;
-  }
 
-  OutputSection *sec = ctx.getOutputSection(b->getChunk());
-  uint8_t *buf = buffer->getBufferStart();
-  uint8_t *secBuf = buf + sec->getFileOff();
-  uint8_t *symBuf = secBuf + (b->getRVA() - sec->getRVA());
-  uint32_t expectedAlign = ctx.config.is64() ? 8 : 4;
-  if (b->getChunk()->getAlignment() < expectedAlign)
-    Warn(ctx) << "'_load_config_used' is misaligned (expected alignment to be "
-              << expectedAlign << " bytes, got "
-              << b->getChunk()->getAlignment() << " instead)";
-  else if (!isAligned(Align(expectedAlign), b->getRVA()))
-    Warn(ctx) << "'_load_config_used' is misaligned (RVA is 0x"
-              << Twine::utohexstr(b->getRVA()) << " not aligned to "
-              << expectedAlign << " bytes)";
+  OutputSection *sec =
+      ctx.getOutputSection(ctx.symtab.loadConfigSym->getChunk());
+  uint8_t *secBuf = buffer->getBufferStart() + sec->getFileOff();
+  uint8_t *symBuf =
+      secBuf + (ctx.symtab.loadConfigSym->getRVA() - sec->getRVA());
 
   if (ctx.config.is64())
     prepareLoadConfig(reinterpret_cast<coff_load_configuration64 *>(symBuf));
diff --git a/lld/test/COFF/guard-warnings.s b/lld/test/COFF/guard-warnings.s
@@ -38,7 +38,7 @@
 
 # RUN: llvm-mc -triple x86_64-windows-msvc %t/loadcfg-misaligned2.s -filetype=obj -o %t/loadcfg-misaligned2.obj
 # RUN: lld-link %t/main.obj %t/loadcfg-misaligned2.obj -guard:cf,longjmp,ehcont -out:%t-misaligned2.exe -entry:main %basename_t-exp.lib 2>&1 | FileCheck %s --check-prefix=WARN_ALIGN2
-# WARN_ALIGN2: warning: '_load_config_used' is misaligned (RVA is 0x{{[0-9A-F]*}}2 not aligned to 8 bytes)
+# WARN_ALIGN2: warning: '_load_config_used' is misaligned (section offset is 0x{{[0-9A-F]*}}2 not aligned to 8 bytes)
 
 # RUN: llvm-mc -triple x86_64-windows-msvc %t/loadcfg-full.s -filetype=obj -o %t/loadcfg-full.obj
 # RUN: lld-link %t/main.obj %t/loadcfg-full.obj -guard:cf,longjmp,ehcont -out:%t.exe -entry:main %basename_t-exp.lib 2>&1 | FileCheck %s --check-prefix=NOWARN --allow-empty
diff --git a/lld/test/COFF/loadcfg-short.test b/lld/test/COFF/loadcfg-short.test
@@ -0,0 +1,33 @@
+# RUN: yaml2obj %s -o %t.obj
+# RUN: not lld-link -out:%t.dll %t.obj -dll -noentry 2>&1 | FileCheck %s
+# CHECK: lld-link: error: _load_config_used section chunk is too small
+
+--- !COFF
+header:
+  Machine:         IMAGE_FILE_MACHINE_AMD64
+  Characteristics: []
+sections:
+  - Name:            .rdata
+    Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ]
+    Alignment:       16
+    SectionData:     '030000'
+symbols:
+  - Name:            .rdata
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+    SectionDefinition:
+      Length:          112
+      NumberOfRelocations: 0
+      NumberOfLinenumbers: 0
+      CheckSum:        0
+      Number:          3
+  - Name:            _load_config_used
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+...
diff --git a/lld/test/COFF/loadcfg-size.test b/lld/test/COFF/loadcfg-size.test
@@ -0,0 +1,33 @@
+# RUN: yaml2obj %s -o %t.obj
+# RUN: not lld-link -out:%t.dll %t.obj -dll -noentry 2>&1 | FileCheck %s
+# CHECK: lld-link: error: _load_config_used specifies a size larger than its containing section chunk
+
+--- !COFF
+header:
+  Machine:         IMAGE_FILE_MACHINE_AMD64
+  Characteristics: []
+sections:
+  - Name:            .rdata
+    Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ]
+    Alignment:       16
+    SectionData:     '0c00000000000000'
+symbols:
+  - Name:            .rdata
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+    SectionDefinition:
+      Length:          112
+      NumberOfRelocations: 0
+      NumberOfLinenumbers: 0
+      CheckSum:        0
+      Number:          3
+  - Name:            _load_config_used
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+...
diff --git a/lld/test/COFF/loadcfg-uninitialized.test b/lld/test/COFF/loadcfg-uninitialized.test
@@ -0,0 +1,33 @@
+# RUN: yaml2obj %s -o %t.obj
+# RUN: not lld-link -out:%t.dll %t.obj -dll -noentry 2>&1 | FileCheck %s
+# CHECK: lld-link: error: _load_config_used points to uninitialized data
+
+--- !COFF
+header:
+  Machine:         IMAGE_FILE_MACHINE_AMD64
+  Characteristics: []
+sections:
+  - Name:            .rdata
+    Characteristics: [ IMAGE_SCN_CNT_UNINITIALIZED_DATA, IMAGE_SCN_MEM_READ ]
+    Alignment:       16
+    VirtualSize:     0x140
+symbols:
+  - Name:            .rdata
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+    SectionDefinition:
+      Length:          112
+      NumberOfRelocations: 0
+      NumberOfLinenumbers: 0
+      CheckSum:        0
+      Number:          3
+  - Name:            _load_config_used
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+...
