[lld-macho] Rework length check when opening input files

This reverts diff D97610 (commit 0223ab0) and adds a one-line fix to verify that a `MemoryBufferRef` has sufficient length before reading a 4-byte magic number.

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

Author: gkmhub

lld/MachO/Driver.cpp

@@ -263,7 +263,7 @@ static std::vector<ArchiveMember> getArchiveMembers(MemoryBufferRef mb) {
 
 static InputFile *addFile(StringRef path, bool forceLoadArchive,
                           bool isBundleLoader = false) {
-  Optional<MemoryBufferRef> buffer = readLinkableFile(path);
+  Optional<MemoryBufferRef> buffer = readFile(path);
   if (!buffer)
     return nullptr;
   MemoryBufferRef mbref = *buffer;
@@ -279,7 +279,7 @@ static InputFile *addFile(StringRef path, bool forceLoadArchive,
       error(path + ": archive has no index; run ranlib to add one");
 
     if (config->allLoad || forceLoadArchive) {
-      if (Optional<MemoryBufferRef> buffer = readLinkableFile(path)) {
+      if (Optional<MemoryBufferRef> buffer = readFile(path)) {
         for (const ArchiveMember &member : getArchiveMembers(*buffer)) {
           if (Optional<InputFile *> file = loadArchiveMember(
                   member.mbref, member.modTime, path, /*objCOnly=*/false)) {
@@ -300,7 +300,7 @@ static InputFile *addFile(StringRef path, bool forceLoadArchive,
       // we already found that it contains an ObjC symbol. We should also
       // consider creating a LazyObjFile class in order to avoid double-loading
       // these files here and below (as part of the ArchiveFile).
-      if (Optional<MemoryBufferRef> buffer = readLinkableFile(path)) {
+      if (Optional<MemoryBufferRef> buffer = readFile(path)) {
         for (const ArchiveMember &member : getArchiveMembers(*buffer)) {
           if (Optional<InputFile *> file = loadArchiveMember(
                   member.mbref, member.modTime, path, /*objCOnly=*/true)) {
@@ -403,7 +403,7 @@ void macho::parseLCLinkerOption(InputFile* f, unsigned argc, StringRef data) {
 }
 
 static void addFileList(StringRef path) {
-  Optional<MemoryBufferRef> buffer = readRawFile(path);
+  Optional<MemoryBufferRef> buffer = readFile(path);
   if (!buffer)
     return;
   MemoryBufferRef mbref = *buffer;
@@ -426,7 +426,7 @@ static void addFileList(StringRef path) {
 //
 // The file can also have line comments that start with '#'.
 static void parseOrderFile(StringRef path) {
-  Optional<MemoryBufferRef> buffer = readRawFile(path);
+  Optional<MemoryBufferRef> buffer = readFile(path);
   if (!buffer) {
     error("Could not read order file at " + path);
     return;
@@ -940,7 +940,7 @@ bool macho::link(ArrayRef<const char *> argsArr, bool canExitEarly,
     StringRef segName = arg->getValue(0);
     StringRef sectName = arg->getValue(1);
     StringRef fileName = arg->getValue(2);
-    Optional<MemoryBufferRef> buffer = readRawFile(fileName);
+    Optional<MemoryBufferRef> buffer = readFile(fileName);
     if (buffer)
       inputFiles.insert(make<OpaqueFile>(*buffer, segName, sectName));
   }

lld/MachO/DriverUtils.cpp

@@ -132,7 +132,7 @@ std::string macho::createResponseFile(const opt::InputArgList &args) {
       os << "-o " << quote(path::filename(arg->getValue())) << "\n";
       break;
     case OPT_filelist:
-      if (Optional<MemoryBufferRef> buffer = readRawFile(arg->getValue()))
+      if (Optional<MemoryBufferRef> buffer = readFile(arg->getValue()))
         for (StringRef path : args::getLines(*buffer))
           os << quote(rewritePath(path)) << "\n";
       break;

lld/MachO/InputFiles.cpp

@@ -91,8 +91,7 @@ std::unique_ptr<TarWriter> macho::tar;
 int InputFile::idCount = 0;
 
 // Open a given file path and return it as a memory-mapped file.
-// Perform no sanity checks--just open, map & return.
-Optional<MemoryBufferRef> macho::readRawFile(StringRef path) {
+Optional<MemoryBufferRef> macho::readFile(StringRef path) {
   // Open a file.
   auto mbOrErr = MemoryBuffer::getFile(path);
   if (auto ec = mbOrErr.getError()) {
@@ -103,32 +102,12 @@ Optional<MemoryBufferRef> macho::readRawFile(StringRef path) {
   std::unique_ptr<MemoryBuffer> &mb = *mbOrErr;
   MemoryBufferRef mbref = mb->getMemBufferRef();
   make<std::unique_ptr<MemoryBuffer>>(std::move(mb)); // take mb ownership
-  return mbref;
-}
-
-// Open a given file path and return it as a memory-mapped file.
-// Assume the file has one of a variety of linkable formats and
-// perform some basic sanity checks, notably minimum length.
-Optional<MemoryBufferRef> macho::readLinkableFile(StringRef path) {
-  Optional<MemoryBufferRef> maybeMbref = readRawFile(path);
-  if (!maybeMbref) {
-    return None;
-  }
-  MemoryBufferRef mbref = *maybeMbref;
-
-  // LD64 hard-codes 20 as minimum header size, which is presumably
-  // the smallest header among the the various linkable input formats
-  // LLD are less demanding. We insist on having only enough data for
-  // a magic number.
-  if (mbref.getBufferSize() < sizeof(uint32_t)) {
-    error("file is too small to contain a magic number: " + path);
-    return None;
-  }
 
   // If this is a regular non-fat file, return it.
   const char *buf = mbref.getBufferStart();
   auto *hdr = reinterpret_cast<const MachO::fat_header *>(buf);
-  if (read32be(&hdr->magic) != MachO::FAT_MAGIC) {
+  if (mbref.getBufferSize() < sizeof(uint32_t) ||
+      read32be(&hdr->magic) != MachO::FAT_MAGIC) {
     if (tar)
       tar->append(relativeToRoot(path), mbref.getBuffer());
     return mbref;
@@ -565,7 +544,7 @@ void ObjFile::parseDebugInfo() {
 
 // The path can point to either a dylib or a .tbd file.
 static Optional<DylibFile *> loadDylib(StringRef path, DylibFile *umbrella) {
-  Optional<MemoryBufferRef> mbref = readLinkableFile(path);
+  Optional<MemoryBufferRef> mbref = readFile(path);
   if (!mbref) {
     error("could not read dylib file at " + path);
     return {};

lld/MachO/InputFiles.h

@@ -173,8 +173,7 @@ class BitcodeFile : public InputFile {
 
 extern llvm::SetVector<InputFile *> inputFiles;
 
-llvm::Optional<MemoryBufferRef> readRawFile(StringRef path);
-llvm::Optional<MemoryBufferRef> readLinkableFile(StringRef path);
+llvm::Optional<MemoryBufferRef> readFile(StringRef path);
 
 const llvm::MachO::load_command *
 findCommand(const llvm::MachO::mach_header_64 *, uint32_t type);

lld/test/MachO/invalid/tiny-input.s

@@ -14,7 +14,7 @@
 # BAD1-DAG: error: invalid name for segment or section: S/ASHY_SEG
 # BAD1-DAG: error: invalid name for segment or section: st*rry_sect
 # BAD1-DAG: error: invalid name for segment or section: -o
-# BAD1-DAG: error: file is too small to contain a magic number:
+# BAD1-DAG: error: {{.*}}: unhandled file type
 
 # RUN: not %lld \
 # RUN:     -rename_segment H#SHY_SEG PL+SSY_SEG \
@@ -24,7 +24,7 @@
 # BAD2-DAG: error: invalid name for segment or section: H#SHY_SEG
 # BAD2-DAG: error: invalid name for segment or section: PL+SSY_SEG
 # BAD2-DAG: error: invalid name for segment or section: -o
-# BAD2-DAG: error: file is too small to contain a magic number:
+# BAD2-DAG: error: {{.*}}: unhandled file type
 
 ## Check that section and segment renames happen
 # RUN: %lld \