Snap for 4620899 from 52ab9662607ec56d7396cbc2e2ff01fb40530ffc to pi-release

Change-Id: I17ecd3ec26e9fc40af483a6b52e1d661cd9677fd

11 files changed, 396 insertions, 108 deletions

diff --git a/Android.bp b/Android.bp
index d95bc49..5d99dcb 100644
--- a/Android.bp
+++ b/Android.bp
@@ -98,6 +98,7 @@ cc_test {
         "src/sample_generator.cc",
         "src/stream_unittest.cc",
         "src/testrunner.cc",
+        "src/unittest_common.cc",
         "src/utils_unittest.cc",
     ],
     shared_libs: [
diff --git a/puffin.gyp b/puffin.gyp
index 33f8688..4903e00 100644
--- a/puffin.gyp
+++ b/puffin.gyp
@@ -171,6 +171,7 @@
             'src/puff_io_unittest.cc',
             'src/puffin_unittest.cc',
             'src/stream_unittest.cc',
+            'src/unittest_common.cc',
             'src/utils_unittest.cc',
           ],
         },
diff --git a/scripts/test_corpus.py b/scripts/test_corpus.py
new file mode 100755
index 0000000..e62b680
--- /dev/null
+++ b/scripts/test_corpus.py
@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+#
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+#
+
+"""A tool for running puffin tests in a corpus of deflate compressed files."""
+
+import argparse
+import filecmp
+import logging
+import os
+import subprocess
+import sys
+import tempfile
+
+_PUFFHUFF = 'puffhuff'
+_PUFFDIFF = 'puffdiff'
+TESTS = (_PUFFHUFF, _PUFFDIFF)
+
+
+class Error(Exception):
+  """Puffin general processing error."""
+
+
+def ParseArguments(argv):
+  """Parses and Validates command line arguments.
+
+  Args:
+    argv: command line arguments to parse.
+
+  Returns:
+    The arguments list.
+  """
+  parser = argparse.ArgumentParser()
+
+  parser.add_argument('corpus', metavar='CORPUS',
+                      help='A corpus directory containing compressed files')
+  parser.add_argument('-d', '--disabled_tests', default=(), metavar='',
+                      nargs='*',
+                      help=('Space separated list of tests to disable. '
+                            'Allowed options include: ' + ', '.join(TESTS)),
+                      choices=TESTS)
+  parser.add_argument('--cache_size', type=int, metavar='SIZE',
+                      help='The size (in bytes) of the cache for puffpatch '
+                      'operations.')
+  parser.add_argument('--debug', action='store_true',
+                      help='Turns on verbosity.')
+
+  # Parse command-line arguments.
+  args = parser.parse_args(argv)
+
+  if not os.path.isdir(args.corpus):
+    raise Error('Corpus directory {} is non-existent or inaccesible'
+                .format(args.corpus))
+  return args
+
+
+def main(argv):
+  """The main function."""
+  args = ParseArguments(argv[1:])
+
+  if args.debug:
+    logging.getLogger().setLevel(logging.DEBUG)
+
+  # Construct list of appropriate files.
+  files = list(filter(os.path.isfile, [os.path.join(args.corpus, f)
+                                       for f in os.listdir(args.corpus)]))
+
+  # For each file in corpus run puffhuff.
+  if _PUFFHUFF not in args.disabled_tests:
+    for src in files:
+      with tempfile.NamedTemporaryFile() as tgt_file:
+
+        operation = 'puffhuff'
+        logging.debug('Running %s on %s', operation, src)
+        cmd = ['puffin',
+               '--operation={}'.format(operation),
+               '--src_file={}'.format(src),
+               '--dst_file={}'.format(tgt_file.name)]
+        if subprocess.call(cmd) != 0:
+          raise Error('Puffin failed to do {} command: {}'
+                      .format(operation, cmd))
+
+        if not filecmp.cmp(src, tgt_file.name):
+          raise Error('The generated file {} is not equivalent to the original '
+                      'file {} after {} operation.'
+                      .format(tgt_file.name, src, operation))
+
+  if _PUFFDIFF not in args.disabled_tests:
+    # Run puffdiff and puffpatch for each pairs of files in the corpus.
+    for src in files:
+      for tgt in files:
+        with tempfile.NamedTemporaryFile() as patch, \
+             tempfile.NamedTemporaryFile() as new_tgt:
+
+          operation = 'puffdiff'
+          logging.debug('Running %s on %s (%d) and %s (%d)',
+                        operation,
+                        os.path.basename(src), os.stat(src).st_size,
+                        os.path.basename(tgt), os.stat(tgt).st_size)
+          cmd = ['puffin',
+                 '--operation={}'.format(operation),
+                 '--src_file={}'.format(src),
+                 '--dst_file={}'.format(tgt),
+                 '--patch_file={}'.format(patch.name)]
+
+          # Running the puffdiff operation
+          if subprocess.call(cmd) != 0:
+            raise Error('Puffin failed to do {} command: {}'
+                        .format(operation, cmd))
+
+          logging.debug('Patch size is: %d', os.stat(patch.name).st_size)
+
+          operation = 'puffpatch'
+          logging.debug('Running %s on src file %s and patch %s',
+                        operation, os.path.basename(src), patch.name)
+          cmd = ['puffin',
+                 '--operation={}'.format(operation),
+                 '--src_file={}'.format(src),
+                 '--dst_file={}'.format(new_tgt.name),
+                 '--patch_file={}'.format(patch.name)]
+          if args.cache_size:
+            cmd += ['--cache_size={}'.format(args.cache_size)]
+
+          # Running the puffpatch operation
+          if subprocess.call(cmd) != 0:
+            raise Error('Puffin failed to do {} command: {}'
+                        .format(operation, cmd))
+
+          if not filecmp.cmp(tgt, new_tgt.name):
+            raise Error('The generated file {} is not equivalent to the '
+                        'original file {} after puffpatch operation.'
+                        .format(new_tgt.name, tgt))
+
+  return 0
+
+
+if __name__ == '__main__':
+  sys.exit(main(sys.argv))
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..dd57c05 100644
--- a/src/main.cc
+++ b/src/main.cc
@@ -63,6 +63,97 @@ vector<T> StringToExtents(const string& str) {
 
 const size_t kDefaultPuffCacheSize = 50 * 1024 * 1024;  // 50 MB
 
+// An enum representing the type of compressed files.
+enum class FileType { kDeflate, kZlib, kGzip, kZip, kRaw, kUnknown };
+
+// Returns a file type based on the input string |file_type| (normally the final
+// extension of the file).
+FileType StringToFileType(const string& file_type) {
+  if (file_type == "raw") {
+    return FileType::kRaw;
+  }
+  if (file_type == "deflate") {
+    return FileType::kDeflate;
+  } else if (file_type == "zlib") {
+    return FileType::kZlib;
+  } else if (file_type == "gzip" || file_type == "gz" || file_type == "tgz") {
+    return FileType::kGzip;
+  } else if (file_type == "zip" || file_type == "apk" || file_type == "jar") {
+    return FileType::kZip;
+  }
+  return FileType::kUnknown;
+}
+
+// Finds the location of deflates in |stream|. If |file_type_to_override| is
+// non-empty, it infers the file type based on that, otherwise, it infers the
+// file type based on the final extension of |file_name|. It returns false if
+// file type cannot be inferred from any of the input arguments. |deflates|
+// is filled with byte-aligned location of deflates.
+bool LocateDeflatesBasedOnFileType(const UniqueStreamPtr& stream,
+                                   const string& file_name,
+                                   const string& file_type_to_override,
+                                   vector<ByteExtent>* deflates) {
+  auto file_type = FileType::kUnknown;
+
+  auto last_dot = file_name.find_last_of(".");
+  if (last_dot == string::npos) {
+    // Could not find a dot so we assume there is no extension.
+    return false;
+  }
+  auto extension = file_name.substr(last_dot + 1);
+  file_type = StringToFileType(extension);
+
+  if (!file_type_to_override.empty()) {
+    auto override_file_type = StringToFileType(file_type_to_override);
+    if (override_file_type == FileType::kUnknown) {
+      LOG(ERROR) << "Overriden file type " << file_type_to_override
+                 << " does not exist.";
+      return false;
+    }
+    if (file_type != FileType::kUnknown && file_type != override_file_type) {
+      LOG(WARNING) << "Based on the file name, the file type is " << extension
+                   << ", But the overriden file type is "
+                   << file_type_to_override << ". Is this intentional?";
+    }
+    file_type = override_file_type;
+  }
+
+  if (file_type == FileType::kRaw) {
+    // Do not need to populate |deflates|.
+    return true;
+  }
+
+  size_t stream_size;
+  TEST_AND_RETURN_FALSE(stream->GetSize(&stream_size));
+  if (file_type == FileType::kDeflate) {
+    // Assume the whole stream is a deflate block.
+    *deflates = {ByteExtent(0, stream_size)};
+    return true;
+  }
+
+  Buffer data(stream_size);
+  TEST_AND_RETURN_FALSE(stream->Read(data.data(), data.size()));
+  switch (file_type) {
+    case FileType::kZlib:
+      TEST_AND_RETURN_FALSE(puffin::LocateDeflatesInZlib(data, deflates));
+      break;
+    case FileType::kGzip:
+      TEST_AND_RETURN_FALSE(puffin::LocateDeflatesInGzip(data, deflates));
+      break;
+    case FileType::kZip:
+      TEST_AND_RETURN_FALSE(puffin::LocateDeflatesInZipArchive(data, deflates));
+      break;
+    default:
+      LOG(ERROR) << "Unknown file type: (" << file_type_to_override << ") nor ("
+                 << extension << ").";
+      return false;
+  }
+  // Return the stream to its zero offset in case we used it.
+  TEST_AND_RETURN_FALSE(stream->Seek(0));
+
+  return true;
+}
+
 }  // namespace
 
 #define SETUP_FLAGS                                                        \
