[COFF] Fix /export:foo=bar when bar is a weak alias

[rnk]

Summary

When handling exports from the command line or from .def files, the
linker does a "fuzzy" string lookup to allow finding mangled symbols.
However, when the symbol is re-exported under a new name, the linker has
to transfer the decorations from the exported symbol over to the new
name. This is implemented by taking the mangled symbol that was found in
the object and replacing the original symbol name with the export name.

Before this patch, LLD implemented the fuzzy search by adding an
undefined symbol with the unmangled name, and then during symbol
resolution, checking if similar mangled symbols had been added after the
last round of symbol resolution. If so, LLD makes the original symbol a
weak alias of the mangled symbol. Later, to get the original symbol
name, LLD would look through the weak alias and forward it on to the
import library writer, which copies the symbol decorations. This
approach doesn't work when bar is itself a weak alias, as is the case in
asan. It's especially bad when the aliasee of bar contains the string
"bar", consider "bar_default". In this case, we would end up exporting
the symbol "foo_default" when we should've exported just "foo".

To fix this, don't look through weak aliases to find the mangled name.
Save the mangled name earlier during fuzzy symbol lookup.

Fixes PR42074

[mstorsjo]

Looks good and thought through. I had a few nits though.

lld/COFF/SymbolTable.cpp

475

