Reland "[llvm-objcopy] Add support for --strip-sections to remove all section headers leaving only program headers and loadable segment data"

ubsan caught an issue I made where I was converting a null pointer to a
reference.

elf utils implements a particularly extreme form of stripping that I'd
like to support. eu-strip has an option called "strip-sections" that
removes all section headers and leaves only program headers and the
segment data. I have implemented this option partly as a test but mainly
because in Fuchsia we would like to use this option to minimize the size
of our executables. The other strip options that are on my list include
--strip-all and --strip-debug. This is a preliminary implementation that
I'd like to start using in Fuchsia builds if possible. This change
implements such a stripping option for llvm-objcopy

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@315484 91177308-0d34-0410-b5e6-96231b3b80d8

4 files changed, 118 insertions, 22 deletions

diff --git a/test/tools/llvm-objcopy/strip-sections.test b/test/tools/llvm-objcopy/strip-sections.test
new file mode 100644
index 00000000000..66cb34e89b0
--- /dev/null
+++ b/test/tools/llvm-objcopy/strip-sections.test
@@ -0,0 +1,66 @@
+# RUN: yaml2obj %s > %t
+# RUN: llvm-objcopy --strip-sections %t %t2
+# RUN: llvm-readobj -file-headers -program-headers %t2 | FileCheck %s
+# RUN: od -t x1 -j 4096 %t2 | FileCheck %s --check-prefix=DATA
+
+!ELF
+FileHeader:
+  Class:           ELFCLASS64
+  Data:            ELFDATA2LSB
+  Type:            ET_EXEC
+  Machine:         EM_X86_64
+Sections:
+  - Name:            .text
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    AddressAlign:    0x0000000000001000
+    Content:         "DEADBEEF"
+ProgramHeaders:
+- Type: PT_LOAD
+  Flags: [ PF_X, PF_R ]
+  Sections:
+    - Section: .text
+
+#DATA: 0010000 de ad be ef
+
+#CHECK: ElfHeader {
+#CHECK-NEXT:  Ident {
+#CHECK-NEXT:     Magic: (7F 45 4C 46)
+#CHECK-NEXT:     Class: 64-bit (0x2)
+#CHECK-NEXT:     DataEncoding: LittleEndian (0x1)
+#CHECK-NEXT:     FileVersion: 1
+#CHECK-NEXT:     OS/ABI: SystemV (0x0)
+#CHECK-NEXT:     ABIVersion: 0
+#CHECK-NEXT:     Unused: (00 00 00 00 00 00 00)
+#CHECK-NEXT:   }
+#CHECK-NEXT:   Type: Executable (0x2)
+#CHECK-NEXT:   Machine: EM_X86_64 (0x3E)
+#CHECK-NEXT:   Version: 1
+#CHECK-NEXT:   Entry: 0x0
+#CHECK-NEXT:   ProgramHeaderOffset: 0x40
+#CHECK-NEXT:   SectionHeaderOffset: 0x0
+#CHECK-NEXT:   Flags [ (0x0)
+#CHECK-NEXT:   ]
+#CHECK-NEXT:   HeaderSize: 64
+#CHECK-NEXT:   ProgramHeaderEntrySize: 56
+#CHECK-NEXT:   ProgramHeaderCount: 1
+#CHECK-NEXT:   SectionHeaderEntrySize: 64
+#CHECK-NEXT:   SectionHeaderCount: 0
+#CHECK-NEXT:   StringTableSectionIndex: 0
+#CHECK-NEXT: }
+
+#CHECK:     ProgramHeaders [
+#CHECK-NEXT:  ProgramHeader {
+#CHECK-NEXT:    Type: PT_LOAD (0x1)
+#CHECK-NEXT:    Offset: 0x1000
+#CHECK-NEXT:    VirtualAddress: 0x0
+#CHECK-NEXT:    PhysicalAddress: 0x0
+#CHECK-NEXT:    FileSize: 4
+#CHECK-NEXT:    MemSize: 4
+#CHECK-NEXT:    Flags [ (0x5)
+#CHECK-NEXT:      PF_R (0x4)
+#CHECK-NEXT:      PF_X (0x1)
+#CHECK-NEXT:    ]
+#CHECK-NEXT:    Alignment: 4096
+#CHECK-NEXT:  }
+#CHECK-NEXT:]
diff --git a/tools/llvm-objcopy/Object.cpp b/tools/llvm-objcopy/Object.cpp
index 2bd3aa16f0b..f9acf001ae9 100644
--- a/tools/llvm-objcopy/Object.cpp
+++ b/tools/llvm-objcopy/Object.cpp
@@ -575,14 +575,20 @@ void Object<ELFT>::writeHeader(FileOutputBuffer &Out) const {
   Ehdr.e_version = Version;
   Ehdr.e_entry = Entry;
   Ehdr.e_phoff = sizeof(Elf_Ehdr);
-  Ehdr.e_shoff = SHOffset;
   Ehdr.e_flags = Flags;
   Ehdr.e_ehsize = sizeof(Elf_Ehdr);
   Ehdr.e_phentsize = sizeof(Elf_Phdr);
   Ehdr.e_phnum = Segments.size();
   Ehdr.e_shentsize = sizeof(Elf_Shdr);
-  Ehdr.e_shnum = Sections.size() + 1;
-  Ehdr.e_shstrndx = SectionNames->Index;
+  if (WriteSectionHeaders) {
+    Ehdr.e_shoff = SHOffset;
+    Ehdr.e_shnum = Sections.size() + 1;
+    Ehdr.e_shstrndx = SectionNames->Index;
+  } else {
+    Ehdr.e_shoff = 0;
+    Ehdr.e_shnum = 0;
+    Ehdr.e_shstrndx = 0;
+  }
 }
 
 template <class ELFT>
