[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

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() {

lld/COFF/InputFiles.h

@@ -207,8 +207,8 @@ public:
 
   static std::vector<ImportFile *> Instances;
 
-  DefinedImportData *ImpSym = nullptr;
-  DefinedImportThunk *ThunkSym = nullptr;
+  Symbol *ImpSym = nullptr;
+  Symbol *ThunkSym = nullptr;
   std::string DLLName;
 
 private:

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);

lld/COFF/SymbolTable.h

@@ -92,9 +92,9 @@ public:
   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);
 

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);
     }
   }
 

lld/test/COFF/Inputs/otherFunc.s

@@ -0,0 +1,7 @@
+.global otherFunc
+.global MessageBoxA
+.text
+otherFunc:
+  ret
+MessageBoxA:
+  ret

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