[llvm-objcopy] Adding support for decompressing zlib compressed dwarf sections.

Summary: I had added support for compressing dwarf sections in a prior commit,
         this one adds support for decompressing. Usage is:

         llvm-objcopy --decompress-debug-sections input.o output.o

Reviewers: jakehehrlich, jhenderson, alexshap	

Reviewed By: jhenderson

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

llvm-svn: 343451

8 files changed, 205 insertions, 14 deletions

diff --git a/llvm/test/tools/llvm-objcopy/compress-and-decompress-debug-sections-error.test b/llvm/test/tools/llvm-objcopy/compress-and-decompress-debug-sections-error.test
new file mode 100644
index 00000000000..0478c22f325
--- /dev/null
+++ b/llvm/test/tools/llvm-objcopy/compress-and-decompress-debug-sections-error.test
@@ -0,0 +1,5 @@
+# RUN: yaml2obj %p/Inputs/compress-debug-sections.yaml -o %t.o
+# RUN: not llvm-objcopy --compress-debug-sections=zlib --decompress-debug-sections %t.o 2>&1 | FileCheck %s
+
+# CHECK: Cannot specify --compress-debug-sections at the same time as --decompress-debug-sections at the same time.
+
diff --git a/llvm/test/tools/llvm-objcopy/compress-debug-sections-zlib-gnu.test b/llvm/test/tools/llvm-objcopy/compress-debug-sections-zlib-gnu.test
index eb697522de2..7199b2e5e11 100644
--- a/llvm/test/tools/llvm-objcopy/compress-debug-sections-zlib-gnu.test
+++ b/llvm/test/tools/llvm-objcopy/compress-debug-sections-zlib-gnu.test
@@ -2,12 +2,26 @@
 
 # RUN: yaml2obj %p/Inputs/compress-debug-sections.yaml -o %t.o
 # RUN: llvm-objcopy --compress-debug-sections=zlib-gnu %t.o %t-compressed.o
+# RUN: llvm-objcopy --decompress-debug-sections %t-compressed.o %t-decompressed.o
 
 # RUN: llvm-objdump -s %t.o -section=.debug_foo | FileCheck %s
 # RUN: llvm-objdump -s %t-compressed.o | FileCheck %s --check-prefix=CHECK-COMPRESSED
 # RUN: llvm-readobj -relocations -s %t-compressed.o | FileCheck %s --check-prefix=CHECK-FLAGS
+# RUN: llvm-readobj -relocations -s %t-decompressed.o | FileCheck %s --check-prefix=CHECK-HEADER
+# RUN: llvm-readobj -relocations -s %t.o | FileCheck %s --check-prefix=CHECK-HEADER
+# RUN: llvm-objdump -s %t-decompressed.o -section=.debug_foo | FileCheck %s
 
 # CHECK: .debug_foo:
+# CHECK-NEXT: 0000 00000000 00000000
+
+# CHECK-HEADER: Index: 1
+# CHECK-HEADER-NEXT: Name: .debug_foo
+# CHECK-HEADER-NEXT: Type: SHT_PROGBITS
+# CHECK-HEADER-NEXT: Flags [
+# CHECK-HEADER-NEXT: ]
+# CHECK-HEADER-NEXT: Address:
+# CHECK-HEADER-NEXT: Offset:
+# CHECK-HEADER-NEXT: Size: 8
 
 # CHECK-COMPRESSED: .zdebug_foo:
 # CHECK-COMPRESSED: ZLIB
diff --git a/llvm/test/tools/llvm-objcopy/compress-debug-sections-zlib.test b/llvm/test/tools/llvm-objcopy/compress-debug-sections-zlib.test
index 843dbf267d7..d2ca6ef249b 100644
--- a/llvm/test/tools/llvm-objcopy/compress-debug-sections-zlib.test
+++ b/llvm/test/tools/llvm-objcopy/compress-debug-sections-zlib.test
@@ -2,12 +2,26 @@
 
 # RUN: yaml2obj %p/Inputs/compress-debug-sections.yaml -o %t.o
 # RUN: llvm-objcopy --compress-debug-sections=zlib %t.o %t-compressed.o
