Locate zlib stream from a buffer

Previously we were locating deflates inside zlib streams based on the given
location of zlib streams. That was not consistent with gzip and zip
streams. This patch changes the internal API and implementation to makes it more
consistent.

Bug: 73487244
Test: unittests
Test: 'brillo_update_payload generate'

Change-Id: Iab0172fe102ab83944f8aaf90d68d9f08f36ae40

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));