[LLD][COFF] Don't resolve weak aliases when performing local import

[cjacek]

Fixes crashes reported in #151255.

The alias may have already been stored for later resolution, which can lead to treating a resolved alias as if it were still undefined. Instead, use the alias target directly for the import.

Also extended the test to make reproducing the problem more likely, and added an assert that catches the issue.

[llvmbot]

@llvm/pr-subscribers-platform-windows
@llvm/pr-subscribers-lld

@llvm/pr-subscribers-lld-coff

Author: Jacek Caban (cjacek)

Changes

Fixes crashes reported in #151255.

The alias may have already been stored for later resolution, which can lead to treating a resolved alias as if it were still undefined. Instead, use the alias target directly for the import.

Also extended the test to make reproducing the problem more likely, and added an assert that catches the issue.

---

Full diff: https://github.com/llvm/llvm-project/pull/152000.diff

3 Files Affected:

• (modified) lld/COFF/Driver.cpp (+3-1)
• (modified) lld/COFF/SymbolTable.cpp (+5-12)
• (modified) lld/test/COFF/import_weak_alias.test (+10-3)

diff --git a/lld/COFF/Driver.cpp b/lld/COFF/Driver.cpp
index 7580b469e4f87..ff616d0ce2bff 100644
--- a/lld/COFF/Driver.cpp
+++ b/lld/COFF/Driver.cpp
@@ -2714,8 +2714,10 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   }
 
   ctx.symtab.initializeSameAddressThunks();
-  for (auto alias : aliases)
+  for (auto alias : aliases) {
+    assert(alias->kind() == Symbol::UndefinedKind);
     alias->resolveWeakAlias();
+  }
 
   if (config->mingw) {
     // Make sure the crtend.o object is the last object file. This object
diff --git a/lld/COFF/SymbolTable.cpp b/lld/COFF/SymbolTable.cpp
index de04cdff6483d..e11d2c6dac83e 100644
--- a/lld/COFF/SymbolTable.cpp
+++ b/lld/COFF/SymbolTable.cpp
@@ -478,17 +478,11 @@ void SymbolTable::resolveRemainingUndefines(std::vector<Undefined *> &aliases) {
     if (name.starts_with("__imp_")) {
       auto findLocalSym = [&](StringRef n) {
         Symbol *sym = find(n);
-        if (auto undef = dyn_cast_or_null<Undefined>(sym)) {
-          // The unprefixed symbol might come later in symMap, so handle it now
-          // if needed.
-          if (!undef->resolveWeakAlias())
-            sym = nullptr;
-        }
-        return sym;
+        return sym ? sym->getDefined() : nullptr;
       };
 
       StringRef impName = name.substr(strlen("__imp_"));
-      Symbol *imp = findLocalSym(impName);
+      Defined *imp = findLocalSym(impName);
       if (!imp && isEC()) {
         // Try to use the mangled symbol on ARM64EC.
         std::optional<std::string> mangledName =
@@ -502,11 +496,10 @@ void SymbolTable::resolveRemainingUndefines(std::vector<Undefined *> &aliases) {
             imp = findLocalSym(*mangledName);
         }
       }
-      if (imp && isa<Defined>(imp)) {
-        auto *d = cast<Defined>(imp);
-        replaceSymbol<DefinedLocalImport>(sym, ctx, name, d);
+      if (imp) {
+        replaceSymbol<DefinedLocalImport>(sym, ctx, name, imp);
         localImportChunks.push_back(cast<DefinedLocalImport>(sym)->getChunk());
-        localImports[sym] = d;
+        localImports[sym] = imp;
         continue;
       }
     }
diff --git a/lld/test/COFF/import_weak_alias.test b/lld/test/COFF/import_weak_alias.test
index ae1817c67a20a..ff07022741d1c 100644
--- a/lld/test/COFF/import_weak_alias.test
+++ b/lld/test/COFF/import_weak_alias.test
@@ -4,17 +4,24 @@
 # RUN: llvm-mc --filetype=obj -triple=x86_64-windows-msvc %t.dir/foo.s -o %t.foo.obj
 # RUN: llvm-mc --filetype=obj -triple=x86_64-windows-msvc %t.dir/qux.s -o %t.qux.obj
 # RUN: lld-link %t.qux.obj %t.foo.obj -out:%t.dll -dll
+# RUN: lld-link %t.foo.obj %t.qux.obj -out:%t.dll -dll
 #
 #--- foo.s
 .text
 bar:
   ret
 
-.weak foo
-.set foo, bar
+.weak a
+.weak b
+.weak c
+.set a, bar
+.set b, bar
+.set c, bar
 #--- qux.s
 .text
 .global _DllMainCRTStartup
 _DllMainCRTStartup:
-  call *__imp_foo(%rip)
+  call *__imp_a(%rip)
+  call *__imp_b(%rip)
+  call *__imp_c(%rip)
   ret

[mstorsjo]

LGTM