@@ -95,6 +186,8 @@ const size_t kDefaultPuffCacheSize = 50 * 1024 * 1024;  // 50 MB
   DEFINE_string(src_file_type, "",                                         \
                 "Type of the input source file: deflate, gzip, "           \
                 "zlib or zip");                                            \
+  DEFINE_string(dst_file_type, "",                                         \
+                "Same as src_file_type but for the target file");          \
   DEFINE_bool(verbose, false,                                              \
               "Logs all the given parameters including internally "        \
               "generated ones");                                           \
@@ -136,41 +229,12 @@ int main(int argc, char** argv) {
     TEST_AND_RETURN_VALUE(src_stream, -1);
   }
 
-  if (!FLAGS_src_file_type.empty()) {
+  if (FLAGS_operation == "puff" || FLAGS_operation == "puffhuff") {
     TEST_AND_RETURN_VALUE(
-        FLAGS_operation == "puff" || FLAGS_operation == "puffhuff", -1);
-    size_t stream_size;
-    TEST_AND_RETURN_VALUE(src_stream->GetSize(&stream_size), -1);
-    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),
-                            -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()),
-                            -1);
-      TEST_AND_RETURN_VALUE(
-          puffin::LocateDeflatesInGzip(src_data, &src_deflates_byte), -1);
-    } else if (FLAGS_src_file_type == "zip") {
-      Buffer src_data(stream_size);
-      TEST_AND_RETURN_VALUE(src_stream->Read(src_data.data(), src_data.size()),
-                            -1);
-      TEST_AND_RETURN_VALUE(
-          puffin::LocateDeflatesInZipArchive(src_data, &src_deflates_byte), -1);
-    } else {
-      LOG(ERROR) << "Unknown file type: " << FLAGS_src_file_type;
-      return -1;
-    }
-  }
-
-  // Return the stream to its zero offset in case we used it.
-  TEST_AND_RETURN_VALUE(src_stream->Seek(0), -1);
+        LocateDeflatesBasedOnFileType(src_stream, FLAGS_src_file,
+                                      FLAGS_src_file_type, &src_deflates_byte),
+        -1);
 
