Locate zlib stream from a buffer am: 75a7f2c89d

am: 17434e1dcb

Change-Id: I6ae461a61f1397d33965bcb43340898f82ab4c3d

4 files changed, 111 insertions, 71 deletions

diff --git a/src/include/puffin/utils.h b/src/include/puffin/utils.h
index 4ff9c16..65f8e63 100644
--- a/src/include/puffin/utils.h
+++ b/src/include/puffin/utils.h
@@ -29,14 +29,14 @@ PUFFIN_EXPORT std::string ExtentsToString(const T& extents) {
   return str;
 }
 
-// Locates deflate buffer locations for a set of zlib buffers |zlibs| in
-// |src|. It performs by removing header and footer bytes from the zlib stream.
-bool LocateDeflatesInZlibBlocks(const UniqueStreamPtr& src,
-                                const std::vector<ByteExtent>& zlibs,
-                                std::vector<BitExtent>* deflates);
+// Locates deflate locations for a zlib buffer |data|. It locates by removing
+// header and footer bytes from the zlib stream.
+bool LocateDeflatesInZlib(const Buffer& data,
+                          std::vector<ByteExtent>* deflate_blocks);
 
 // Similar to the function above, except that it accepts the file path to the
-// source.
+// source and a list of zlib blocks and returns the deflate addresses in bit
+// extents.
 PUFFIN_EXPORT
 bool LocateDeflatesInZlibBlocks(const std::string& file_path,
                                 const std::vector<ByteExtent>& zlibs,
diff --git a/src/main.cc b/src/main.cc
index 6ff7bbe..c43320f 100644
--- a/src/main.cc
+++ b/src/main.cc
@@ -144,10 +144,11 @@ int main(int argc, char** argv) {
     if (FLAGS_src_file_type == "deflate") {
       src_deflates_byte = {ByteExtent(0, stream_size)};
     } else if (FLAGS_src_file_type == "zlib") {
-      std::vector<ByteExtent> zlibs = {ByteExtent(0, stream_size)};
-      TEST_AND_RETURN_VALUE(puffin::LocateDeflatesInZlibBlocks(
-                                src_stream, zlibs, &src_deflates_bit),
+      Buffer src_data(stream_size);
+      TEST_AND_RETURN_VALUE(src_stream->Read(src_data.data(), src_data.size()),
                             -1);
+      TEST_AND_RETURN_VALUE(
+          puffin::LocateDeflatesInZlib(src_data, &src_deflates_byte), -1);
     } else if (FLAGS_src_file_type == "gzip") {
       Buffer src_data(stream_size);
       TEST_AND_RETURN_VALUE(src_stream->Read(src_data.data(), src_data.size()),
diff --git a/src/utils.cc b/src/utils.cc
index ba10570..2b41220 100644
--- a/src/utils.cc
+++ b/src/utils.cc
@@ -88,56 +88,38 @@ size_t BytesInByteExtents(const vector<ByteExtent>& extents) {
 }
 
 // This function uses RFC1950 (https://www.ietf.org/rfc/rfc1950.txt) for the
-// definition of a zlib stream.
-bool LocateDeflatesInZlibBlocks(const UniqueStreamPtr& src,
-                                const vector<ByteExtent>& zlibs,
-                                vector<BitExtent>* deflates) {
-  for (auto& zlib : zlibs) {
-    TEST_AND_RETURN_FALSE(src->Seek(zlib.offset));
-    uint16_t zlib_header;
-    TEST_AND_RETURN_FALSE(src->Read(&zlib_header, 2));
-    BufferBitReader bit_reader(reinterpret_cast<uint8_t*>(&zlib_header), 2);
-
-    TEST_AND_RETURN_FALSE(bit_reader.CacheBits(8));
-    auto cmf = bit_reader.ReadBits(8);
-    auto cm = bit_reader.ReadBits(4);
-    if (cm != 8 && cm != 15) {
-      LOG(ERROR) << "Invalid compression method! cm: " << cm;
-      return false;
-    }
-    bit_reader.DropBits(4);
-    auto cinfo = bit_reader.ReadBits(4);
-    if (cinfo > 7) {
-      LOG(ERROR) << "cinfo greater than 7 is not allowed in deflate";
-      return false;
-    }
-    bit_reader.DropBits(4);
-
-    TEST_AND_RETURN_FALSE(bit_reader.CacheBits(8));
-    auto flg = bit_reader.ReadBits(8);
-    if (((cmf << 8) + flg) % 31) {
-      LOG(ERROR) << "Invalid zlib header on offset: " << zlib.offset;
-      return false;
-    }
-    bit_reader.ReadBits(5);  // FCHECK
-    bit_reader.DropBits(5);
-
-    auto fdict = bit_reader.ReadBits(1);
-    bit_reader.DropBits(1);
-
-    bit_reader.ReadBits(2);  // FLEVEL
-    bit_reader.DropBits(2);
-
-    auto header_len = 2;
-    if (fdict) {
-      TEST_AND_RETURN_FALSE(bit_reader.CacheBits(32));
-      bit_reader.DropBits(32);
-      header_len += 4;
-    }
-
-    ByteExtent deflate(zlib.offset + header_len, zlib.length - header_len - 4);
-    TEST_AND_RETURN_FALSE(FindDeflateSubBlocks(src, {deflate}, deflates));
+// definition of a zlib stream.  For finding the deflate blocks, we relying on
+// the proper size of the zlib stream in |data|. Basically the size of the zlib
+// stream should be known before hand. Otherwise we need to parse the stream and
+// find the location of compressed blocks using CalculateSizeOfDeflateBlock().
+bool LocateDeflatesInZlib(const Buffer& data,
+                          std::vector<ByteExtent>* deflate_blocks) {
+  // A zlib stream has the following format:
+  // 0           1     compression method and flag
+  // 1           1     flag
+  // 2           4     preset dictionary (optional)
+  // 2 or 6      n     compressed data
+  // n+(2 or 6)  4     Adler-32 checksum
+  TEST_AND_RETURN_FALSE(data.size() >= 6 + 4);  // Header + Footer
+  uint16_t cmf = data[0];
+  auto compression_method = cmf & 0x0F;
+  // For deflate compression_method should be 8.
+  TEST_AND_RETURN_FALSE(compression_method == 8);
+
+  auto cinfo = (cmf & 0xF0) >> 4;
+  // Value greater than 7 is not allowed in deflate.
+  TEST_AND_RETURN_FALSE(cinfo <= 7);
+
+  auto flag = data[1];
+  TEST_AND_RETURN_FALSE(((cmf << 8) + flag) % 31 == 0);
+
+  size_t header_len = 2;
+  if (flag & 0x20) {
+    header_len += 4;  // 4 bytes for the preset dictionary.
   }
+
+  // 4 is for ADLER32.
+  deflate_blocks->emplace_back(header_len, data.size() - header_len - 4);
   return true;
 }
 
@@ -173,7 +155,27 @@ bool LocateDeflatesInZlibBlocks(const string& file_path,
                                 vector<BitExtent>* deflates) {
   auto src = FileStream::Open(file_path, true, false);
   TEST_AND_RETURN_FALSE(src);
-  return LocateDeflatesInZlibBlocks(src, zlibs, deflates);
+
+  Buffer buffer;
+  for (auto& zlib : zlibs) {
+    buffer.resize(zlib.length);
+    TEST_AND_RETURN_FALSE(src->Seek(zlib.offset));
+    TEST_AND_RETURN_FALSE(src->Read(buffer.data(), buffer.size()));
+
+    vector<ByteExtent> deflate_blocks;
+    TEST_AND_RETURN_FALSE(LocateDeflatesInZlib(buffer, &deflate_blocks));
+
+    vector<BitExtent> deflate_subblocks;
+    auto zlib_blc_src = MemoryStream::CreateForRead(buffer);
+    TEST_AND_RETURN_FALSE(
+        FindDeflateSubBlocks(zlib_blc_src, deflate_blocks, &deflate_subblocks));
+
+    // Relocated based on the offset of the zlib.
+    for (const auto& def : deflate_subblocks) {
+      deflates->emplace_back(zlib.offset * 8 + def.offset, def.length);
+    }
+  }
+  return true;
 }
 
 // For more information about gzip format, refer to RFC 1952 located at:
diff --git a/src/utils_unittest.cc b/src/utils_unittest.cc
index 55f2fa1..8d972d2 100644
--- a/src/utils_unittest.cc
+++ b/src/utils_unittest.cc
@@ -2,10 +2,13 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include <unistd.h>
+
 #include <vector>
 
 #include "gtest/gtest.h"
 
+#include "puffin/src/file_stream.h"
 #include "puffin/src/include/puffin/common.h"
 #include "puffin/src/include/puffin/utils.h"
 #include "puffin/src/memory_stream.h"
@@ -13,6 +16,7 @@
 
 namespace puffin {
 
+using std::string;
 using std::vector;
 
 namespace {
@@ -73,12 +77,25 @@ const uint8_t kGzipEntryWithExtraField[] = {
     0x34, 0x32, 0x36, 0x31, 0x35, 0x33, 0xb7, 0xb0, 0xe4, 0x02, 0x00, 0xd1,
     0xe5, 0x76, 0x40, 0x0b, 0x00, 0x00, 0x00};
 
+// echo "0123456789" | zlib-flate -compress |
+// hexdump -v -e '12/1 "0x%02x, " "\n"'
+const uint8_t kZlibEntry[] = {
+    0x78, 0x9c, 0x33, 0x30, 0x34, 0x32, 0x36, 0x31, 0x35, 0x33, 0xb7, 0xb0,
+    0xe4, 0x02, 0x00, 0x0d, 0x17, 0x02, 0x18};
+
 void FindDeflatesInZlibBlocks(const Buffer& src,
                               const vector<ByteExtent>& zlibs,
                               const vector<BitExtent>& deflates) {
-  auto src_stream = MemoryStream::CreateForRead(src);
+  string tmp_file;
+  ASSERT_TRUE(MakeTempFile(&tmp_file, nullptr));
+  ScopedPathUnlinker unlinker(tmp_file);
+  auto src_stream = FileStream::Open(tmp_file, false, true);
+  ASSERT_TRUE(src_stream);
+  ASSERT_TRUE(src_stream->Write(src.data(), src.size()));
+  ASSERT_TRUE(src_stream->Close());
+
   vector<BitExtent> deflates_out;
-  ASSERT_TRUE(LocateDeflatesInZlibBlocks(src_stream, zlibs, &deflates_out));
+  ASSERT_TRUE(LocateDeflatesInZlibBlocks(tmp_file, zlibs, &deflates_out));
   ASSERT_EQ(deflates, deflates_out);
 }
 
@@ -95,15 +112,7 @@ void CheckFindPuffLocation(const Buffer& compressed,
 }
 }  // namespace
 
-TEST(UtilsTest, LocateDeflatesInZlibsTest) {
-  Buffer empty;
-  vector<ByteExtent> empty_zlibs;
-  vector<BitExtent> empty_deflates;
-  FindDeflatesInZlibBlocks(empty, empty_zlibs, empty_deflates);
-}
-
 // Test Simple Puffing of the source.
-
 TEST(UtilsTest, FindPuffLocations1Test) {
   CheckFindPuffLocation(kDeflates8, kSubblockDeflateExtents8, kPuffExtents8,
                         kPuffs8.size());
@@ -114,8 +123,36 @@ TEST(UtilsTest, FindPuffLocations2Test) {
                         kPuffs9.size());
 }
 
-// TODO(ahassani): Test a proper zlib format.
-// TODO(ahassani): Test zlib format with wrong header.
+TEST(UtilsTest, LocateDeflatesInZlib) {
+  Buffer zlib_data(kZlibEntry, std::end(kZlibEntry));
+  vector<ByteExtent> deflates;
+  EXPECT_TRUE(LocateDeflatesInZlib(zlib_data, &deflates));
+  EXPECT_EQ(static_cast<size_t>(1), deflates.size());
+  EXPECT_EQ(ByteExtent(2, 13), deflates[0]);
+}
+
+TEST(UtilsTest, LocateDeflatesInEmptyZlib) {
+  Buffer empty;
+  vector<ByteExtent> empty_zlibs;
+  vector<BitExtent> empty_deflates;
+  FindDeflatesInZlibBlocks(empty, empty_zlibs, empty_deflates);
+}
+
+TEST(UtilsTest, LocateDeflatesInZlibWithInvalidFields) {
+  Buffer zlib_data(kZlibEntry, std::end(kZlibEntry));
+  auto cmf = zlib_data[0];
+  auto flag = zlib_data[1];
+
+  vector<ByteExtent> deflates;
+  zlib_data[0] = cmf & 0xF0;
+  EXPECT_FALSE(LocateDeflatesInZlib(zlib_data, &deflates));
+  zlib_data[0] = cmf | (8 << 4);
+  EXPECT_FALSE(LocateDeflatesInZlib(zlib_data, &deflates));
+  zlib_data[0] = cmf;  // Correct it.
+
+  zlib_data[1] = flag & 0xF0;
+  EXPECT_FALSE(LocateDeflatesInZlib(zlib_data, &deflates));
+}
 
 TEST(UtilsTest, LocateDeflatesInZipArchiveSmoke) {
   Buffer zip_entries(kZipEntries, std::end(kZipEntries));