[lld/mac] Fix issues around thin archives

- most importantly, fix a use-after-free when using thin archives,
  by putting the archive unique_ptr to the arena allocator. This
  ports D65565 to MachO

- correctly demangle symbol namess from archives in diagnostics

- add a test for thin archives -- it finds this UaF, but only when
  running it under asan (it also finds the demangling fix)

- make forceLoadArchive() use addFile() with a bool to have the archive
  loading code in fewer places. no behavior change; matches COFF port a
  bit better

Differential Revision: https://reviews.llvm.org/D92360

5 files changed, 80 insertions, 30 deletions

diff --git a/lld/MachO/Driver.cpp b/lld/MachO/Driver.cpp
index 295f2c412a7f..a02be0c9a67c 100644
--- a/lld/MachO/Driver.cpp
+++ b/lld/MachO/Driver.cpp
@@ -225,10 +225,12 @@ static std::vector<ArchiveMember> getArchiveMembers(MemoryBufferRef mb) {
   std::unique_ptr<Archive> file =
       CHECK(Archive::create(mb),
             mb.getBufferIdentifier() + ": failed to parse archive");
+  Archive *archive = file.get();
+  make<std::unique_ptr<Archive>>(std::move(file)); // take ownership
 
   std::vector<ArchiveMember> v;
   Error err = Error::success();
-  for (const Archive::Child &c : file->children(err)) {
+  for (const Archive::Child &c : archive->children(err)) {
     MemoryBufferRef mbref =
         CHECK(c.getMemoryBufferRef(),
               mb.getBufferIdentifier() +
@@ -246,17 +248,7 @@ static std::vector<ArchiveMember> getArchiveMembers(MemoryBufferRef mb) {
   return v;
 }
 
-static void forceLoadArchive(StringRef path) {
-  if (Optional<MemoryBufferRef> buffer = readFile(path)) {
-    for (const ArchiveMember &member : getArchiveMembers(*buffer)) {
-      auto file = make<ObjFile>(member.mbref, member.modTime);
-      file->archiveName = buffer->getBufferIdentifier();
-      inputFiles.push_back(file);
-    }
-  }
-}
-
-static InputFile *addFile(StringRef path) {
+static InputFile *addFile(StringRef path, bool forceLoadArchive) {
   Optional<MemoryBufferRef> buffer = readFile(path);
   if (!buffer)
     return nullptr;
@@ -271,8 +263,14 @@ static InputFile *addFile(StringRef path) {
     if (!file->isEmpty() && !file->hasSymbolTable())
       error(path + ": archive has no index; run ranlib to add one");
 
-    if (config->allLoad) {
-      forceLoadArchive(path);
+    if (config->allLoad || forceLoadArchive) {
+      if (Optional<MemoryBufferRef> buffer = readFile(path)) {
+        for (const ArchiveMember &member : getArchiveMembers(*buffer)) {
+          auto file = make<ObjFile>(member.mbref, member.modTime);
+          file->archiveName = buffer->getBufferIdentifier();
+          inputFiles.push_back(file);
+        }
+      }
     } else if (config->forceLoadObjC) {
       for (const object::Archive::Symbol &sym : file->symbols())
         if (sym.getName().startswith(objc::klass))
@@ -320,7 +318,7 @@ static void addFileList(StringRef path) {
     return;
   MemoryBufferRef mbref = *buffer;
   for (StringRef path : args::getLines(mbref))
-    addFile(path);
+    addFile(path, false);
 }
 
 static std::array<StringRef, 6> archNames{"arm",    "arm64", "i386",
@@ -671,10 +669,11 @@ bool macho::link(llvm::ArrayRef<const char *> argsArr, bool canExitEarly,
     // TODO: are any of these better handled via filtered() or getLastArg()?
     switch (opt.getID()) {
     case OPT_INPUT:
-      addFile(arg->getValue());
+      addFile(arg->getValue(), false);
       break;
     case OPT_weak_library: {
-      auto *dylibFile = dyn_cast_or_null<DylibFile>(addFile(arg->getValue()));
+      auto *dylibFile =
+          dyn_cast_or_null<DylibFile>(addFile(arg->getValue(), false));
       if (dylibFile)
         dylibFile->forceWeakImport = true;
       break;
@@ -683,13 +682,13 @@ bool macho::link(llvm::ArrayRef<const char *> argsArr, bool canExitEarly,
       addFileList(arg->getValue());
       break;
     case OPT_force_load:
-      forceLoadArchive(arg->getValue());
+      addFile(arg->getValue(), true);
       break;
     case OPT_l:
     case OPT_weak_l: {
       StringRef name = arg->getValue();
       if (Optional<std::string> path = findLibrary(name)) {
-        auto *dylibFile = dyn_cast_or_null<DylibFile>(addFile(*path));
+        auto *dylibFile = dyn_cast_or_null<DylibFile>(addFile(*path, false));
         if (opt.getID() == OPT_weak_l && dylibFile)
           dylibFile->forceWeakImport = true;
         break;
@@ -701,7 +700,7 @@ bool macho::link(llvm::ArrayRef<const char *> argsArr, bool canExitEarly,
     case OPT_weak_framework: {
       StringRef name = arg->getValue();
       if (Optional<std::string> path = findFramework(name)) {
-        auto *dylibFile = dyn_cast_or_null<DylibFile>(addFile(*path));
+        auto *dylibFile = dyn_cast_or_null<DylibFile>(addFile(*path, false));
         if (opt.getID() == OPT_weak_framework && dylibFile)
           dylibFile->forceWeakImport = true;
         break;
diff --git a/lld/MachO/InputFiles.cpp b/lld/MachO/InputFiles.cpp
index 8a451d0dc217..cfb30cfc35f5 100644
--- a/lld/MachO/InputFiles.cpp
+++ b/lld/MachO/InputFiles.cpp
@@ -584,7 +584,7 @@ void ArchiveFile::fetch(const object::Archive::Symbol &sym) {
   object::Archive::Child c =
       CHECK(sym.getMember(), toString(this) +
                                  ": could not get the member for symbol " +
-                                 sym.getName());
+                                 toMachOString(sym));
 
   if (!seen.insert(c.getChildOffset()).second)
     return;
@@ -593,13 +593,13 @@ void ArchiveFile::fetch(const object::Archive::Symbol &sym) {
       CHECK(c.getMemoryBufferRef(),
             toString(this) +
                 ": could not get the buffer for the member defining symbol " +
-                sym.getName());
+                toMachOString(sym));
 
   uint32_t modTime = toTimeT(
       CHECK(c.getLastModified(), toString(this) +
                                      ": could not get the modification time "
                                      "for the member defining symbol " +
-                                     sym.getName()));
+                                     toMachOString(sym)));
 
   auto file = make<ObjFile>(mb, modTime);
   file->archiveName = getName();
diff --git a/lld/MachO/Symbols.cpp b/lld/MachO/Symbols.cpp
index 87bbab00901f..4c83188fd259 100644
--- a/lld/MachO/Symbols.cpp
+++ b/lld/MachO/Symbols.cpp
@@ -14,6 +14,21 @@ using namespace llvm;
 using namespace lld;
 using namespace lld::macho;
 
+// Returns a symbol for an error message.
+static std::string demangle(StringRef symName) {
+  if (config->demangle)
+    return demangleItanium(symName);
+  return std::string(symName);
+}
+
+std::string lld::toString(const Symbol &sym) {
+  return demangle(sym.getName());
+}
+
+std::string lld::toMachOString(const object::Archive::Symbol &b) {
+  return demangle(b.getName());
+}
+
 uint64_t Defined::getVA() const {
   if (isAbsolute())
     return value;
@@ -31,13 +46,6 @@ uint64_t Defined::getFileOffset() const {
 
 void LazySymbol::fetchArchiveMember() { file->fetch(sym); }
 
-// Returns a symbol for an error message.
-std::string lld::toString(const Symbol &sym) {
-  if (config->demangle)
-    return demangleItanium(sym.getName());
-  return std::string(sym.getName());
-}
-
 uint64_t DSOHandle::getVA() const { return header->addr; }
 
 uint64_t DSOHandle::getFileOffset() const { return header->fileOff; }
diff --git a/lld/MachO/Symbols.h b/lld/MachO/Symbols.h
index f0d77d5ce0fe..671a1ec200fe 100644
--- a/lld/MachO/Symbols.h
+++ b/lld/MachO/Symbols.h
@@ -235,6 +235,8 @@ T *replaceSymbol(Symbol *s, ArgT &&... arg) {
 } // namespace macho
 
 std::string toString(const macho::Symbol &);
+std::string toMachOString(const llvm::object::Archive::Symbol &);
+
 } // namespace lld
 
 #endif
diff --git a/lld/test/MachO/thin-archive.s b/lld/test/MachO/thin-archive.s
new file mode 100644
index 000000000000..f12506eecab4
--- /dev/null
+++ b/lld/test/MachO/thin-archive.s
@@ -0,0 +1,41 @@
+# REQUIRES: x86
+
+# RUN: split-file %s %t
+# RUN: llvm-mc -filetype=obj -triple=x86_64-apple-macos -o %t/main.o %t/main.s
+
+# RUN: llvm-mc -filetype=obj -triple=x86_64-apple-macos -o %t/lib.o \
+# RUN:     %t/mangled-symbol.s
+# RUN: llvm-ar csr  %t/lib.a      %t/lib.o
+# RUN: llvm-ar csrT %t/lib_thin.a %t/lib.o
+
+# RUN: %lld %t/main.o %t/lib.a -o %t/out
+# RUN: llvm-nm %t/out
+# RUN: %lld %t/main.o %t/lib_thin.a -o %t/out
+# RUN: llvm-nm %t/out
+# RUN: %lld /%t/main.o -force_load %t/lib_thin.a -o %t/out
+# RUN: llvm-nm %t/out
+
+# RUN: rm %t/lib.o
+# RUN: %lld %t/main.o %t/lib.a -o %t/out
+# RUN: llvm-nm %t/out
+# RUN: not %lld %t/main.o %t/lib_thin.a -demangle -o %t/out 2>&1 | \
+# RUN:     FileCheck --check-prefix=NOOBJ %s
+# RUN: not %lld %t/main.o %t/lib_thin.a -o %t/out 2>&1 | \
+# RUN:     FileCheck --check-prefix=NOOBJNODEMANGLE %s
+
+# CHECK: __Z1fv
+# CHECK: _main
+# NOOBJ: error: {{.*}}lib_thin.a: could not get the buffer for the member defining symbol f(): '{{.*}}lib.o':
+# NOOBJNODEMANGLE: error: {{.*}}lib_thin.a: could not get the buffer for the member defining symbol __Z1fv: '{{.*}}lib.o':
+
+#--- mangled-symbol.s
+.globl  __Z1fv
+__Z1fv:
+  retq
+
+#--- main.s
+.global _main
+_main:
+  callq __Z1fv
+  mov $0, %rax
+  retq