+# RUN: llvm-objcopy --decompress-debug-sections %t-compressed.o %t-decompressed.o
 
 # RUN: llvm-objdump -s %t.o -section=.debug_foo | FileCheck %s
 # RUN: llvm-objdump -s %t-compressed.o | FileCheck %s --check-prefix=CHECK-COMPRESSED
 # RUN: llvm-readobj -relocations -s %t-compressed.o | FileCheck %s --check-prefix=CHECK-FLAGS
+# RUN: llvm-readobj -relocations -s %t-decompressed.o | FileCheck %s --check-prefix=CHECK-HEADER
+# RUN: llvm-readobj -relocations -s %t.o | FileCheck %s --check-prefix=CHECK-HEADER
+# RUN: llvm-objdump -s %t-decompressed.o -section=.debug_foo | FileCheck %s
 
 # CHECK: .debug_foo:
+# CHECK-NEXT: 0000 00000000 00000000
+
+# CHECK-HEADER: Index: 1
+# CHECK-HEADER-NEXT: Name: .debug_foo
+# CHECK-HEADER-NEXT: Type: SHT_PROGBITS
+# CHECK-HEADER-NEXT: Flags [
+# CHECK-HEADER-NEXT: ]
+# CHECK-HEADER-NEXT: Address:
+# CHECK-HEADER-NEXT: Offset:
+# CHECK-HEADER-NEXT: Size: 8
 
 # CHECK-COMPRESSED: .debug_foo:
 # CHECK-COMPRESSED: .notdebug_foo:
diff --git a/llvm/test/tools/llvm-objcopy/compress-debug-sections.test b/llvm/test/tools/llvm-objcopy/compress-debug-sections.test
new file mode 100644
index 00000000000..6fadad3f383
--- /dev/null
+++ b/llvm/test/tools/llvm-objcopy/compress-debug-sections.test
@@ -0,0 +1,24 @@
+# REQUIRES: zlib
+
+# RUN: yaml2obj %p/Inputs/compress-debug-sections.yaml -o %t.o
+# RUN: cp %t.o %t.copy.o
+
+# RUN: llvm-objcopy --compress-debug-sections=zlib     %t.o %tz.o
+# RUN: llvm-objcopy --compress-debug-sections=zlib-gnu %t.o %tzg.o
+# RUN: cp %tz.o %tz.copy.o
+# RUN: cp %tzg.o %tzg.copy.o
+
+# RUN: llvm-objcopy --decompress-debug-sections %tz.o  %t2.o
+# RUN: llvm-objcopy --decompress-debug-sections %tzg.o %t3.o
+
+# Using redirects to avoid llvm-objdump from printing the filename.
+# RUN: llvm-objdump -s -section=.debug_str - < %t.o  > %t.txt
+# RUN: llvm-objdump -s -section=.debug_str - < %t2.o > %t2.txt
+# RUN: llvm-objdump -s -section=.debug_str - < %t3.o > %t3.txt
+
+# RUN: diff %t.txt %t2.txt
+# RUN: diff %t.txt %t3.txt
+
+# RUN: cmp %t.o %t.copy.o
+# RUN: cmp %tz.o %tz.copy.o
+# RUN: cmp %tzg.o %tzg.copy.o
diff --git a/llvm/tools/llvm-objcopy/ObjcopyOpts.td b/llvm/tools/llvm-objcopy/ObjcopyOpts.td
index 469f76d24d0..3a3c82c4573 100644
--- a/llvm/tools/llvm-objcopy/ObjcopyOpts.td
+++ b/llvm/tools/llvm-objcopy/ObjcopyOpts.td
@@ -23,6 +23,8 @@ def compress_debug_sections_eq : Joined<["--", "-"], "compress-debug-sections=">
                                  HelpText<"Compress DWARF debug sections using "
                                           "specified style. Supported styles: "
                                           "'zlib-gnu' and 'zlib'">;
+def decompress_debug_sections : Flag<["-", "--"], "decompress-debug-sections">,
+                                HelpText<"Decompress DWARF debug sections.">;
 def O : JoinedOrSeparate<["-"], "O">,
         Alias<output_target>;
 defm split_dwo : Eq<"split-dwo">,
diff --git a/llvm/tools/llvm-objcopy/Object.cpp b/llvm/tools/llvm-objcopy/Object.cpp
index e57a626fcad..ddf811a769b 100644
--- a/llvm/tools/llvm-objcopy/Object.cpp
+++ b/llvm/tools/llvm-objcopy/Object.cpp
@@ -136,6 +136,63 @@ void SectionWriter::visit(const OwnedDataSection &Sec) {
   std::copy(std::begin(Sec.Data), std::end(Sec.Data), Buf);
 }
 
