[COFF] Store import symbol pointers as pointers to the base class

Future symbol insertions can potentially change the type of these symbols - keep pointers to the base class to reflect this, and use dynamic casts to inspect them before using as the subclass type. This fixes crashes that were possible before, by touching these symbols that now are populated as e.g. a DefinedRegular, via the old pointers with DefinedImportThunk type. Differential Revision: https://reviews.llvm.org/D48953 llvm-svn: 336652
llvm · Jul 10, 2018 · 474be00 · 474be00
1 parent bc6702a
commit 474be00
Show file tree

Hide file tree

Showing 7 changed files with 45 additions and 16 deletions.
diff --git a/lld/COFF/InputFiles.cpp b/lld/COFF/InputFiles.cpp
@@ -431,7 +431,8 @@ void ImportFile::parse() {
   // address pointed by the __imp_ symbol. (This allows you to call
   // DLL functions just like regular non-DLL functions.)
   if (Hdr->getType() == llvm::COFF::IMPORT_CODE)
-    ThunkSym = Symtab->addImportThunk(Name, ImpSym, Hdr->Machine);
+    ThunkSym = Symtab->addImportThunk(
+        Name, cast_or_null<DefinedImportData>(ImpSym), Hdr->Machine);
 }
 
 void BitcodeFile::parse() {

diff --git a/lld/COFF/InputFiles.h b/lld/COFF/InputFiles.h
@@ -207,8 +207,8 @@ class ImportFile : public InputFile {
 
   static std::vector<ImportFile *> Instances;
 
-  DefinedImportData *ImpSym = nullptr;
-  DefinedImportThunk *ThunkSym = nullptr;
+  Symbol *ImpSym = nullptr;
+  Symbol *ThunkSym = nullptr;
   std::string DLLName;
 
 private:

diff --git a/lld/COFF/SymbolTable.cpp b/lld/COFF/SymbolTable.cpp
@@ -342,30 +342,29 @@ Symbol *SymbolTable::addCommon(InputFile *F, StringRef N, uint64_t Size,
   return S;
 }
 
-DefinedImportData *SymbolTable::addImportData(StringRef N, ImportFile *F) {
+Symbol *SymbolTable::addImportData(StringRef N, ImportFile *F) {
   Symbol *S;
   bool WasInserted;
   std::tie(S, WasInserted) = insert(N);
   S->IsUsedInRegularObj = true;
   if (WasInserted || isa<Undefined>(S) || isa<Lazy>(S)) {
     replaceSymbol<DefinedImportData>(S, N, F);
-    return cast<DefinedImportData>(S);
+    return S;
   }
 
   reportDuplicate(S, F);
   return nullptr;
 }
 
-DefinedImportThunk *SymbolTable::addImportThunk(StringRef Name,
-                                               DefinedImportData *ID,
-                                               uint16_t Machine) {
+Symbol *SymbolTable::addImportThunk(StringRef Name, DefinedImportData *ID,
+                                    uint16_t Machine) {
   Symbol *S;
   bool WasInserted;
   std::tie(S, WasInserted) = insert(Name);
   S->IsUsedInRegularObj = true;
   if (WasInserted || isa<Undefined>(S) || isa<Lazy>(S)) {
     replaceSymbol<DefinedImportThunk>(S, Name, ID, Machine);
-    return cast<DefinedImportThunk>(S);
+    return S;
   }
 
   reportDuplicate(S, ID->File);

diff --git a/lld/COFF/SymbolTable.h b/lld/COFF/SymbolTable.h
@@ -92,9 +92,9 @@ class SymbolTable {
   Symbol *addCommon(InputFile *F, StringRef N, uint64_t Size,
                     const llvm::object::coff_symbol_generic *S = nullptr,
                     CommonChunk *C = nullptr);
-  DefinedImportData *addImportData(StringRef N, ImportFile *F);
-  DefinedImportThunk *addImportThunk(StringRef Name, DefinedImportData *S,
-                                     uint16_t Machine);
+  Symbol *addImportData(StringRef N, ImportFile *F);
+  Symbol *addImportThunk(StringRef Name, DefinedImportData *S,
+                         uint16_t Machine);
 
   void reportDuplicate(Symbol *Existing, InputFile *NewFile);
 

diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
@@ -544,17 +544,24 @@ void Writer::createImportTables() {
     if (Config->DLLOrder.count(DLL) == 0)
       Config->DLLOrder[DLL] = Config->DLLOrder.size();
 
-    if (DefinedImportThunk *Thunk = File->ThunkSym)
+    if (File->ThunkSym) {
+      if (!isa<DefinedImportThunk>(File->ThunkSym))
+        fatal(toString(*File->ThunkSym) + " was replaced");
+      DefinedImportThunk *Thunk = cast<DefinedImportThunk>(File->ThunkSym);
       if (File->ThunkLive)
         TextSec->addChunk(Thunk->getChunk());
+    }
 
+    if (File->ImpSym && !isa<DefinedImportData>(File->ImpSym))
+      fatal(toString(*File->ImpSym) + " was replaced");
+    DefinedImportData *ImpSym = cast_or_null<DefinedImportData>(File->ImpSym);
     if (Config->DelayLoads.count(StringRef(File->DLLName).lower())) {
       if (!File->ThunkSym)
         fatal("cannot delay-load " + toString(File) +
-              " due to import of data: " + toString(*File->ImpSym));
-      DelayIdata.add(File->ImpSym);
+              " due to import of data: " + toString(*ImpSym));
+      DelayIdata.add(ImpSym);
     } else {
-      Idata.add(File->ImpSym);
+      Idata.add(ImpSym);
     }
   }
 

diff --git a/lld/test/COFF/Inputs/otherFunc.s b/lld/test/COFF/Inputs/otherFunc.s
@@ -0,0 +1,7 @@
+.global otherFunc
+.global MessageBoxA
+.text
+otherFunc:
+  ret
+MessageBoxA:
+  ret
diff --git a/lld/test/COFF/thunk-replace.s b/lld/test/COFF/thunk-replace.s
@@ -0,0 +1,15 @@
+# REQUIRES: x86
+
+# RUN: llvm-mc -triple=x86_64-win32 %s -filetype=obj -o %t.main.obj
+# RUN: llvm-mc -triple=x86_64-win32 %p/Inputs/otherFunc.s -filetype=obj -o %t.other.obj
+# RUN: llvm-ar rcs %t.other.lib %t.other.obj
+# RUN: not lld-link -out:%t.exe -entry:main %t.main.obj %p/Inputs/std64.lib %t.other.lib -opt:noref 2>&1 | FileCheck %s
+# CHECK: MessageBoxA was replaced
+
+.global main
+.text
+main:
+  callq MessageBoxA
+  callq ExitProcess
+  callq otherFunc
+  ret