From 38cb238f752cefe5fde8bf6cb73dacc3fec8f89d Mon Sep 17 00:00:00 2001
From: James Henderson <jh7370@my.bristol.ac.uk>
Date: Tue, 2 Apr 2019 14:11:13 +0000
Subject: [PATCH] [llvm-objcopy]Allow llvm-objcopy to be used on an ELF file
 with no section headers

This patch fixes https://bugs.llvm.org/show_bug.cgi?id=41293 and
https://bugs.llvm.org/show_bug.cgi?id=41045. llvm-objcopy assumed that
it could always read a section header string table. This isn't the case
when the sections were previously all stripped, and the e_shstrndx field
was set to 0. This patch fixes this. It also fixes a double space in an
error message relating to this issue, and prevents llvm-objcopy from
adding extra space for non-existent section headers, meaning that
--strip-sections on the output of a previous --strip-sections run
produces identical output, simplifying the test.

Reviewed by: rupprecht, grimar

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

llvm-svn: 357475
---
 .../ELF/copy-after-strip-sections.test        | 26 +++++++++++++++++
 llvm/tools/llvm-objcopy/ELF/Object.cpp        | 28 ++++++++++++-------
 llvm/tools/llvm-objcopy/ELF/Object.h          |  6 ++--
 3 files changed, 47 insertions(+), 13 deletions(-)
 create mode 100644 llvm/test/tools/llvm-objcopy/ELF/copy-after-strip-sections.test

diff --git a/llvm/test/tools/llvm-objcopy/ELF/copy-after-strip-sections.test b/llvm/test/tools/llvm-objcopy/ELF/copy-after-strip-sections.test
new file mode 100644
index 0000000000000..3b61649b2639c
--- /dev/null
+++ b/llvm/test/tools/llvm-objcopy/ELF/copy-after-strip-sections.test
@@ -0,0 +1,26 @@
+# llvm-objcopy's --strip-sections removes the section headers. It should be
+# possible to run the tool on the output after this operation. Performing any
+# subsequent stripping operation, or copying the object, should produce
+# identical output.
+
+# RUN: yaml2obj %s -o %t.in
+# RUN: llvm-objcopy %t.in %t.stripped --strip-sections
+# RUN: llvm-objcopy %t.stripped %t.stripped2 --strip-sections
+# RUN: llvm-objcopy %t.stripped2 %t.out
+# RUN: cmp %t.stripped %t.stripped2
+# RUN: cmp %t.stripped %t.out
+
+--- !ELF
+FileHeader:
+  Class:   ELFCLASS64
+  Data:    ELFDATA2LSB
+  Type:    ET_EXEC
+  Machine: EM_X86_64
+Sections:
+  - Name: .text
+    Type: SHT_PROGBITS
+    Content: 'facefeed'
+ProgramHeaders:
+  - Type: PT_LOAD
+    Sections:
+      - Section: .text
diff --git a/llvm/tools/llvm-objcopy/ELF/Object.cpp b/llvm/tools/llvm-objcopy/ELF/Object.cpp
index 321adf6fff0fc..de51851ecfe50 100644
--- a/llvm/tools/llvm-objcopy/ELF/Object.cpp
+++ b/llvm/tools/llvm-objcopy/ELF/Object.cpp
@@ -1212,13 +1212,16 @@ template <class ELFT> void ELFBuilder<ELFT>::build() {
   if (ShstrIndex == SHN_XINDEX)
     ShstrIndex = unwrapOrError(ElfFile.getSection(0))->sh_link;
 
-  Obj.SectionNames =
-      Obj.sections().template getSectionOfType<StringTableSection>(
-          ShstrIndex,
-          "e_shstrndx field value " + Twine(Ehdr.e_shstrndx) +
-              " in elf header " + " is invalid",
-          "e_shstrndx field value " + Twine(Ehdr.e_shstrndx) +
-              " in elf header " + " is not a string table");
+  if (ShstrIndex == SHN_UNDEF)
+    Obj.HadShdrs = false;
+  else
+    Obj.SectionNames =
+        Obj.sections().template getSectionOfType<StringTableSection>(
+            ShstrIndex,
+            "e_shstrndx field value " + Twine(Ehdr.e_shstrndx) +
+                " in elf header is invalid",
+            "e_shstrndx field value " + Twine(Ehdr.e_shstrndx) +
+                " in elf header is not a string table");
 }
 
 Writer::~Writer() {}
@@ -1372,6 +1375,10 @@ template <class ELFT> void ELFWriter<ELFT>::writeSegmentData() {
   }
 }
 
+template <class ELFT>
+ELFWriter<ELFT>::ELFWriter(Object &Obj, Buffer &Buf, bool WSH)
+    : Writer(Obj, Buf), WriteSectionHeaders(WSH && Obj.HadShdrs) {}
+
 Error Object::removeSections(
     std::function<bool(const SectionBase &)> ToRemove) {
 
@@ -1558,9 +1565,10 @@ template <class ELFT> void ELFWriter<ELFT>::assignOffsets() {
 template <class ELFT> size_t ELFWriter<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 = WriteSectionHeaders ? sizeof(Elf_Shdr) : 0;
-  return Obj.SHOffset + Obj.sections().size() * sizeof(Elf_Shdr) +
-         NullSectionSize;
+  if (!WriteSectionHeaders)
+    return Obj.SHOffset;
+  size_t ShdrCount = Obj.sections().size() + 1; // Includes null shdr.
+  return Obj.SHOffset + ShdrCount * sizeof(Elf_Shdr);
 }
 
 template <class ELFT> Error ELFWriter<ELFT>::write() {
diff --git a/llvm/tools/llvm-objcopy/ELF/Object.h b/llvm/tools/llvm-objcopy/ELF/Object.h
index f85d4b7001855..4e0e23e5f49b8 100644
--- a/llvm/tools/llvm-objcopy/ELF/Object.h
+++ b/llvm/tools/llvm-objcopy/ELF/Object.h
@@ -226,12 +226,11 @@ template <class ELFT> class ELFWriter : public Writer {
 
 public:
   virtual ~ELFWriter() {}
-  bool WriteSectionHeaders = true;
+  bool WriteSectionHeaders;
 
   Error finalize() override;
   Error write() override;
-  ELFWriter(Object &Obj, Buffer &Buf, bool WSH)
-      : Writer(Obj, Buf), WriteSectionHeaders(WSH) {}
+  ELFWriter(Object &Obj, Buffer &Buf, bool WSH);
 };
 
 class BinaryWriter : public Writer {
@@ -810,6 +809,7 @@ class Object {
   uint32_t Version;
   uint32_t Flags;
 
+  bool HadShdrs = true;
   StringTableSection *SectionNames = nullptr;
   SymbolTableSection *SymbolTable = nullptr;
   SectionIndexSection *SectionIndexTable = nullptr;