Snap for 4610834 from 397d26828f8ce6e2c65c81517e791e24a72d7199 to pi-release

Change-Id: I9ab206d1aa938d3233ddb151e26f38a612969c5f

4 files changed, 228 insertions, 40 deletions

diff --git a/src/include/puffin/utils.h b/src/include/puffin/utils.h
index d5f1832..4ff9c16 100644
--- a/src/include/puffin/utils.h
+++ b/src/include/puffin/utils.h
@@ -42,6 +42,11 @@ bool LocateDeflatesInZlibBlocks(const std::string& file_path,
                                 const std::vector<ByteExtent>& zlibs,
                                 std::vector<BitExtent>* deflates);
 
+// Searches for deflate locations in a gzip file. The results are
+// saved in |deflate_blocks|.
+bool LocateDeflatesInGzip(const Buffer& data,
+                          std::vector<ByteExtent>* deflate_blocks);
+
 // Search for the deflates in a zip archive, and put the result in
 // |deflate_blocks|.
 bool LocateDeflatesInZipArchive(const Buffer& data,
diff --git a/src/main.cc b/src/main.cc
index 19ecd08..6ff7bbe 100644
--- a/src/main.cc
+++ b/src/main.cc
@@ -13,27 +13,31 @@
 #include "gflags/gflags.h"
 #endif
 
+#include "puffin/src/extent_stream.h"
+#include "puffin/src/file_stream.h"
 #include "puffin/src/include/puffin/common.h"
 #include "puffin/src/include/puffin/huffer.h"
 #include "puffin/src/include/puffin/puffdiff.h"
 #include "puffin/src/include/puffin/puffer.h"
 #include "puffin/src/include/puffin/puffpatch.h"
 #include "puffin/src/include/puffin/utils.h"
-#include "puffin/src/extent_stream.h"
-#include "puffin/src/file_stream.h"
+#include "puffin/src/memory_stream.h"
 #include "puffin/src/puffin_stream.h"
 #include "puffin/src/set_errors.h"
 
-using std::vector;
-using std::string;
 using puffin::BitExtent;
+using puffin::Buffer;
 using puffin::ByteExtent;
-using puffin::ExtentStream;
 using puffin::Error;
+using puffin::ExtentStream;
 using puffin::FileStream;
 using puffin::Huffer;
+using puffin::MemoryStream;
 using puffin::Puffer;
+using puffin::PuffinStream;
 using puffin::UniqueStreamPtr;
+using std::string;
+using std::vector;
 
 namespace {
 
@@ -86,7 +90,11 @@ const size_t kDefaultPuffCacheSize = 50 * 1024 * 1024;  // 50 MB
   DEFINE_string(dst_extents, "",                                           \
                 "Target extents in the format of offset:length,...");      \
   DEFINE_string(operation, "",                                             \
-                "Type of the operation: puff, huff, puffdiff, puffpatch"); \
+                "Type of the operation: puff, huff, puffdiff, puffpatch, " \
+                "puffhuff");                                               \
+  DEFINE_string(src_file_type, "",                                         \
+                "Type of the input source file: deflate, gzip, "           \
+                "zlib or zip");                                            \
   DEFINE_bool(verbose, false,                                              \
               "Logs all the given parameters including internally "        \
               "generated ones");                                           \
@@ -120,21 +128,6 @@ int main(int argc, char** argv) {
   auto src_extents = StringToExtents<ByteExtent>(FLAGS_src_extents);
   auto dst_extents = StringToExtents<ByteExtent>(FLAGS_dst_extents);
 
-  if (FLAGS_verbose) {
-    LOG(INFO) << "src_deflates_byte: "
-              << puffin::ExtentsToString(src_deflates_byte);
-    LOG(INFO) << "dst_deflates_byte: "
-              << puffin::ExtentsToString(dst_deflates_byte);
-    LOG(INFO) << "src_deflates_bit: "
-              << puffin::ExtentsToString(src_deflates_bit);
-    LOG(INFO) << "dst_deflates_bit: "
-              << puffin::ExtentsToString(dst_deflates_bit);
-    LOG(INFO) << "src_puffs: " << puffin::ExtentsToString(src_puffs);
-    LOG(INFO) << "dst_puffs: " << puffin::ExtentsToString(dst_puffs);
-    LOG(INFO) << "src_extents: " << puffin::ExtentsToString(src_extents);
-    LOG(INFO) << "dst_extents: " << puffin::ExtentsToString(dst_extents);
-  }
-
   auto src_stream = FileStream::Open(FLAGS_src_file, true, false);
   TEST_AND_RETURN_VALUE(src_stream, -1);
   if (!src_extents.empty()) {
@@ -143,11 +136,41 @@ int main(int argc, char** argv) {
     TEST_AND_RETURN_VALUE(src_stream, -1);
   }
 
-  vector<ByteExtent> puffs;
-  if (FLAGS_operation == "puff") {
+  if (!FLAGS_src_file_type.empty()) {
+    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);
+
+  if (FLAGS_operation == "puff" || FLAGS_operation == "puffhuff") {
     auto puffer = std::make_shared<Puffer>();
-    auto dst_stream = FileStream::Open(FLAGS_dst_file, false, true);
-    TEST_AND_RETURN_VALUE(dst_stream, -1);
     if (src_deflates_bit.empty() && src_deflates_byte.empty()) {
       LOG(WARNING) << "You should pass source deflates, is this intentional?";
     }
@@ -161,23 +184,51 @@ int main(int argc, char** argv) {
     TEST_AND_RETURN_VALUE(FindPuffLocations(src_stream, src_deflates_bit,
                                             &dst_puffs, &dst_puff_size),
                           -1);
-    if (FLAGS_verbose) {
-      LOG(INFO) << "out_dst_puffs: " << puffin::ExtentsToString(dst_puffs);
-    }
-    // Puff using the given puff_size.
-    auto reader = puffin::PuffinStream::CreateForPuff(
-        std::move(src_stream), puffer, dst_puff_size, src_deflates_bit,
-        dst_puffs);
-    puffin::Buffer buffer(1024 * 1024);
+
+    auto dst_stream = FileStream::Open(FLAGS_dst_file, false, true);
+    TEST_AND_RETURN_VALUE(dst_stream, -1);
+
+    auto reader =
+        PuffinStream::CreateForPuff(std::move(src_stream), puffer,
+                                    dst_puff_size, src_deflates_bit, dst_puffs);
+
+    Buffer puff_buffer;
+    auto writer = FLAGS_operation == "puffhuff"
+                      ? MemoryStream::CreateForWrite(&puff_buffer)
+                      : std::move(dst_stream);
+
+    Buffer buffer(1024 * 1024);
     size_t bytes_wrote = 0;
     while (bytes_wrote < dst_puff_size) {
       auto write_size = std::min(
           buffer.size(), static_cast<size_t>(dst_puff_size - bytes_wrote));
       TEST_AND_RETURN_VALUE(reader->Read(buffer.data(), write_size), -1);
-      TEST_AND_RETURN_VALUE(dst_stream->Write(buffer.data(), write_size), -1);
+      TEST_AND_RETURN_VALUE(writer->Write(buffer.data(), write_size), -1);
       bytes_wrote += write_size;
     }
 
+    // puffhuff operation puffs a stream and huffs it back to the target stream
+    // to make sure we can get to the original stream.
+    if (FLAGS_operation == "puffhuff") {
+      src_puffs = dst_puffs;
+      dst_deflates_byte = src_deflates_byte;
+      dst_deflates_bit = src_deflates_bit;
+
+      auto read_puff_stream = MemoryStream::CreateForRead(puff_buffer);
+      auto huffer = std::make_shared<Huffer>();
+      auto huff_writer = PuffinStream::CreateForHuff(
+          std::move(dst_stream), huffer, dst_puff_size, dst_deflates_bit,
+          src_puffs);
+
+      size_t bytes_read = 0;
+      while (bytes_read < dst_puff_size) {
+        auto read_size = std::min(buffer.size(), dst_puff_size - bytes_read);
+        TEST_AND_RETURN_VALUE(read_puff_stream->Read(buffer.data(), read_size),
+                              -1);
+        TEST_AND_RETURN_VALUE(huff_writer->Write(buffer.data(), read_size), -1);
+        bytes_read += read_size;
+      }
+    }
   } else if (FLAGS_operation == "huff") {
     if (dst_deflates_bit.empty() && src_puffs.empty()) {
       LOG(WARNING) << "You should pass source puffs and destination deflates"
@@ -190,11 +241,11 @@ int main(int argc, char** argv) {
     TEST_AND_RETURN_VALUE(dst_file, -1);
 
     auto huffer = std::make_shared<Huffer>();
-    auto dst_stream = puffin::PuffinStream::CreateForHuff(
-        std::move(dst_file), huffer, src_stream_size, dst_deflates_bit,
-        src_puffs);
+    auto dst_stream = PuffinStream::CreateForHuff(std::move(dst_file), huffer,
+                                                  src_stream_size,
+                                                  dst_deflates_bit, src_puffs);
 
-    puffin::Buffer buffer(1024 * 1024);
+    Buffer buffer(1024 * 1024);
     size_t bytes_read = 0;
     while (bytes_read < src_stream_size) {
       auto read_size = std::min(buffer.size(), src_stream_size - bytes_read);
@@ -229,7 +280,7 @@ int main(int argc, char** argv) {
                             -1);
     }
 
-    puffin::Buffer puffdiff_delta;
+    Buffer puffdiff_delta;
     TEST_AND_RETURN_VALUE(
         puffin::PuffDiff(std::move(src_stream), std::move(dst_stream),
                          src_deflates_bit, dst_deflates_bit, "/tmp/patch.tmp",
@@ -248,7 +299,7 @@ int main(int argc, char** argv) {
     size_t patch_size;
     TEST_AND_RETURN_VALUE(patch_stream->GetSize(&patch_size), -1);
 
-    puffin::Buffer puffdiff_delta(patch_size);
+    Buffer puffdiff_delta(patch_size);
     TEST_AND_RETURN_VALUE(
         patch_stream->Read(puffdiff_delta.data(), puffdiff_delta.size()), -1);
     auto dst_stream = FileStream::Open(FLAGS_dst_file, false, true);
@@ -266,5 +317,20 @@ int main(int argc, char** argv) {
                           FLAGS_cache_size),  // max_cache_size
         -1);
   }
+
+  if (FLAGS_verbose) {
+    LOG(INFO) << "src_deflates_byte: "
+              << puffin::ExtentsToString(src_deflates_byte);
+    LOG(INFO) << "dst_deflates_byte: "
+              << puffin::ExtentsToString(dst_deflates_byte);
+    LOG(INFO) << "src_deflates_bit: "
+              << puffin::ExtentsToString(src_deflates_bit);
+    LOG(INFO) << "dst_deflates_bit: "
+              << puffin::ExtentsToString(dst_deflates_bit);
+    LOG(INFO) << "src_puffs: " << puffin::ExtentsToString(src_puffs);
+    LOG(INFO) << "dst_puffs: " << puffin::ExtentsToString(dst_puffs);
+    LOG(INFO) << "src_extents: " << puffin::ExtentsToString(src_extents);
+    LOG(INFO) << "dst_extents: " << puffin::ExtentsToString(dst_extents);
+  }
   return 0;
 }
diff --git a/src/utils.cc b/src/utils.cc
index 35a39e9..ba10570 100644
--- a/src/utils.cc
+++ b/src/utils.cc
@@ -176,6 +176,78 @@ bool LocateDeflatesInZlibBlocks(const string& file_path,
   return LocateDeflatesInZlibBlocks(src, zlibs, deflates);
 }
 
+// For more information about gzip format, refer to RFC 1952 located at:
+// https://www.ietf.org/rfc/rfc1952.txt
+bool LocateDeflatesInGzip(const Buffer& data,
+                          vector<ByteExtent>* deflate_blocks) {
+  size_t member_start = 0;
+  while (member_start < data.size()) {
+    // Each member entry has the following format
+    // 0      1     0x1F
+    // 1      1     0x8B
+    // 2      1     compression method (8 denotes deflate)
+    // 3      1     set of flags
+    // 4      4     modification time
+    // 8      1     extra flags
+    // 9      1     operating system
+    TEST_AND_RETURN_FALSE(member_start + 10 <= data.size());
+    TEST_AND_RETURN_FALSE(data[member_start + 0] == 0x1F);
+    TEST_AND_RETURN_FALSE(data[member_start + 1] == 0x8B);
+    TEST_AND_RETURN_FALSE(data[member_start + 2] == 8);
+
+    size_t offset = member_start + 10;
+    int flag = data[member_start + 3];
+    // Extra field
+    if (flag & 4) {
+      TEST_AND_RETURN_FALSE(offset + 2 <= data.size());
+      uint16_t extra_length = data[offset++];
+      extra_length |= static_cast<uint16_t>(data[offset++]) << 8;
+      TEST_AND_RETURN_FALSE(offset + extra_length <= data.size());
+      offset += extra_length;
+    }
+    // File name field
+    if (flag & 8) {
+      while (true) {
+        TEST_AND_RETURN_FALSE(offset + 1 <= data.size());
+        if (data[offset++] == 0) {
+          break;
+        }
+      }
+    }
+    // File comment field
+    if (flag & 16) {
+      while (true) {
+        TEST_AND_RETURN_FALSE(offset + 1 <= data.size());
+        if (data[offset++] == 0) {
+          break;
+        }
+      }
+    }
+    // CRC16 field
+    if (flag & 2) {
+      offset += 2;
+    }
+
+    size_t compressed_size, uncompressed_size;
+    TEST_AND_RETURN_FALSE(CalculateSizeOfDeflateBlock(
+        data, offset, &compressed_size, &uncompressed_size));
+    TEST_AND_RETURN_FALSE(offset + compressed_size <= data.size());
+    deflate_blocks->push_back(ByteExtent(offset, compressed_size));
+    offset += compressed_size;
+
+    // Ignore CRC32;
+    TEST_AND_RETURN_FALSE(offset + 8 <= data.size());
+    offset += 4;
+    uint32_t u_size = 0;
+    for (size_t i = 0; i < 4; i++) {
+      u_size |= static_cast<uint32_t>(data[offset++]) << (i * 8);
+    }
+    TEST_AND_RETURN_FALSE(uncompressed_size % (1 << 31) == u_size);
+    member_start = offset;
+  }
+  return true;
+}
+
 // For more information about the zip format, refer to
 // https://support.pkware.com/display/PKZIP/APPNOTE
 bool LocateDeflatesInZipArchive(const Buffer& data,
diff --git a/src/utils_unittest.cc b/src/utils_unittest.cc
index 10a4c0d..55f2fa1 100644
--- a/src/utils_unittest.cc
+++ b/src/utils_unittest.cc
@@ -47,6 +47,32 @@ const uint8_t kZipEntryWithDataDescriptor[] = {
     0x2e, 0x00, 0xb4, 0xa0, 0xf2, 0x36, 0x06, 0x00, 0x00, 0x00, 0x07, 0x00,
     0x00, 0x00};
 
+// echo "0123456789" > test1.txt && echo "9876543210" > test2.txt &&
+// gzip -kf test1.txt test2.txt && cat test1.txt.gz test2.txt.gz |
+// hexdump -v -e '12/1 "0x%02x, " "\n"'
+const uint8_t kGzipEntryWithMultipleMembers[] = {
+    0x1f, 0x8b, 0x08, 0x08, 0x77, 0xd5, 0x84, 0x5a, 0x00, 0x03, 0x74, 0x65,
+    0x73, 0x74, 0x31, 0x2e, 0x74, 0x78, 0x74, 0x00, 0x33, 0x30, 0x34, 0x32,
+    0x36, 0x31, 0x35, 0x33, 0xb7, 0xb0, 0xe4, 0x02, 0x00, 0xd1, 0xe5, 0x76,
+    0x40, 0x0b, 0x00, 0x00, 0x00, 0x1f, 0x8b, 0x08, 0x08, 0x77, 0xd5, 0x84,
+    0x5a, 0x00, 0x03, 0x74, 0x65, 0x73, 0x74, 0x32, 0x2e, 0x74, 0x78, 0x74,
+    0x00, 0xb3, 0xb4, 0x30, 0x37, 0x33, 0x35, 0x31, 0x36, 0x32, 0x34, 0xe0,
+    0x02, 0x00, 0x20, 0x9c, 0x5f, 0x89, 0x0b, 0x00, 0x00, 0x00};
+
+// echo "0123456789" > test1.txt && gzip -kf test1.txt && cat test1.txt.gz |
+// hexdump -v -e '12/1 "0x%02x, " "\n"'
+// And manually insert extra field with two byte length (10) followed by:
+// echo "extrafield" | hexdump -v -e '12/1 "0x%02x, " "\n"'
+// Then change the forth byte of array to -x0c to enable the extra field.
+const uint8_t kGzipEntryWithExtraField[] = {
+    0x1f, 0x8b, 0x08, 0x0c, 0xcf, 0x0e, 0x86, 0x5a, 0x00, 0x03,
+    // Extra field begin
+    0x0A, 0x00, 0x65, 0x78, 0x74, 0x72, 0x61, 0x66, 0x69, 0x65, 0x6c, 0x64,
+    // Extra field end
+    0x74, 0x65, 0x73, 0x74, 0x31, 0x2e, 0x74, 0x78, 0x74, 0x00, 0x33, 0x30,
+    0x34, 0x32, 0x36, 0x31, 0x35, 0x33, 0xb7, 0xb0, 0xe4, 0x02, 0x00, 0xd1,
+    0xe5, 0x76, 0x40, 0x0b, 0x00, 0x00, 0x00};
+
 void FindDeflatesInZlibBlocks(const Buffer& src,
                               const vector<ByteExtent>& zlibs,
                               const vector<BitExtent>& deflates) {
@@ -125,4 +151,23 @@ TEST(UtilsTest, LocateDeflatesInZipArchiveErrorChecks) {
   EXPECT_EQ(static_cast<size_t>(0), deflates_incomplete.size());
 }
 
+TEST(UtilsTest, LocateDeflatesInGzip) {
+  Buffer gzip_data(kGzipEntryWithMultipleMembers,
+                   std::end(kGzipEntryWithMultipleMembers));
+  vector<ByteExtent> deflates;
+  EXPECT_TRUE(LocateDeflatesInGzip(gzip_data, &deflates));
+  EXPECT_EQ(static_cast<size_t>(2), deflates.size());
+  EXPECT_EQ(ByteExtent(20, 13), deflates[0]);
+  EXPECT_EQ(ByteExtent(61, 13), deflates[1]);
+}
+
+TEST(UtilsTest, LocateDeflatesInGzipWithExtraField) {
+  Buffer gzip_data(kGzipEntryWithExtraField,
+                   std::end(kGzipEntryWithExtraField));
+  vector<ByteExtent> deflates;
+  EXPECT_TRUE(LocateDeflatesInGzip(gzip_data, &deflates));
+  EXPECT_EQ(static_cast<size_t>(1), deflates.size());
+  EXPECT_EQ(ByteExtent(32, 13), deflates[0]);
+}
+
 }  // namespace puffin