+static const std::vector<uint8_t> ZlibGnuMagic = {'Z', 'L', 'I', 'B'};
+
+static bool isDataGnuCompressed(ArrayRef<uint8_t> Data) {
+  return Data.size() > ZlibGnuMagic.size() &&
+         std::equal(ZlibGnuMagic.begin(), ZlibGnuMagic.end(), Data.data());
+}
+
+template <class ELFT>
+static std::tuple<uint64_t, uint64_t>
+getDecompressedSizeAndAlignment(ArrayRef<uint8_t> Data) {
+  const bool IsGnuDebug = isDataGnuCompressed(Data);
+  const uint64_t DecompressedSize =
+      IsGnuDebug
+          ? support::endian::read64be(reinterpret_cast<const uint64_t *>(
+                Data.data() + ZlibGnuMagic.size()))
+          : reinterpret_cast<const Elf_Chdr_Impl<ELFT> *>(Data.data())->ch_size;
+  const uint64_t DecompressedAlign =
+      IsGnuDebug ? 1
+                 : reinterpret_cast<const Elf_Chdr_Impl<ELFT> *>(Data.data())
+                       ->ch_addralign;
+
+  return std::make_tuple(DecompressedSize, DecompressedAlign);
+}
+
+template <class ELFT>
+void ELFSectionWriter<ELFT>::visit(const DecompressedSection &Sec) {
+  uint8_t *Buf = Out.getBufferStart() + Sec.Offset;
+
+  if (!zlib::isAvailable()) {
+    std::copy(Sec.OriginalData.begin(), Sec.OriginalData.end(), Buf);
+    return;
+  }
+
+  const size_t DataOffset = isDataGnuCompressed(Sec.OriginalData)
+                                ? (ZlibGnuMagic.size() + sizeof(Sec.Size))
+                                : sizeof(Elf_Chdr_Impl<ELFT>);
+
+  StringRef CompressedContent(
+      reinterpret_cast<const char *>(Sec.OriginalData.data()) + DataOffset,
+      Sec.OriginalData.size() - DataOffset);
+
+  SmallVector<char, 128> DecompressedContent;
+  if (Error E = zlib::uncompress(CompressedContent, DecompressedContent,
+                                 static_cast<size_t>(Sec.Size)))
+    reportError(Sec.Name, std::move(E));
+
+  std::copy(DecompressedContent.begin(), DecompressedContent.end(), Buf);
+}
+
+void BinarySectionWriter::visit(const DecompressedSection &Sec) {
+  error("Cannot write compressed section '" + Sec.Name + "' ");
+}
+
+void DecompressedSection::accept(SectionVisitor &Visitor) const {
+  Visitor.visit(*this);
+}
+
 void OwnedDataSection::accept(SectionVisitor &Visitor) const {
   Visitor.visit(*this);
 }
@@ -206,6 +263,14 @@ CompressedSection::CompressedSection(const SectionBase &Sec,
   Align = 8;
 }
 
+CompressedSection::CompressedSection(ArrayRef<uint8_t> CompressedData,
+                                     uint64_t DecompressedSize,
+                                     uint64_t DecompressedAlign)
+    : CompressionType(DebugCompressionType::None),
+      DecompressedSize(DecompressedSize), DecompressedAlign(DecompressedAlign) {
+  OriginalData = CompressedData;
+}
+
 void CompressedSection::accept(SectionVisitor &Visitor) const {
   Visitor.visit(*this);
 }
@@ -969,10 +1034,20 @@ SectionBase &ELFBuilder<ELFT>::makeSection(const Elf_Shdr &Shdr) {
   }
   case SHT_NOBITS:
     return Obj.addSection<Section>(Data);
-  default:
+  default: {
     Data = unwrapOrError(ElfFile.getSectionContents(&Shdr));
+
+    if (isDataGnuCompressed(Data) || (Shdr.sh_flags & ELF::SHF_COMPRESSED)) {
+      uint64_t DecompressedSize, DecompressedAlign;
+      std::tie(DecompressedSize, DecompressedAlign) =
+          getDecompressedSizeAndAlignment<ELFT>(Data);
+      return Obj.addSection<CompressedSection>(Data, DecompressedSize,
+                                               DecompressedAlign);
+    }
+
     return Obj.addSection<Section>(Data);
   }
