1 files changed, 195 insertions, 53 deletions

diff --git a/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp b/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
index 4578d0bb75d4f..5ca5b133572b4 100644
--- a/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
+++ b/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
@@ -10,6 +10,7 @@
 #include "../CopyConfig.h"
 #include "MachOReader.h"
 #include "MachOWriter.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
 
@@ -18,21 +19,44 @@ namespace objcopy {
 namespace macho {
 
 using namespace object;
-using SectionPred = std::function<bool(const Section &Sec)>;
+using SectionPred = std::function<bool(const std::unique_ptr<Section> &Sec)>;
+using LoadCommandPred = std::function<bool(const LoadCommand &LC)>;
+
+#ifndef NDEBUG
+static bool isLoadCommandWithPayloadString(const LoadCommand &LC) {
+  // TODO: Add support for LC_REEXPORT_DYLIB, LC_LOAD_UPWARD_DYLIB and
+  // LC_LAZY_LOAD_DYLIB
+  return LC.MachOLoadCommand.load_command_data.cmd == MachO::LC_RPATH ||
+         LC.MachOLoadCommand.load_command_data.cmd == MachO::LC_ID_DYLIB ||
+         LC.MachOLoadCommand.load_command_data.cmd == MachO::LC_LOAD_DYLIB ||
+         LC.MachOLoadCommand.load_command_data.cmd == MachO::LC_LOAD_WEAK_DYLIB;
+}
+#endif
+
+static StringRef getPayloadString(const LoadCommand &LC) {
+  assert(isLoadCommandWithPayloadString(LC) &&
+         "unsupported load command encountered");
+
+  return StringRef(reinterpret_cast<const char *>(LC.Payload.data()),
+                   LC.Payload.size())
+      .rtrim('\0');
+}
 
-static void removeSections(const CopyConfig &Config, Object &Obj) {
-  SectionPred RemovePred = [](const Section &) { return false; };
+static Error removeSections(const CopyConfig &Config, Object &Obj) {
+  SectionPred RemovePred = [](const std::unique_ptr<Section> &) {
+    return false;
+  };
 
   if (!Config.ToRemove.empty()) {
-    RemovePred = [&Config, RemovePred](const Section &Sec) {
-      return Config.ToRemove.matches(Sec.CanonicalName);
+    RemovePred = [&Config, RemovePred](const std::unique_ptr<Section> &Sec) {
+      return Config.ToRemove.matches(Sec->CanonicalName);
     };
   }
 
   if (Config.StripAll || Config.StripDebug) {
     // Remove all debug sections.
-    RemovePred = [RemovePred](const Section &Sec) {
-      if (Sec.Segname == "__DWARF")
+    RemovePred = [RemovePred](const std::unique_ptr<Section> &Sec) {
+      if (Sec->Segname == "__DWARF")
         return true;
 
       return RemovePred(Sec);
@@ -41,8 +65,8 @@ static void removeSections(const CopyConfig &Config, Object &Obj) {
 
   if (!Config.OnlySection.empty()) {
     // Overwrite RemovePred because --only-section takes priority.
-    RemovePred = [&Config](const Section &Sec) {
-      return !Config.OnlySection.matches(Sec.CanonicalName);
+    RemovePred = [&Config](const std::unique_ptr<Section> &Sec) {
+      return !Config.OnlySection.matches(Sec->CanonicalName);
     };
   }
 
@@ -60,41 +84,158 @@ static void updateAndRemoveSymbols(const CopyConfig &Config, Object &Obj) {
   for (SymbolEntry &Sym : Obj.SymTable) {
     auto I = Config.SymbolsToRename.find(Sym.Name);
     if (I != Config.SymbolsToRename.end())
-      Sym.Name = I->getValue();
+      Sym.Name = std::string(I->getValue());
   }
 
-  auto RemovePred = [Config](const std::unique_ptr<SymbolEntry> &N) {
+  auto RemovePred = [Config, &Obj](const std::unique_ptr<SymbolEntry> &N) {
     if (N->Referenced)
       return false;
-    return Config.StripAll;
+    if (Config.StripAll)
+      return true;
+    if (Config.DiscardMode == DiscardType::All && !(N->n_type & MachO::N_EXT))
+      return true;
+    // This behavior is consistent with cctools' strip.
+    if (Config.StripSwiftSymbols && (Obj.Header.Flags & MachO::MH_DYLDLINK) &&
+        Obj.SwiftVersion && *Obj.SwiftVersion && N->isSwiftSymbol())
+      return true;
+    return false;
   };
 
   Obj.SymTable.removeSymbols(RemovePred);
 }
 
+template <typename LCType>
+static void updateLoadCommandPayloadString(LoadCommand &LC, StringRef S) {
+  assert(isLoadCommandWithPayloadString(LC) &&
+         "unsupported load command encountered");
+
+  uint32_t NewCmdsize = alignTo(sizeof(LCType) + S.size() + 1, 8);
+
+  LC.MachOLoadCommand.load_command_data.cmdsize = NewCmdsize;
+  LC.Payload.assign(NewCmdsize - sizeof(LCType), 0);
+  std::copy(S.begin(), S.end(), LC.Payload.begin());
+}
+
 static LoadCommand buildRPathLoadCommand(StringRef Path) {
   LoadCommand LC;
   MachO::rpath_command RPathLC;
   RPathLC.cmd = MachO::LC_RPATH;
   RPathLC.path = sizeof(MachO::rpath_command);
-  RPathLC.cmdsize = alignTo(sizeof(MachO::rpath_command) + Path.size(), 8);
+  RPathLC.cmdsize = alignTo(sizeof(MachO::rpath_command) + Path.size() + 1, 8);
   LC.MachOLoadCommand.rpath_command_data = RPathLC;
   LC.Payload.assign(RPathLC.cmdsize - sizeof(MachO::rpath_command), 0);
   std::copy(Path.begin(), Path.end(), LC.Payload.begin());
   return LC;
 }
 
+static Error processLoadCommands(const CopyConfig &Config, Object &Obj) {
+  // Remove RPaths.
+  DenseSet<StringRef> RPathsToRemove(Config.RPathsToRemove.begin(),
+                                     Config.RPathsToRemove.end());
+
+  LoadCommandPred RemovePred = [&RPathsToRemove](const LoadCommand &LC) {
+    if (LC.MachOLoadCommand.load_command_data.cmd == MachO::LC_RPATH) {
+      StringRef RPath = getPayloadString(LC);
+      if (RPathsToRemove.count(RPath)) {
+        RPathsToRemove.erase(RPath);
+        return true;
+      }
+    }
+    return false;
+  };
+
+  if (Error E = Obj.removeLoadCommands(RemovePred))
+    return E;
+
+  // Emit an error if the Mach-O binary does not contain an rpath path name
+  // specified in -delete_rpath.
+  for (StringRef RPath : Config.RPathsToRemove) {
+    if (RPathsToRemove.count(RPath))
+      return createStringError(errc::invalid_argument,
+                               "no LC_RPATH load command with path: %s",
+                               RPath.str().c_str());
+  }
+
+  DenseSet<StringRef> RPaths;
+
+  // Get all existing RPaths.
+  for (LoadCommand &LC : Obj.LoadCommands) {
+    if (LC.MachOLoadCommand.load_command_data.cmd == MachO::LC_RPATH)
+      RPaths.insert(getPayloadString(LC));
+  }
+
+  // Throw errors for invalid RPaths.
+  for (const auto &OldNew : Config.RPathsToUpdate) {
+    StringRef Old = OldNew.getFirst();
+    StringRef New = OldNew.getSecond();
+    if (RPaths.count(Old) == 0)
+      return createStringError(errc::invalid_argument,
+                               "no LC_RPATH load command with path: " + Old);
+    if (RPaths.count(New) != 0)
+      return createStringError(errc::invalid_argument,
+                               "rpath " + New +
+                                   " would create a duplicate load command");
+  }
+
+  // Update load commands.
+  for (LoadCommand &LC : Obj.LoadCommands) {
+    switch (LC.MachOLoadCommand.load_command_data.cmd) {
+    case MachO::LC_ID_DYLIB:
+      if (Config.SharedLibId) {
+        StringRef Id = Config.SharedLibId.getValue();
+        if (Id.empty())
+          return createStringError(errc::invalid_argument,
+                                   "cannot specify an empty id");
+        updateLoadCommandPayloadString<MachO::dylib_command>(LC, Id);
+      }
+      break;
+
+    case MachO::LC_RPATH: {
+      StringRef RPath = getPayloadString(LC);
+      StringRef NewRPath = Config.RPathsToUpdate.lookup(RPath);
+      if (!NewRPath.empty())
+        updateLoadCommandPayloadString<MachO::rpath_command>(LC, NewRPath);
+      break;
+    }
+
+    // TODO: Add LC_REEXPORT_DYLIB, LC_LAZY_LOAD_DYLIB, and LC_LOAD_UPWARD_DYLIB
+    // here once llvm-objcopy supports them.
+    case MachO::LC_LOAD_DYLIB:
+    case MachO::LC_LOAD_WEAK_DYLIB:
+      StringRef InstallName = getPayloadString(LC);
+      StringRef NewInstallName =
+          Config.InstallNamesToUpdate.lookup(InstallName);
+      if (!NewInstallName.empty())
+        updateLoadCommandPayloadString<MachO::dylib_command>(LC,
+                                                             NewInstallName);
+      break;
+    }
+  }
+
+  // Add new RPaths.
+  for (StringRef RPath : Config.RPathToAdd) {
+    if (RPaths.count(RPath) != 0)
+      return createStringError(errc::invalid_argument,
+                               "rpath " + RPath +
+                                   " would create a duplicate load command");
+    RPaths.insert(RPath);
+    Obj.addLoadCommand(buildRPathLoadCommand(RPath));
+  }
+
+  return Error::success();
+}
+
 static Error dumpSectionToFile(StringRef SecName, StringRef Filename,
                                Object &Obj) {
   for (LoadCommand &LC : Obj.LoadCommands)
-    for (Section &Sec : LC.Sections) {
-      if (Sec.CanonicalName == SecName) {
+    for (const std::unique_ptr<Section> &Sec : LC.Sections) {
+      if (Sec->CanonicalName == SecName) {
         Expected<std::unique_ptr<FileOutputBuffer>> BufferOrErr =
-            FileOutputBuffer::create(Filename, Sec.Content.size());
+            FileOutputBuffer::create(Filename, Sec->Content.size());
         if (!BufferOrErr)
           return BufferOrErr.takeError();
         std::unique_ptr<FileOutputBuffer> Buf = std::move(*BufferOrErr);
-        llvm::copy(Sec.Content, Buf->getBufferStart());
+        llvm::copy(Sec->Content, Buf->getBufferStart());
 
         if (Error E = Buf->commit())
           return E;
@@ -122,7 +263,7 @@ static Error addSection(StringRef SecName, StringRef Filename, Object &Obj) {
   for (LoadCommand &LC : Obj.LoadCommands) {
     Optional<StringRef> SegName = LC.getSegmentName();
     if (SegName && SegName == TargetSegName) {
-      LC.Sections.push_back(Sec);
+      LC.Sections.push_back(std::make_unique<Section>(Sec));
       return Error::success();
     }
   }
@@ -130,7 +271,7 @@ static Error addSection(StringRef SecName, StringRef Filename, Object &Obj) {
   // There's no segment named TargetSegName. Create a new load command and
   // Insert a new section into it.
   LoadCommand &NewSegment = Obj.addSegment(TargetSegName);
-  NewSegment.Sections.push_back(Sec);
+  NewSegment.Sections.push_back(std::make_unique<Section>(Sec));
   return Error::success();
 }
 
@@ -167,17 +308,27 @@ static Error handleArgs(const CopyConfig &Config, Object &Obj) {
       !Config.SectionsToRename.empty() ||
       !Config.UnneededSymbolsToRemove.empty() ||
       !Config.SetSectionAlignment.empty() || !Config.SetSectionFlags.empty() ||
-      Config.ExtractDWO || Config.KeepFileSymbols || Config.LocalizeHidden ||
-      Config.PreserveDates || Config.StripAllGNU || Config.StripDWO ||
-      Config.StripNonAlloc || Config.StripSections || Config.Weaken ||
-      Config.DecompressDebugSections || Config.StripNonAlloc ||
-      Config.StripSections || Config.StripUnneeded ||
-      Config.DiscardMode != DiscardType::None || !Config.SymbolsToAdd.empty() ||
-      Config.EntryExpr) {
+      Config.ExtractDWO || Config.LocalizeHidden || Config.PreserveDates ||
+      Config.StripAllGNU || Config.StripDWO || Config.StripNonAlloc ||
+      Config.StripSections || Config.Weaken || Config.DecompressDebugSections ||
+      Config.StripNonAlloc || Config.StripSections || Config.StripUnneeded ||
+      Config.DiscardMode == DiscardType::Locals ||
+      !Config.SymbolsToAdd.empty() || Config.EntryExpr) {
     return createStringError(llvm::errc::invalid_argument,
                              "option not supported by llvm-objcopy for MachO");
   }
-  removeSections(Config, Obj);
+
+  // Dump sections before add/remove for compatibility with GNU objcopy.
+  for (StringRef Flag : Config.DumpSection) {
+    StringRef SectionName;
+    StringRef FileName;
+    std::tie(SectionName, FileName) = Flag.split('=');
+    if (Error E = dumpSectionToFile(SectionName, FileName, Obj))
+      return E;
+  }
+
+  if (Error E = removeSections(Config, Obj))
+    return E;
 
   // Mark symbols to determine which symbols are still needed.
   if (Config.StripAll)
@@ -187,16 +338,8 @@ static Error handleArgs(const CopyConfig &Config, Object &Obj) {
 
   if (Config.StripAll)
     for (LoadCommand &LC : Obj.LoadCommands)
-      for (Section &Sec : LC.Sections)
-        Sec.Relocations.clear();
-
-  for (const StringRef &Flag : Config.DumpSection) {
-    std::pair<StringRef, StringRef> SecPair = Flag.split("=");
-    StringRef SecName = SecPair.first;
-    StringRef File = SecPair.second;
-    if (Error E = dumpSectionToFile(SecName, File, Obj))
-      return E;
-  }
+      for (std::unique_ptr<Section> &Sec : LC.Sections)
+        Sec->Relocations.clear();
 
   for (const auto &Flag : Config.AddSection) {
     std::pair<StringRef, StringRef> SecPair = Flag.split("=");
@@ -208,19 +351,9 @@ static Error handleArgs(const CopyConfig &Config, Object &Obj) {
       return E;
   }
 
-  for (StringRef RPath : Config.RPathToAdd) {
-    for (LoadCommand &LC : Obj.LoadCommands) {
-      if (LC.MachOLoadCommand.load_command_data.cmd == MachO::LC_RPATH &&
-          RPath == StringRef(reinterpret_cast<char *>(LC.Payload.data()),
-                             LC.Payload.size())
-                       .trim(0)) {
-        return createStringError(errc::invalid_argument,
-                                 "rpath " + RPath +
-                                     " would create a duplicate load command");
-      }
-    }
-    Obj.addLoadCommand(buildRPathLoadCommand(RPath));
-  }
+  if (Error E = processLoadCommands(Config, Obj))
+    return E;
+
   return Error::success();
 }
 
@@ -237,9 +370,18 @@ Error executeObjcopyOnBinary(const CopyConfig &Config,
   if (Error E = handleArgs(Config, *O))
     return createFileError(Config.InputFilename, std::move(E));
 
-  // TODO: Support 16KB pages which are employed in iOS arm64 binaries:
-  //       https://github.com/llvm/llvm-project/commit/1bebb2832ee312d3b0316dacff457a7a29435edb
-  const uint64_t PageSize = 4096;
+  // Page size used for alignment of segment sizes in Mach-O executables and
+  // dynamic libraries.
+  uint64_t PageSize;
+  switch (In.getArch()) {
+  case Triple::ArchType::arm:
+  case Triple::ArchType::aarch64:
+  case Triple::ArchType::aarch64_32:
+    PageSize = 16384;
+    break;
+  default:
+    PageSize = 4096;
+  }
 
   MachOWriter Writer(*O, In.is64Bit(), In.isLittleEndian(), PageSize, Out);
   if (auto E = Writer.finalize())