Add support for zstd compressed debug sections (#318)

src/bloaty.cc

@@ -25,6 +25,7 @@ typedef size_t z_size_t;
 #include <signal.h>
 #include <stdlib.h>
 #include <zlib.h>
+#include <zstd.h>
 
 #include <atomic>
 #include <cmath>
@@ -1401,6 +1402,34 @@ absl::string_view RangeSink::ZlibDecompress(absl::string_view data,
   return sv;
 }
 
+absl::string_view RangeSink::ZstdDecompress(absl::string_view data,
+                                            uint64_t uncompressed_size) {
+  if (!arena_) {
+    THROW("This range sink isn't prepared to zlib decompress.");
+  }
+  uint64_t mb = 1 << 20;
+  // Limit for uncompressed size is 30x the compressed size + 128MB.
+  if (uncompressed_size >
+      static_cast<uint64_t>(data.size()) * 30 + (128 * mb)) {
+    fprintf(stderr,
+            "warning: ignoring compressed debug data, implausible uncompressed "
+            "size (compressed: %zu, uncompressed: %" PRIu64 ")\n",
+            data.size(), uncompressed_size);
+    return absl::string_view();
+  }
+  char* dbuf =
+      arena_->proto2::Arena::CreateArray<char>(
+          arena_, uncompressed_size);
+  size_t zstd_len = uncompressed_size;
+  size_t decompress_len =
+      ZSTD_decompress(dbuf, uncompressed_size, data.data(), data.size());
+  if (ZSTD_isError(decompress_len) || decompress_len != zstd_len) {
+    THROW("Error decompressing debug info");
+  }
+  string_view sv(dbuf, zstd_len);
+  return sv;
+}
+
 // ThreadSafeIterIndex /////////////////////////////////////////////////////////
 
 class ThreadSafeIterIndex {

src/bloaty.h

@@ -212,6 +212,12 @@ class RangeSink {
   absl::string_view ZlibDecompress(absl::string_view contents,
                                    uint64_t uncompressed_size);
 
+  // Decompresses zstd-formatted data and returns the decompressed data.
+  // Since the decompressed data is not actually part of the file, any
+  // Add*Range() calls to this region will be no-ops.
+  absl::string_view ZstdDecompress(absl::string_view contents,
+                                   uint64_t uncompressed_size);
+
   static constexpr uint64_t kUnknownSize = RangeMap::kUnknownSize;
 
  private:

src/elf.cc

@@ -29,6 +29,10 @@
 
 using absl::string_view;
 
+#ifndef ELFCOMPRESS_ZSTD
+#define ELFCOMPRESS_ZSTD 2
+#endif
+
 namespace bloaty {
 
 namespace {
@@ -58,6 +62,12 @@ void AdvancePastStruct(string_view* data) {
   *data = data->substr(sizeof(T));
 }
 
+enum class SectionCompressionType {
+  Zlib,
+  Zstd,
+  None
+};
+
 // ElfFile /////////////////////////////////////////////////////////////////////
 
 // For parsing the pieces we need out of an ELF file (.o, .so, and binaries).
@@ -1196,16 +1206,24 @@ void ReadDWARFSections(const InputFile &file, dwarf::File *dwarf,
     string_view name = section.GetName();
     string_view contents = section.contents();
     uint64_t uncompressed_size = 0;
+    SectionCompressionType compression_type = SectionCompressionType::None;
 
     if (section.header().sh_flags & SHF_COMPRESSED) {
       // Standard ELF section compression, produced when you link with
       //   --compress-debug-sections=zlib-gabi
       Elf64_Chdr chdr;
       absl::string_view range;
       elf.ReadStruct<Elf32_Chdr>(contents, 0, ChdrMunger(), &range, &chdr);
-      if (chdr.ch_type != ELFCOMPRESS_ZLIB) {
-        // Unknown compression format.
-        continue;
+      switch (chdr.ch_type) {
+        case ELFCOMPRESS_ZLIB:
+          compression_type = SectionCompressionType::Zlib;
+          break;
+        case ELFCOMPRESS_ZSTD:
+          compression_type = SectionCompressionType::Zstd;
+          break;
+        default:
+          // Unknown compression format.
+          continue;
       }
       uncompressed_size = chdr.ch_size;
       contents.remove_prefix(range.size());
@@ -1220,6 +1238,7 @@ void ReadDWARFSections(const InputFile &file, dwarf::File *dwarf,
       if (ReadBytes(4, &contents) != "ZLIB") {
         continue;  // Bad compression header.
       }
+      compression_type = SectionCompressionType::Zlib;
       uncompressed_size = ReadBigEndian<uint64_t>(&contents);
     }
 
@@ -1230,10 +1249,15 @@ void ReadDWARFSections(const InputFile &file, dwarf::File *dwarf,
     }
 
     if (string_view* member = dwarf->GetFieldByName(name)) {
-      if (uncompressed_size) {
-        *member = sink->ZlibDecompress(contents, uncompressed_size);
-      } else {
-        *member = section.contents();
+      switch (compression_type) {
+        case SectionCompressionType::Zlib:
+          *member = sink->ZlibDecompress(contents, uncompressed_size);
+          break;
+        case SectionCompressionType::Zstd:
+          *member = sink->ZstdDecompress(contents, uncompressed_size);
+          break;
+        default:
+          *member = section.contents();
       }
     }
   }