+  }
 }
 
 template <class ELFT> void ELFBuilder<ELFT>::readSectionHeaders() {
diff --git a/llvm/tools/llvm-objcopy/Object.h b/llvm/tools/llvm-objcopy/Object.h
index 42b2733dc1a..5fb03a5501e 100644
--- a/llvm/tools/llvm-objcopy/Object.h
+++ b/llvm/tools/llvm-objcopy/Object.h
@@ -41,6 +41,7 @@ class GnuDebugLinkSection;
 class GroupSection;
 class SectionIndexSection;
 class CompressedSection;
+class DecompressedSection;
 class Segment;
 class Object;
 struct Symbol;
@@ -89,6 +90,7 @@ public:
   virtual void visit(const GroupSection &Sec) = 0;
   virtual void visit(const SectionIndexSection &Sec) = 0;
   virtual void visit(const CompressedSection &Sec) = 0;
+  virtual void visit(const DecompressedSection &Sec) = 0;
 };
 
 class SectionWriter : public SectionVisitor {
@@ -108,6 +110,7 @@ public:
   virtual void visit(const GroupSection &Sec) override = 0;
   virtual void visit(const SectionIndexSection &Sec) override = 0;
   virtual void visit(const CompressedSection &Sec) override = 0;
+  virtual void visit(const DecompressedSection &Sec) override = 0;
 
   explicit SectionWriter(Buffer &Buf) : Out(Buf) {}
 };
@@ -127,6 +130,7 @@ public:
   void visit(const GroupSection &Sec) override;
   void visit(const SectionIndexSection &Sec) override;
   void visit(const CompressedSection &Sec) override;
+  void visit(const DecompressedSection &Sec) override;
 
   explicit ELFSectionWriter(Buffer &Buf) : SectionWriter(Buf) {}
 };
@@ -145,6 +149,7 @@ public:
   void visit(const GroupSection &Sec) override;
   void visit(const SectionIndexSection &Sec) override;
   void visit(const CompressedSection &Sec) override;
+  void visit(const DecompressedSection &Sec) override;
 
   explicit BinarySectionWriter(Buffer &Buf) : SectionWriter(Buf) {}
 };
@@ -370,6 +375,33 @@ class CompressedSection : public SectionBase {
 public:
   CompressedSection(const SectionBase &Sec,
                     DebugCompressionType CompressionType);
+  CompressedSection(ArrayRef<uint8_t> CompressedData, uint64_t DecompressedSize,
+                    uint64_t DecompressedAlign);
+
+  uint64_t getDecompressedSize() const { return DecompressedSize; }
+  uint64_t getDecompressedAlign() const { return DecompressedAlign; }
+
+  void accept(SectionVisitor &Visitor) const override;
+
+  static bool classof(const SectionBase *S) {
+    return (S->Flags & ELF::SHF_COMPRESSED) ||
+           (StringRef(S->Name).startswith(".zdebug"));
+  }
+};
+
+class DecompressedSection : public SectionBase {
+  MAKE_SEC_WRITER_FRIEND
+
+public:
+  explicit DecompressedSection(const CompressedSection &Sec)
+      : SectionBase(Sec) {
+    Size = Sec.getDecompressedSize();
+    Align = Sec.getDecompressedAlign();
+    Flags = (Flags & ~ELF::SHF_COMPRESSED);
+    if (StringRef(Name).startswith(".zdebug"))
+      Name = "." + Name.substr(2);
+  }
+
   void accept(SectionVisitor &Visitor) const override;
 };
 
diff --git a/llvm/tools/llvm-objcopy/llvm-objcopy.cpp b/llvm/tools/llvm-objcopy/llvm-objcopy.cpp
index ce2a1f1f259..345a11d8ba7 100644
--- a/llvm/tools/llvm-objcopy/llvm-objcopy.cpp
+++ b/llvm/tools/llvm-objcopy/llvm-objcopy.cpp
@@ -178,6 +178,7 @@ struct CopyConfig {
   bool StripSections = false;
   bool StripUnneeded = false;
   bool Weaken = false;
+  bool DecompressDebugSections = false;
   DebugCompressionType CompressionType = DebugCompressionType::None;
 };
 