@@ -626,19 +632,19 @@ void Object<ELFT>::removeSections(
       std::begin(Sections), std::end(Sections), [=](const SecPtr &Sec) {
         if (ToRemove(*Sec))
           return false;
-        if (auto RelSec = dyn_cast<RelocationSectionBase>(Sec.get()))
-          return !ToRemove(*RelSec->getSection());
+        if (auto RelSec = dyn_cast<RelocationSectionBase>(Sec.get())) {
+          if (auto ToRelSec = RelSec->getSection())
+            return !ToRemove(*ToRelSec);
+        }
         return true;
       });
   if (SymbolTable != nullptr && ToRemove(*SymbolTable))
     SymbolTable = nullptr;
   if (ToRemove(*SectionNames)) {
-    // Right now llvm-objcopy always outputs section headers. This will not
-    // always be the case. Eventully the section header table will become
-    // optional and if no section header is output then there dosn't need to be
-    // a section header string table.
-    error("Cannot remove " + SectionNames->Name +
-          " because it is the section header string table.");
+    if (WriteSectionHeaders)
+      error("Cannot remove " + SectionNames->Name +
+            " because it is the section header string table.");
+    SectionNames = nullptr;
   }
   // Now make sure there are no remaining references to the sections that will
   // be removed. Sometimes it is impossible to remove a reference so we emit
@@ -732,29 +738,34 @@ template <class ELFT> void ELFObject<ELFT>::assignOffsets() {
     }
   }
 