-  if (FLAGS_operation == "puff" || FLAGS_operation == "puffhuff") {
-    auto puffer = std::make_shared<Puffer>();
     if (src_deflates_bit.empty() && src_deflates_byte.empty()) {
       LOG(WARNING) << "You should pass source deflates, is this intentional?";
     }
@@ -187,7 +251,7 @@ int main(int argc, char** argv) {
 
     auto dst_stream = FileStream::Open(FLAGS_dst_file, false, true);
     TEST_AND_RETURN_VALUE(dst_stream, -1);
-
+    auto puffer = std::make_shared<Puffer>();
     auto reader =
         PuffinStream::CreateForPuff(std::move(src_stream), puffer,
                                     dst_puff_size, src_deflates_bit, dst_puffs);
@@ -254,14 +318,24 @@ int main(int argc, char** argv) {
       bytes_read += read_size;
     }
   } else if (FLAGS_operation == "puffdiff") {
+    auto dst_stream = FileStream::Open(FLAGS_dst_file, true, false);
+    TEST_AND_RETURN_VALUE(dst_stream, -1);
+
+    TEST_AND_RETURN_VALUE(
+        LocateDeflatesBasedOnFileType(src_stream, FLAGS_src_file,
+                                      FLAGS_src_file_type, &src_deflates_byte),
+        -1);
+    TEST_AND_RETURN_VALUE(
+        LocateDeflatesBasedOnFileType(dst_stream, FLAGS_dst_file,
+                                      FLAGS_dst_file_type, &dst_deflates_byte),
+        -1);
+
     if (src_deflates_bit.empty() && src_deflates_byte.empty()) {
       LOG(WARNING) << "You should pass source deflates, is this intentional?";
     }
     if (dst_deflates_bit.empty() && dst_deflates_byte.empty()) {
       LOG(WARNING) << "You should pass target deflates, is this intentional?";
     }
-    auto dst_stream = FileStream::Open(FLAGS_dst_file, true, false);
-    TEST_AND_RETURN_VALUE(dst_stream, -1);
     if (!dst_extents.empty()) {
       dst_stream =
           ExtentStream::CreateForWrite(std::move(dst_stream), dst_extents);
diff --git a/src/patching_unittest.cc b/src/patching_unittest.cc
index d478262..d6327cf 100644
--- a/src/patching_unittest.cc
+++ b/src/patching_unittest.cc
@@ -30,7 +30,8 @@ void TestPatching(const Buffer& src_buf,
                   const vector<BitExtent>& dst_deflates,
                   const Buffer patch) {
   Buffer patch_out;
-  string patch_path = "/tmp/patch.tmp";
+  string patch_path;
+  ASSERT_TRUE(MakeTempFile(&patch_path, nullptr));
   ScopedPathUnlinker scoped_unlinker(patch_path);
   ASSERT_TRUE(PuffDiff(src_buf, dst_buf, src_deflates, dst_deflates, patch_path,
                        &patch_out));
diff --git a/src/puffin_stream.cc b/src/puffin_stream.cc
index e26e2c9..3fb54a1 100644
--- a/src/puffin_stream.cc
+++ b/src/puffin_stream.cc
@@ -334,7 +334,7 @@ bool PuffinStream::Write(const void* buffer, size_t length) {
   auto bytes = static_cast<const uint8_t*>(buffer);
   size_t bytes_wrote = 0;
   while (bytes_wrote < length) {
-    if (deflate_bit_pos_ < (cur_deflate_->offset & ~7u)) {
+    if (deflate_bit_pos_ < (cur_deflate_->offset & ~7ull)) {
       // Between two puffs or before the first puff. We know that we are
       // starting from the byte boundary because we have already processed the
       // non-deflate bits of the last byte of the last deflate. Here we don't
@@ -397,7 +397,7 @@ bool PuffinStream::Write(const void* buffer, size_t length) {
         if (extra_byte_ == 1) {
           deflate_buffer_->data()[bytes_to_write - 1] |=
               puff_buffer_->data()[cur_puff_->length] << (deflate_bit_pos_ & 7);
-          deflate_bit_pos_ = (deflate_bit_pos_ + 7) & ~7u;
+          deflate_bit_pos_ = (deflate_bit_pos_ + 7) & ~7ull;
         } else if ((deflate_bit_pos_ & 7) != 0) {
           // This happens if current and next deflate finish and end on the same
           // byte, then we cannot write into output until we have huffed the
@@ -436,7 +436,7 @@ bool PuffinStream::SetExtraByte() {
     return true;
   }
   size_t end_bit = cur_deflate_->offset + cur_deflate_->length;
-  if ((end_bit & 7) && ((end_bit + 7) & ~7u) <= (cur_deflate_ + 1)->offset) {
+  if ((end_bit & 7) && ((end_bit + 7) & ~7ull) <= (cur_deflate_ + 1)->offset) {
     extra_byte_ = 1;
   } else {
     extra_byte_ = 0;
diff --git a/src/unittest_common.cc b/src/unittest_common.cc
new file mode 100644
index 0000000..b542740
--- /dev/null
+++ b/src/unittest_common.cc
@@ -0,0 +1,26 @@
+// Copyright 2018 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "puffin/src/unittest_common.h"
+
+namespace puffin {
+
+using std::string;
+
+bool MakeTempFile(string* filename, int* fd) {
+  char tmp_template[] = "/tmp/puffin-XXXXXX";
+  int mkstemp_fd = mkstemp(tmp_template);
+  TEST_AND_RETURN_FALSE(mkstemp_fd >= 0);
+  if (filename) {
+    *filename = tmp_template;
+  }
+  if (fd) {
+    *fd = mkstemp_fd;
+  } else {
+    close(mkstemp_fd);
+  }
+  return true;
+}
+
+}  // namespace puffin
diff --git a/src/unittest_common.h b/src/unittest_common.h
index ac1e71c..d0bbb02 100644
--- a/src/unittest_common.h
+++ b/src/unittest_common.h
@@ -29,6 +29,11 @@ class ScopedPathUnlinker {
   DISALLOW_COPY_AND_ASSIGN(ScopedPathUnlinker);
 };
 
+// Makes a temporary file as /tmp/puffin-XXXXXX. Both |filename| and |fd| are
+// optional, but if given, they will be populated with the new temporary file's
+// values.
+bool MakeTempFile(std::string* filename, int* fd);
+
 // clang-format off
 // Uncompressed deflate block.
 const Buffer kRaw1 = {0x01, 0x02, 0x03, 0x04, 0x05};
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));