@@ -430,33 +431,34 @@ static bool isCompressable(const SectionBase &Section) {
          Section.Name != ".gdb_index";
 }
 
-static void compressSections(const CopyConfig &Config, Object &Obj,
-                             SectionPred &RemovePred) {
-  SmallVector<SectionBase *, 13> ToCompress;
+static void replaceDebugSections(
+    const CopyConfig &Config, Object &Obj, SectionPred &RemovePred,
+    function_ref<bool(const SectionBase &)> shouldReplace,
+    function_ref<SectionBase *(const SectionBase *)> addSection) {
+  SmallVector<SectionBase *, 13> ToReplace;
   SmallVector<RelocationSection *, 13> RelocationSections;
   for (auto &Sec : Obj.sections()) {
     if (RelocationSection *R = dyn_cast<RelocationSection>(&Sec)) {
-      if (isCompressable(*R->getSection()))
+      if (shouldReplace(*R->getSection()))
         RelocationSections.push_back(R);
       continue;
     }
 
-    if (isCompressable(Sec))
-      ToCompress.push_back(&Sec);
+    if (shouldReplace(Sec))
+      ToReplace.push_back(&Sec);
   }
 
-  for (SectionBase *S : ToCompress) {
-    CompressedSection &CS =
-        Obj.addSection<CompressedSection>(*S, Config.CompressionType);
+  for (SectionBase *S : ToReplace) {
+    SectionBase *NewSection = addSection(S);
 
     for (RelocationSection *RS : RelocationSections) {
       if (RS->getSection() == S)
-        RS->setSection(&CS);
+        RS->setSection(NewSection);
     }
   }
 
-  RemovePred = [RemovePred](const SectionBase &Sec) {
-    return isCompressable(Sec) || RemovePred(Sec);
+  RemovePred = [shouldReplace, RemovePred](const SectionBase &Sec) {
+    return shouldReplace(Sec) || RemovePred(Sec);
   };
 }
 
@@ -672,7 +674,19 @@ static void handleArgs(const CopyConfig &Config, Object &Obj,
   }
 
   if (Config.CompressionType != DebugCompressionType::None)
-    compressSections(Config, Obj, RemovePred);
+    replaceDebugSections(Config, Obj, RemovePred, isCompressable,
+                         [&Config, &Obj](const SectionBase *S) {
+                           return &Obj.addSection<CompressedSection>(
+                               *S, Config.CompressionType);
+                         });
+  else if (Config.DecompressDebugSections)
+    replaceDebugSections(
+        Config, Obj, RemovePred,
+        [](const SectionBase &S) { return isa<CompressedSection>(&S); },
+        [&Obj](const SectionBase *S) {
+          auto CS = cast<CompressedSection>(S);
+          return &Obj.addSection<DecompressedSection>(*CS);
+        });
 
   Obj.removeSections(RemovePred);
 
@@ -983,6 +997,8 @@ static DriverConfig parseObjcopyOptions(ArrayRef<const char *> ArgsArr) {
   Config.DiscardAll = InputArgs.hasArg(OBJCOPY_discard_all);
   Config.OnlyKeepDebug = InputArgs.hasArg(OBJCOPY_only_keep_debug);
   Config.KeepFileSymbols = InputArgs.hasArg(OBJCOPY_keep_file_symbols);
+  Config.DecompressDebugSections =
+      InputArgs.hasArg(OBJCOPY_decompress_debug_sections);
   for (auto Arg : InputArgs.filtered(OBJCOPY_localize_symbol))
     Config.SymbolsToLocalize.push_back(Arg->getValue());
   for (auto Arg : InputArgs.filtered(OBJCOPY_keep_global_symbol))
@@ -1002,6 +1018,15 @@ static DriverConfig parseObjcopyOptions(ArrayRef<const char *> ArgsArr) {
 
   DriverConfig DC;
   DC.CopyConfigs.push_back(std::move(Config));
+  if (Config.DecompressDebugSections &&
+      Config.CompressionType != DebugCompressionType::None) {
+    error("Cannot specify --compress-debug-sections at the same time as "
+          "--decompress-debug-sections at the same time");
+  }
+
+  if (Config.DecompressDebugSections && !zlib::isAvailable())
+    error("LLVM was not compiled with LLVM_ENABLE_ZLIB: cannot decompress.");
+
   return DC;
 }