-  Offset = alignTo(Offset, sizeof(typename ELFT::Addr));
+  if (this->WriteSectionHeaders) {
+    Offset = alignTo(Offset, sizeof(typename ELFT::Addr));
+  }
   this->SHOffset = Offset;
 }
 
 template <class ELFT> size_t ELFObject<ELFT>::totalSize() const {
   // We already have the section header offset so we can calculate the total
   // size by just adding up the size of each section header.
+  auto NullSectionSize = this->WriteSectionHeaders ? sizeof(Elf_Shdr) : 0;
   return this->SHOffset + this->Sections.size() * sizeof(Elf_Shdr) +
-         sizeof(Elf_Shdr);
+         NullSectionSize;
 }
 
 template <class ELFT> void ELFObject<ELFT>::write(FileOutputBuffer &Out) const {
   this->writeHeader(Out);
   this->writeProgramHeaders(Out);
   this->writeSectionData(Out);
-  this->writeSectionHeaders(Out);
+  if (this->WriteSectionHeaders)
+    this->writeSectionHeaders(Out);
 }
 
 template <class ELFT> void ELFObject<ELFT>::finalize() {
   // Make sure we add the names of all the sections.
-  for (const auto &Section : this->Sections) {
-    this->SectionNames->addString(Section->Name);
-  }
+  if (this->SectionNames != nullptr)
+    for (const auto &Section : this->Sections) {
+      this->SectionNames->addString(Section->Name);
+    }
   // Make sure we add the names of all the symbols.
   if (this->SymbolTable != nullptr)
     this->SymbolTable->addSymbolNames();
@@ -763,14 +774,16 @@ template <class ELFT> void ELFObject<ELFT>::finalize() {
   assignOffsets();
 
   // Finalize SectionNames first so that we can assign name indexes.
-  this->SectionNames->finalize();
+  if (this->SectionNames != nullptr)
+    this->SectionNames->finalize();
   // Finally now that all offsets and indexes have been set we can finalize any
   // remaining issues.
   uint64_t Offset = this->SHOffset + sizeof(Elf_Shdr);
   for (auto &Section : this->Sections) {
     Section->HeaderOffset = Offset;
     Offset += sizeof(Elf_Shdr);
-    Section->NameIndex = this->SectionNames->findIndex(Section->Name);
+    if (this->WriteSectionHeaders)
+      Section->NameIndex = this->SectionNames->findIndex(Section->Name);
     Section->finalize();
   }
 }
diff --git a/tools/llvm-objcopy/Object.h b/tools/llvm-objcopy/Object.h
index 391465db65e..f6088434805 100644
--- a/tools/llvm-objcopy/Object.h
+++ b/tools/llvm-objcopy/Object.h
@@ -346,6 +346,7 @@ public:
   uint32_t Machine;
   uint32_t Version;
   uint32_t Flags;
+  bool WriteSectionHeaders = true;
 
   Object(const llvm::object::ELFObjectFile<ELFT> &Obj);
   void removeSections(std::function<bool(const SectionBase &)> ToRemove);
diff --git a/tools/llvm-objcopy/llvm-objcopy.cpp b/tools/llvm-objcopy/llvm-objcopy.cpp
index d76735482d6..7f55a434b33 100644
--- a/tools/llvm-objcopy/llvm-objcopy.cpp
+++ b/tools/llvm-objcopy/llvm-objcopy.cpp
@@ -56,11 +56,14 @@ cl::opt<std::string> OutputFilename(cl::Positional, cl::desc("<output>"),
 cl::opt<std::string>
     OutputFormat("O", cl::desc("set output format to one of the following:"
                                "\n\tbinary"));
-
 cl::list<std::string> ToRemove("remove-section",
                                cl::desc("Remove a specific section"));
 cl::alias ToRemoveA("R", cl::desc("Alias for remove-section"),
                     cl::aliasopt(ToRemove));
+cl::opt<bool> StripSections("strip-sections",
+                            cl::desc("Remove all section headers"));
+
+typedef std::function<bool(const SectionBase &Sec)> SectionPred;
 
 void CopyBinary(const ELFObjectFile<ELF64LE> &ObjFile) {
   std::unique_ptr<FileOutputBuffer> Buffer;
@@ -71,12 +74,25 @@ void CopyBinary(const ELFObjectFile<ELF64LE> &ObjFile) {
     Obj = llvm::make_unique<BinaryObject<ELF64LE>>(ObjFile);
   else
     Obj = llvm::make_unique<ELFObject<ELF64LE>>(ObjFile);
+
+  SectionPred RemovePred = [](const SectionBase &) { return false; };
+
   if (!ToRemove.empty()) {
-    Obj->removeSections([&](const SectionBase &Sec) {
+    RemovePred = [&](const SectionBase &Sec) {
       return std::find(std::begin(ToRemove), std::end(ToRemove), Sec.Name) !=
              std::end(ToRemove);
-    });
+    };
   }
+
+  if (StripSections) {
+    RemovePred = [RemovePred](const SectionBase &Sec) {
+      return RemovePred(Sec) || (Sec.Flags & SHF_ALLOC) == 0;
+    };
+    Obj->WriteSectionHeaders = false;
+  }
+
+  Obj->removeSections(RemovePred);
+
   Obj->finalize();
   ErrorOr<std::unique_ptr<FileOutputBuffer>> BufferOrErr =
       FileOutputBuffer::create(OutputFilename, Obj->totalSize(),