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

[mstorsjo]

Summary

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.

Alternatively, just disallow the replacements of these symbols (implemented in D48952).

[mstorsjo]

Fixed a case where ImpSym was assumed to be non-null, using cast_or_null instead.

[ruiu]

LGTM

[mstorsjo]

In D48953#1153297, @ruiu wrote:

LGTM

Thanks! Just to be sure, did you look at D48952 as well, but prefer this alternative?

[mstorsjo]

In D48953#1153940, @mstorsjo wrote:

In D48953#1153297, @ruiu wrote:

LGTM

Thanks! Just to be sure, did you look at D48952 as well, but prefer this alternative?

Ping @ruiu, just a quick check if you looked at both alternatives to this solution, before I'll commit it.

[ruiu]

Sorry for the belated response. Yes, I did read the other patch too, and I prefer this one because this patch doesn't change the code for symbol name resolution.

[mstorsjo]

In D48953#1156528, @ruiu wrote:

Sorry for the belated response. Yes, I did read the other patch too, and I prefer this one because this patch doesn't change the code for symbol name resolution.

Ok, thanks for confirming and clarifying! I'll land this one soon then.

COFF/InputFiles.h

 class ImportFile : public InputFile {
 public:
   explicit ImportFile(MemoryBufferRef M) : InputFile(ImportKind, M) {}
 
   static bool classof(const InputFile *F) { return F->kind() == ImportKind; }
 
   static std::vector<ImportFile *> Instances;
 
-  DefinedImportData *ImpSym = nullptr;
-  DefinedImportThunk *ThunkSym = nullptr;
+  Symbol *ImpSym = nullptr;
+  Symbol *ThunkSym = nullptr;
   std::string DLLName;
 
 private:
   void parse() override;
 
 public:
   StringRef ExternalName;
   const coff_import_header *Hdr;

COFF/InputFiles.cpp

void ImportFile::parse() {
 
   if (Hdr->getType() == llvm::COFF::IMPORT_CONST)
     static_cast<void>(Symtab->addImportData(Name, this));
 
   // If type is function, we need to create a thunk which jump to an
   // 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() {
   Obj = check(lto::InputFile::create(MemoryBufferRef(
       MB.getBuffer(), Saver.save(ParentName + MB.getBufferIdentifier()))));
   std::vector<std::pair<Symbol *, bool>> Comdat(Obj->getComdatTable().size());
   for (size_t I = 0; I != Obj->getComdatTable().size(); ++I)
     Comdat[I] = Symtab->addComdat(this, Saver.save(Obj->getComdatTable()[I]));

COFF/SymbolTable.h

Symbol *addRegular(InputFile *F, StringRef N,
                      const llvm::object::coff_symbol_generic *S = nullptr,
                      SectionChunk *C = nullptr);
   std::pair<Symbol *, bool>
   addComdat(InputFile *F, StringRef N,
             const llvm::object::coff_symbol_generic *S = nullptr);
   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,
 
   void reportDuplicate(Symbol *Existing, InputFile *NewFile);
 
   // A list of chunks which to be added to .rdata.
   std::vector<Chunk *> LocalImportChunks;
 
   // Iterates symbols in non-determinstic hash table order.
   template <typename T> void forEachSymbol(T Callback) {

COFF/SymbolTable.cpp

Symbol *SymbolTable::addCommon(InputFile *F, StringRef N, uint64_t Size,
   if (WasInserted || !isa<DefinedCOFF>(S))
     replaceSymbol<DefinedCommon>(S, F, N, Size, Sym, C);
   else if (auto *DC = dyn_cast<DefinedCommon>(S))
     if (Size > DC->getSize())
       replaceSymbol<DefinedCommon>(S, F, N, Size, Sym, C);
   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,
   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);
   return nullptr;
 }
 
 std::vector<Chunk *> SymbolTable::getChunks() {
   std::vector<Chunk *> Res;

COFF/Writer.cpp

void Writer::createImportTables() {
   for (ImportFile *File : ImportFile::Instances) {
     if (!File->Live)
       continue;
 
     std::string DLL = StringRef(File->DLLName).lower();
     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);
     }
   }
 
   if (!Idata.empty())
     for (Chunk *C : Idata.getChunks())
       IdataSec->addChunk(C);
 
   if (!DelayIdata.empty()) {

test/COFF/Inputs/otherFunc.s

+.global otherFunc
+.global MessageBoxA
+.text
+otherFunc:
+  ret
+MessageBoxA:
+  ret

test/COFF/thunk-replace.s

+# 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