This regex is a little misleading/incomplete, as it can also match by skipping the first char of Prefix. (I guess that can't be expressed in regex syntax easily...)

lld/test/COFF/entry-mangled.test

5

Is this -verbose needed? I see no checks for its contents.

[ruiu]

LGTM

lld/COFF/SymbolTable.cpp

494

I wondered why you needed to call getSymWithPrefix first instead of scanning the entire symbol table. Looks like this is just a performance optimization; find candidates first and then try to find a match from them. Can you add a comment about that?

llvm/include/llvm/Object/COFFImportFile.h

72–87

Thank you for doing this!

[rnk]

Thanks!

lld/test/COFF/entry-mangled.test

5

Oops, it was for debugging.

llvm/include/llvm/Object/COFFImportFile.h

72–87

Hopefully it's correct. :)

lld/COFF/Driver.h

private:
   // We don't want to add the same file more than once.
   // Files are uniquified by their filesystem and file number.
   std::set<llvm::sys::fs::UniqueID> VisitedFiles;
 
   std::set<std::string> VisitedLibs;
 
   Symbol *addUndefined(StringRef Sym);
 
+  StringRef mangleMaybe(Symbol *S);
+
   // Windows specific -- "main" is not the only main function in Windows.
   // You can choose one from these four -- {w,}{WinMain,main}.
   // There are four different entry point functions for them,
   // {w,}{WinMain,main}CRTStartup, respectively. The linker needs to
   // choose the right one depending on which "main" function is defined.
   // This function looks up the symbol table and resolve corresponding
   // entry point name.
   StringRef findDefaultEntry();

lld/COFF/Driver.cpp

 static StringRef mangle(StringRef Sym) {
   assert(Config->Machine != IMAGE_FILE_MACHINE_UNKNOWN);
   if (Config->Machine == I386)
     return Saver.save("_" + Sym);
   return Sym;
 }
 
 static bool findUnderscoreMangle(StringRef Sym) {
-  StringRef Entry = Symtab->findMangle(mangle(Sym));
-  return !Entry.empty() && !isa<Undefined>(Symtab->find(Entry));
+  Symbol *S = Symtab->findMangle(mangle(Sym));
+  return S && !isa<Undefined>(S);
 }
 
 MemoryBufferRef LinkerDriver::takeBuffer(std::unique_ptr<MemoryBuffer> MB) {
   MemoryBufferRef MBRef = *MB;
   make<std::unique_ptr<MemoryBuffer>>(std::move(MB)); // take ownership
 
   if (Driver->Tar)
     Driver->Tar->append(relativeToRoot(MBRef.getBufferIdentifier()),
Symbol *LinkerDriver::addUndefined(StringRef Name) {
   Symbol *B = Symtab->addUndefined(Name);
   if (!B->IsGCRoot) {
     B->IsGCRoot = true;
     Config->GCRoot.push_back(B);
   }
   return B;
 }
 
+StringRef LinkerDriver::mangleMaybe(Symbol *S) {
+  // If the plain symbol name has already been resolved, do nothing.
+  Undefined *Unmangled = dyn_cast<Undefined>(S);
+  if (!Unmangled)
+    return "";
+
+  // Otherwise, see if a similar, mangled symbol exists in the symbol table.
+  Symbol *Mangled = Symtab->findMangle(Unmangled->getName());
+  if (!Mangled)
+    return "";
+
+  // If we find a similar mangled symbol, make this an alias to it and return
+  // its name.
+  log(Unmangled->getName() + " aliased to " + Mangled->getName());
+  Unmangled->WeakAlias = Symtab->addUndefined(Mangled->getName());
+  return Mangled->getName();
+}
+
 // Windows specific -- find default entry point name.
 //
 // There are four different entry point functions for Windows executables,
 // each of which corresponds to a user-defined "main" function. This function
 // infers an entry point from a user-defined "main" function.
 StringRef LinkerDriver::findDefaultEntry() {
   assert(Config->Subsystem != IMAGE_SUBSYSTEM_UNKNOWN &&
          "must handle /subsystem before calling this");
void LinkerDriver::link(ArrayRef<const char *> ArgsArr) {
 
   // This code may add new undefined symbols to the link, which may enqueue more
   // symbol resolution tasks, so we need to continue executing tasks until we
   // converge.
   do {
     // Windows specific -- if entry point is not found,
     // search for its mangled names.
     if (Config->Entry)
-      Symtab->mangleMaybe(Config->Entry);
+      mangleMaybe(Config->Entry);
 
     // Windows specific -- Make sure we resolve all dllexported symbols.
     for (Export &E : Config->Exports) {
       if (!E.ForwardTo.empty())
         continue;
       E.Sym = addUndefined(E.Name);
       if (!E.Directives)
-        Symtab->mangleMaybe(E.Sym);
+        E.SymbolName = mangleMaybe(E.Sym);
     }
 
     // Add weak aliases. Weak aliases is a mechanism to give remaining
     // undefined symbols final chance to be resolved successfully.
     for (auto Pair : Config->AlternateNames) {
       StringRef From = Pair.first;
       StringRef To = Pair.second;
       Symbol *Sym = Symtab->find(From);

lld/COFF/DriverUtils.cpp

void fixupExports() {
   for (Export &E : Config->Exports) {
     if (E.Ordinal == 0)
       continue;
     if (!Ords.insert(E.Ordinal).second)
       fatal("duplicate export ordinal: " + E.Name);
   }
 
   for (Export &E : Config->Exports) {
-    Symbol *Sym = E.Sym;
-    if (!E.ForwardTo.empty() || !Sym) {
-      E.SymbolName = E.Name;
-    } else {
-      if (auto *U = dyn_cast<Undefined>(Sym))
-        if (U->WeakAlias)
-          Sym = U->WeakAlias;
-      E.SymbolName = Sym->getName();
-    }
-  }
-
-  for (Export &E : Config->Exports) {
     if (!E.ForwardTo.empty()) {
       E.ExportName = undecorate(E.Name);
     } else {
       E.ExportName = undecorate(E.ExtName.empty() ? E.Name : E.ExtName);
     }
   }
 
   if (Config->KillAt && Config->Machine == I386) {

lld/COFF/SymbolTable.h

public:
   // Returns a symbol for a given name. Returns a nullptr if not found.
   Symbol *find(StringRef Name);
   Symbol *findUnderscore(StringRef Name);
 
   // Occasionally we have to resolve an undefined symbol to its
   // mangled symbol. This function tries to find a mangled name
   // for U from the symbol table, and if found, set the symbol as
   // a weak alias for U.
-  void mangleMaybe(Symbol *B);
-  StringRef findMangle(StringRef Name);
+  Symbol *findMangle(StringRef Name);
 
   // Build a set of COFF objects representing the combined contents of
   // BitcodeFiles and add them to the symbol table. Called after all files are
   // added and before the writer writes results to a file.
   void addCombinedLTOObjects();
   std::vector<StringRef> compileBitcodeFiles();
 
   // Creates an Undefined symbol for a given name.
for (auto &Pair : SymMap)
       Callback(Pair.second);
   }
 
 private:
   /// Inserts symbol if not already present.
   std::pair<Symbol *, bool> insert(StringRef Name);
   /// Same as insert(Name), but also sets IsUsedInRegularObj.
   std::pair<Symbol *, bool> insert(StringRef Name, InputFile *F);
-  StringRef findByPrefix(StringRef Prefix);
+
+  std::vector<Symbol *> getSymsWithPrefix(StringRef Prefix);
 
   llvm::DenseMap<llvm::CachedHashStringRef, Symbol *> SymMap;
   std::unique_ptr<BitcodeCompiler> LTO;
 };
 
 extern SymbolTable *Symtab;
 
 std::string getSymbolLocations(ObjFile *File, uint32_t SymIndex);
 
 } // namespace coff
 } // namespace lld
 
 #endif

lld/COFF/SymbolTable.cpp

 }
 
 Symbol *SymbolTable::findUnderscore(StringRef Name) {
   if (Config->Machine == I386)
     return find(("_" + Name).str());
   return find(Name);
 }
 
-StringRef SymbolTable::findByPrefix(StringRef Prefix) {
+// Return all symbols matching '.?${Prefix}'.
+std::vector<Symbol *> SymbolTable::getSymsWithPrefix(StringRef Prefix) {
+  std::vector<Symbol *> Syms;
+
   for (auto Pair : SymMap) {
     StringRef Name = Pair.first.val();
-    if (Name.startswith(Prefix))
-      return Name;
+    if (Name.startswith(Prefix) || Name.startswith(Prefix.drop_front()) ||
+        Name.drop_front().startswith(Prefix) ||
+        Name.drop_front().startswith(Prefix.drop_front())) {
+      Syms.push_back(Pair.second);
+    }
   }
-  return "";
+  return Syms;
 }
 
-StringRef SymbolTable::findMangle(StringRef Name) {
+Symbol *SymbolTable::findMangle(StringRef Name) {
   if (Symbol *Sym = find(Name))
     if (!isa<Undefined>(Sym))
-      return Name;
+      return Sym;
+  std::vector<Symbol *> Syms = getSymsWithPrefix(Name);
+  auto FindByPrefix = [&Syms](const Twine &T) -> Symbol * {
+    std::string Prefix = T.str();
+    for (auto *S : Syms)
+      if (S->getName().startswith(Prefix))
+        return S;
+    return nullptr;
+  };
+
+  // For non-x86, just look for C++ functions.
   if (Config->Machine != I386)
-    return findByPrefix(("?" + Name + "@@Y").str());
+    return FindByPrefix("?" + Name + "@@Y");
+
   if (!Name.startswith("_"))
-    return "";
+    return nullptr;
   // Search for x86 stdcall function.
-  StringRef S = findByPrefix((Name + "@").str());
-  if (!S.empty())
+  if (Symbol *S = FindByPrefix(Name + "@"))
     return S;
   // Search for x86 fastcall function.
-  S = findByPrefix(("@" + Name.substr(1) + "@").str());
-  if (!S.empty())
+  if (Symbol *S = FindByPrefix("@" + Name.substr(1) + "@"))
     return S;
   // Search for x86 vectorcall function.
-  S = findByPrefix((Name.substr(1) + "@@").str());
-  if (!S.empty())
+  if (Symbol *S = FindByPrefix(Name.substr(1) + "@@"))
     return S;
   // Search for x86 C++ non-member function.
-  return findByPrefix(("?" + Name.substr(1) + "@@Y").str());
-}
-
-void SymbolTable::mangleMaybe(Symbol *B) {
-  auto *U = dyn_cast<Undefined>(B);
-  if (!U || U->WeakAlias)
-    return;
-  StringRef Alias = findMangle(U->getName());
-  if (!Alias.empty()) {
-    log(U->getName() + " aliased to " + Alias);
-    U->WeakAlias = addUndefined(Alias);
-  }
+  return FindByPrefix("?" + Name.substr(1) + "@@Y");
 }
 
 Symbol *SymbolTable::addUndefined(StringRef Name) {
   return addUndefined(Name, nullptr, false);
 }
 
 std::vector<StringRef> SymbolTable::compileBitcodeFiles() {
   LTO.reset(new BitcodeCompiler);

lld/test/COFF/entry-mangled.test

 # REQUIRES: x86
 # RUN: yaml2obj < %s > %t.obj
 # RUN: lld-link /out:%t.exe /entry:main %t.obj
 # RUN: llvm-as -o %t.lto.obj %S/Inputs/entry-mangled.ll
-# RUN: lld-link /out:%t.exe /entry:main %t.lto.obj
+# RUN: lld-link /out:%t.exe /entry:main %t.lto.obj -verbose
 
 --- !COFF
 header:
   Machine:         IMAGE_FILE_MACHINE_AMD64
   Characteristics: []
 sections:
   - Name:            .text
     Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ]

lld/test/COFF/export-stdcall.s

+# RUN: llvm-mc -triple i686-windows-msvc %s -o %t.obj -filetype=obj
+# RUN: lld-link %t.obj -out:%t.dll -dll -nodefaultlib -noentry -export:foo_std=bar_std -export:foo_fast=bar_fast
+# RUN: llvm-nm %t.lib | FileCheck %s
+
+# MSVC fudges the lookup of 'bar' to allow it to find the stdcall function
+# _bar_std@8, and then exports _foo_std@8. Same for fastcall and other mangling
+# schemes.
+
+# CHECK: export-stdcall.s.tmp.dll:
+# CHECK: 00000000 T @foo_fast@8
+# CHECK: 00000000 T __imp_@foo_fast@8
+
+# CHECK: export-stdcall.s.tmp.dll:
+# CHECK: 00000000 T __imp__foo_std@8
+# CHECK: 00000000 T _foo_std@8
+
+  .text
+  .def   _bar_std@8; .scl 2; .type  32; .endef
+  .globl  _bar_std@8
+_bar_std@8:
+  movl  8(%esp), %eax
+  movl  4(%esp), %ecx
+  leal  42(%ecx,%eax), %eax
+  retl  $8
+
+  .def   @bar_fast@8; .scl  2; .type  32; .endef
+  .globl  @bar_fast@8
+@bar_fast@8:
+  leal  42(%ecx,%eax), %eax
+  retl
+

lld/test/COFF/export-weak-alias.s

+# RUN: llvm-mc -triple x86_64-windows-msvc %s -o %t.obj -filetype=obj
+# RUN: lld-link %t.obj -out:%t.dll -dll -nodefaultlib -noentry
+# RUN: llvm-nm %t.lib | FileCheck %s
+
+# CHECK: export-weak-alias.s.tmp.dll:
+# CHECK: 00000000 T __imp_foo_dll{{$}}
+# CHECK: 00000000 T foo_dll{{$}}
+
+  .text
+  .def   @feat.00;
+  .scl  3;
+  .type 0;
+  .endef
+  .globl  @feat.00
+.set @feat.00, 0
+  .file "t.c"
+  .def   foo_def;
+  .scl  2;
+  .type 32;
+  .endef
+  .globl  foo_def                 # -- Begin function foo_def
+  .p2align  4, 0x90
+foo_def:                                # @foo_def
+# %bb.0:                                # %entry
+  movl  $42, %eax
+  retq
+                                        # -- End function
+  .section  .drectve,"yn"
+  .ascii  " /alternatename:foo=foo_def"
+  .ascii  " /export:foo_dll=foo"
+
+  .addrsig

llvm/include/llvm/Object/COFFImportFile.h

 
 private:
   bool isData() const {
     return getCOFFImportHeader()->getType() == COFF::IMPORT_DATA;
   }
 };
 
 struct COFFShortExport {
+  /// The name of the export as specified in the .def file or on the command
+  /// line, i.e. "foo" in "/EXPORT:foo", and "bar" in "/EXPORT:foo=bar". This
+  /// may lack mangling, such as underscore prefixing and stdcall suffixing.
   std::string Name;
+
+  /// The external, exported name. Only non-empty when export renaming is in
+  /// effect, i.e. "foo" in "/EXPORT:foo=bar".
   std::string ExtName;
+
+  /// The real, mangled symbol name from the object file. Given
+  /// "/export:foo=bar", this could be "_bar@8" if bar is stdcall.
   std::string SymbolName;
+
+  /// Creates a weak alias. This is the name of the weak aliasee. In a .def
+  /// file, this is "baz" in "EXPORTS\nfoo = bar == baz".
   std::string AliasTarget;
 
   uint16_t Ordinal = 0;
   bool Noname = false;
   bool Data = false;
   bool Private = false;
   bool Constant = false;
 
   friend bool operator==(const COFFShortExport &L, const COFFShortExport &R) {