update_engine: Generate PUFFDIFF operation

This patch adds functionality for generating PUFFDIFF operation, but
since the minor verion has not been increase, it actually does not
generate the operation itself. PUFFDIFF is used for patching deflate
based files.

BUG=chromium:767120
TEST=unittest passes; brillo_update_payload {generate|verify} passes;
CQ-DEPEND=CL:726945

Change-Id: I4ffdac8dce3740ef9fc2156cb0dad83a8364c8b5
Reviewed-on: https://chromium-review.googlesource.com/658298
Commit-Ready: Amin Hassani <ahassani@chromium.org>
Tested-by: Amin Hassani <ahassani@chromium.org>
Reviewed-by: Ben Chan <benchan@chromium.org>
diff --git a/Android.mk b/Android.mk
index fcde4f5..08798aa 100644
--- a/Android.mk
+++ b/Android.mk
@@ -611,8 +611,9 @@
 # for any client code.
 ue_libpayload_generator_exported_static_libraries := \
     libbsdiff \
-    libpayload_consumer \
     liblzma \
+    libpayload_consumer \
+    libpuffdiff \
     update_metadata-protos \
     $(ue_libpayload_consumer_exported_static_libraries) \
     $(ue_update_metadata_protos_exported_static_libraries)
@@ -661,8 +662,9 @@
 LOCAL_C_INCLUDES := $(ue_common_c_includes)
 LOCAL_STATIC_LIBRARIES := \
     libbsdiff \
-    libpayload_consumer \
     liblzma \
+    libpayload_consumer \
+    libpuffdiff \
     update_metadata-protos \
     $(ue_common_static_libraries) \
     $(ue_libpayload_consumer_exported_static_libraries) \
diff --git a/payload_generator/deflate_utils.cc b/payload_generator/deflate_utils.cc
index 9dd3fa4..7a5d9a4 100644
--- a/payload_generator/deflate_utils.cc
+++ b/payload_generator/deflate_utils.cc
@@ -33,6 +33,8 @@
 
 using std::string;
 using std::vector;
+using puffin::BitExtent;
+using puffin::ByteExtent;
 
 namespace chromeos_update_engine {
 namespace deflate_utils {
@@ -81,6 +83,9 @@
   for (auto& in_file : *files) {  // We need to modify so no constant.
     TEST_AND_RETURN_FALSE(
         ShiftExtentsOverExtents(file.extents, &in_file.extents));
+    TEST_AND_RETURN_FALSE(
+        ShiftBitExtentsOverExtents(file.extents, &in_file.deflates));
+
     in_file.name = file.name + "/" + in_file.name;
     num_blocks += BlocksInExtents(in_file.extents);
   }
@@ -90,8 +95,20 @@
   return true;
 }
 
+bool IsBitExtentInExtent(const Extent& extent, const BitExtent& bit_extent) {
+  return (bit_extent.offset / 8) >= (extent.start_block() * kBlockSize) &&
+         ((bit_extent.offset + bit_extent.length + 7) / 8) <=
+             ((extent.start_block() + extent.num_blocks()) * kBlockSize);
+}
+
 }  // namespace
 
+ByteExtent ExpandToByteExtent(const BitExtent& extent) {
+  uint64_t offset = extent.offset / 8;
+  uint64_t length = ((extent.offset + extent.length + 7) / 8) - offset;
+  return {offset, length};
+}
+
 bool ShiftExtentsOverExtents(const vector<Extent>& base_extents,
                              vector<Extent>* over_extents) {
   if (BlocksInExtents(base_extents) < BlocksInExtents(*over_extents)) {
@@ -133,8 +150,104 @@
   return true;
 }
 
+bool ShiftBitExtentsOverExtents(const vector<Extent>& base_extents,
+                                vector<BitExtent>* over_extents) {
+  if (over_extents->empty()) {
+    return true;
+  }
+
+  // This check is needed to make sure the number of bytes in |over_extents|
+  // does not exceed |base_extents|.
+  auto last_extent = ExpandToByteExtent(over_extents->back());
+  TEST_AND_RETURN_FALSE(last_extent.offset + last_extent.length <=
+                        BlocksInExtents(base_extents) * kBlockSize);
+
+  for (auto o_ext = over_extents->begin(); o_ext != over_extents->end();) {
+    size_t gap_blocks = base_extents[0].start_block();
+    size_t last_end_block = base_extents[0].start_block();
+    bool o_ext_processed = false;
+    for (auto b_ext : base_extents) {  // We need to modify |b_ext|, so we copy.
+      gap_blocks += b_ext.start_block() - last_end_block;
+      last_end_block = b_ext.start_block() + b_ext.num_blocks();
+      b_ext.set_start_block(b_ext.start_block() - gap_blocks);
+      auto byte_o_ext = ExpandToByteExtent(*o_ext);
+      if (byte_o_ext.offset >= b_ext.start_block() * kBlockSize &&
+          byte_o_ext.offset <
+              (b_ext.start_block() + b_ext.num_blocks()) * kBlockSize) {
+        if ((byte_o_ext.offset + byte_o_ext.length) <=
+            (b_ext.start_block() + b_ext.num_blocks()) * kBlockSize) {
+          // |o_ext| is inside |b_ext|, increase its start block.
+          o_ext->offset += gap_blocks * kBlockSize * 8;
+          ++o_ext;
+        } else {
+          // |o_ext| spills over this |b_ext|, remove it.
+          o_ext = over_extents->erase(o_ext);
+        }
+        o_ext_processed = true;
+        break;  // We processed o_ext, so break the loop;
+      }
+    }
+    TEST_AND_RETURN_FALSE(o_ext_processed);
+  }
+  return true;
+}
+
+vector<BitExtent> FindDeflates(const vector<Extent>& extents,
+                               const vector<BitExtent>& in_deflates) {
+  vector<BitExtent> result;
+  // TODO(ahassani): Replace this with binary_search style search.
+  for (const auto& deflate : in_deflates) {
+    for (const auto& extent : extents) {
+      if (IsBitExtentInExtent(extent, deflate)) {
+        result.push_back(deflate);
+        break;
+      }
+    }
+  }
+  return result;
+}
+
+bool CompactDeflates(const vector<Extent>& extents,
+                     const vector<BitExtent>& in_deflates,
+                     vector<BitExtent>* out_deflates) {
+  size_t bytes_passed = 0;
+  out_deflates->reserve(in_deflates.size());
+  for (const auto& extent : extents) {
+    size_t gap_bytes = extent.start_block() * kBlockSize - bytes_passed;
+    for (const auto& deflate : in_deflates) {
+      if (IsBitExtentInExtent(extent, deflate)) {
+        out_deflates->emplace_back(deflate.offset - (gap_bytes * 8),
+                                   deflate.length);
+      }
+    }
+    bytes_passed += extent.num_blocks() * kBlockSize;
+  }
+
+  // All given |in_deflates| items should've been inside one of the extents in
+  // |extents|.
+  TEST_AND_RETURN_FALSE(in_deflates.size() == out_deflates->size());
+
+  // Make sure all outgoing deflates are ordered and non-overlapping.
+  auto result = std::adjacent_find(out_deflates->begin(),
+                                   out_deflates->end(),
+                                   [](const BitExtent& a, const BitExtent& b) {
+                                     return (a.offset + a.length) > b.offset;
+                                   });
+  TEST_AND_RETURN_FALSE(result == out_deflates->end());
+  return true;
+}
+
+bool FindAndCompactDeflates(const vector<Extent>& extents,
+                            const vector<BitExtent>& in_deflates,
+                            vector<BitExtent>* out_deflates) {
+  auto found_deflates = FindDeflates(extents, in_deflates);
+  TEST_AND_RETURN_FALSE(CompactDeflates(extents, found_deflates, out_deflates));
+  return true;
+}
+
 bool PreprocessParitionFiles(const PartitionConfig& part,
-                             vector<FilesystemInterface::File>* result_files) {
+                             vector<FilesystemInterface::File>* result_files,
+                             bool extract_deflates) {
   // Get the file system files.
   vector<FilesystemInterface::File> tmp_files;
   part.fs_interface->GetFiles(&tmp_files);
@@ -149,15 +262,18 @@
       TEST_AND_RETURN_FALSE(
           CopyExtentsToFile(part.path, file.extents, path.value(), kBlockSize));
       // Test if it is actually a Squashfs file.
-      auto sqfs = SquashfsFilesystem::CreateFromFile(path.value());
+      auto sqfs =
+          SquashfsFilesystem::CreateFromFile(path.value(), extract_deflates);
       if (sqfs) {
         // It is an squashfs file. Get its files to replace with itself.
         vector<FilesystemInterface::File> files;
         sqfs->GetFiles(&files);
 
-        // Replace squashfs file with its files only if |files| has at
-        // least two files.
-        if (files.size() > 1) {
+        // Replace squashfs file with its files only if |files| has at least two
+        // files or if it has some deflates (since it is better to replace it to
+        // take advantage of the deflates.)
+        if (files.size() > 1 ||
+            (files.size() == 1 && !files[0].deflates.empty())) {
           TEST_AND_RETURN_FALSE(RealignSplittedFiles(file, &files));
           result_files->insert(result_files->end(), files.begin(), files.end());
           continue;
@@ -170,7 +286,6 @@
     // TODO(ahassani): Process other types of files like apk, zip, etc.
     result_files->push_back(file);
   }
-
   return true;
 }
 
diff --git a/payload_generator/deflate_utils.h b/payload_generator/deflate_utils.h
index 31082c3..798ce25 100644
--- a/payload_generator/deflate_utils.h
+++ b/payload_generator/deflate_utils.h
@@ -17,6 +17,7 @@
 #ifndef UPDATE_ENGINE_PAYLOAD_GENERATOR_DEFLATE_UTILS_H_
 #define UPDATE_ENGINE_PAYLOAD_GENERATOR_DEFLATE_UTILS_H_
 
+#include <puffin/puffdiff.h>
 #include <vector>
 
 #include "update_engine/payload_generator/filesystem_interface.h"
@@ -25,8 +26,15 @@
 namespace chromeos_update_engine {
 namespace deflate_utils {
 
+// Gets the files from the partition and processes all its files. Processing
+// includes:
+//  - splitting large Squashfs containers into its smaller files.
+bool PreprocessParitionFiles(const PartitionConfig& part,
+                             std::vector<FilesystemInterface::File>* result,
+                             bool extract_deflates);
+
 // Spreads all extents in |over_extents| over |base_extents|. Here we assume the
-// extents are non-overlapping.
+// |over_extents| are non-overlapping and sorted by their offset.
 //
 // |base_extents|:
 // |               -----------------------        ------         --------------
@@ -38,11 +46,51 @@
 bool ShiftExtentsOverExtents(const std::vector<Extent>& base_extents,
                              std::vector<Extent>* over_extents);
 
-// Gets the files from the partition and processes all its files. Processing
-// includes:
-//  - splitting large Squashfs containers into its smaller files.
-bool PreprocessParitionFiles(const PartitionConfig& part,
-                             std::vector<FilesystemInterface::File>* result);
+// Spreads all extents in |over_extents| over |base_extents|. Here we assume the
+// |over_extents| are non-overlapping and sorted by their offset. An item in
+// |over_extents| is removed if it is spread in two or more extents in
+// |base_extents|.
+//
+// |base_extents|:
+// |               -----------------------        ------         --------------
+// |over_extents|:
+// |  ==========  ====    ==========  ======
+// |over_extents| is transforms to:
+// |                 ==========  ====                                 ======
+//
+bool ShiftBitExtentsOverExtents(const std::vector<Extent>& base_extents,
+                                std::vector<puffin::BitExtent>* over_extents);
+
+// Finds all deflate locations in |deflates| that are inside an Extent in
+// |extents|. This function should not change the order of deflates.
+std::vector<puffin::BitExtent> FindDeflates(
+    const std::vector<Extent>& extents,
+    const std::vector<puffin::BitExtent>& deflates);
+
+// Creates a new list of deflate locations (|out_deflates|) from |in_deflates|
+// by assuming all extents in the |extents| have been put together
+// linearly. This function assumes that all deflate locations given in
+// |in_deflates| are located somewhere in the |extents|. |out_deflates| should
+// be empty on call.
+//
+// |extents|:
+// |               -----------------------        ------         --------------
+// |in_deflates|:
+// |                   ========  ====              ====           ======
+// |out_deflates|:
+// |    ========  ====      ====  ======
+//
+bool CompactDeflates(const std::vector<Extent>& extents,
+                     const std::vector<puffin::BitExtent>& in_deflates,
+                     std::vector<puffin::BitExtent>* out_deflates);
+
+// Combines |FindDeflates| and |CompcatDeflates| for ease of use.
+bool FindAndCompactDeflates(const std::vector<Extent>& extents,
+                            const std::vector<puffin::BitExtent>& in_deflates,
+                            std::vector<puffin::BitExtent>* out_deflates);
+
+// Expands a BitExtents to a ByteExtent.
+puffin::ByteExtent ExpandToByteExtent(const puffin::BitExtent& extent);
 
 }  // namespace deflate_utils
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/deflate_utils_unittest.cc b/payload_generator/deflate_utils_unittest.cc
index 8774ef0..cb9476a 100644
--- a/payload_generator/deflate_utils_unittest.cc
+++ b/payload_generator/deflate_utils_unittest.cc
@@ -30,15 +30,27 @@
 #include "update_engine/payload_generator/extent_utils.h"
 
 using std::vector;
+using puffin::BitExtent;
+using puffin::ByteExtent;
 
 namespace chromeos_update_engine {
 namespace deflate_utils {
 
+// This creates a sudo-random BitExtents from ByteExtents for simpler testing.
+vector<BitExtent> ByteToBitExtent(const vector<ByteExtent>& byte_extents) {
+  vector<BitExtent> bit_extents;
+  for (auto& byte_extent : byte_extents) {
+    bit_extents.emplace_back(byte_extent.offset * 8 + (byte_extent.offset & 7),
+                             byte_extent.length * 8 - (byte_extent.length & 7));
+  }
+  return bit_extents;
+}
+
 TEST(DeflateUtilsTest, ExtentsShiftTest) {
   vector<Extent> base_extents = {ExtentForRange(10, 10),
-                                 ExtentForRange(30, 10),
-                                 ExtentForRange(50, 10),
                                  ExtentForRange(70, 10),
+                                 ExtentForRange(50, 10),
+                                 ExtentForRange(30, 10),
                                  ExtentForRange(90, 10)};
   vector<Extent> over_extents = {ExtentForRange(2, 2),
                                  ExtentForRange(5, 2),
@@ -49,10 +61,10 @@
   vector<Extent> out_over_extents = {ExtentForRange(12, 2),
                                      ExtentForRange(15, 2),
                                      ExtentForRange(17, 3),
-                                     ExtentForRange(33, 7),
+                                     ExtentForRange(73, 7),
                                      ExtentForRange(50, 3),
                                      ExtentForRange(55, 5),
-                                     ExtentForRange(70, 10),
+                                     ExtentForRange(30, 10),
                                      ExtentForRange(90, 5),
                                      ExtentForRange(97, 3)};
   EXPECT_TRUE(ShiftExtentsOverExtents(base_extents, &over_extents));
@@ -64,5 +76,115 @@
   EXPECT_FALSE(ShiftExtentsOverExtents(base_extents, &over_extents));
 }
 
+TEST(DeflateUtilsTest, ShiftBitExtentsOverExtentsTest) {
+  vector<Extent> base_extents = {ExtentForRange(3, 1),
+                                 ExtentForRange(1, 1),
+                                 ExtentForRange(5, 1),
+                                 ExtentForRange(7, 1),
+                                 ExtentForRange(9, 1)};
+  vector<BitExtent> over_extents =
+      ByteToBitExtent({{0, 0}, {100, 2000}, {4096, 0}, {5000, 5000}});
+  vector<BitExtent> out_over_extents =
+      ByteToBitExtent({{12288, 0}, {12388, 2000}, {4096, 0}});
+  ASSERT_TRUE(ShiftBitExtentsOverExtents(base_extents, &over_extents));
+  EXPECT_EQ(over_extents, out_over_extents);
+}
+
+TEST(DeflateUtilsTest, ShiftBitExtentsOverExtentsBoundaryTest) {
+  vector<Extent> base_extents = {ExtentForRange(1, 1)};
+  vector<BitExtent> over_extents = ByteToBitExtent({{2, 4096}});
+  vector<BitExtent> out_over_extents = {};
+  EXPECT_FALSE(ShiftBitExtentsOverExtents(base_extents, &over_extents));
+
+  base_extents = {ExtentForRange(1, 1)};
+  over_extents = {};
+  out_over_extents = {};
+  EXPECT_TRUE(ShiftBitExtentsOverExtents(base_extents, &over_extents));
+  EXPECT_EQ(over_extents, out_over_extents);
+
+  base_extents = {};
+  over_extents = {};
+  out_over_extents = {};
+  EXPECT_TRUE(ShiftBitExtentsOverExtents(base_extents, &over_extents));
+  EXPECT_EQ(over_extents, out_over_extents);
+
+  base_extents = {};
+  over_extents = ByteToBitExtent({{0, 1}});
+  out_over_extents = ByteToBitExtent({{0, 1}});
+  EXPECT_FALSE(ShiftBitExtentsOverExtents(base_extents, &over_extents));
+  EXPECT_EQ(over_extents, out_over_extents);
+
+  base_extents = {ExtentForRange(1, 2)};
+  over_extents = ByteToBitExtent({{0, 3 * 4096}, {4 * 4096, 4096}});
+  out_over_extents = ByteToBitExtent({{0, 3 * 4096}, {4 * 4096, 4096}});
+  EXPECT_FALSE(ShiftBitExtentsOverExtents(base_extents, &over_extents));
+  EXPECT_EQ(over_extents, out_over_extents);
+}
+
+TEST(DeflateUtilsTest, FindDeflatesTest) {
+  vector<Extent> extents = {
+      ExtentForRange(1, 1), ExtentForRange(3, 1), ExtentForRange(5, 1)};
+  vector<BitExtent> in_deflates = ByteToBitExtent({{0, 0},
+                                                   {10, 400},
+                                                   {4096, 0},
+                                                   {3000, 2000},
+                                                   {4096, 100},
+                                                   {4097, 100},
+                                                   {8100, 92},
+                                                   {8100, 93},
+                                                   {8100, 6000},
+                                                   {25000, 1}});
+  vector<BitExtent> expected_out_deflates =
+      ByteToBitExtent({{4096, 0}, {4096, 100}, {4097, 100}, {8100, 92}});
+  vector<BitExtent> out_deflates;
+  out_deflates = FindDeflates(extents, in_deflates);
+  EXPECT_EQ(out_deflates, expected_out_deflates);
+}
+
+TEST(DeflateUtilsTest, FindDeflatesBoundaryTest) {
+  vector<Extent> extents = {};
+  vector<BitExtent> in_deflates = ByteToBitExtent({{0, 0}, {8100, 93}});
+  vector<BitExtent> expected_out_deflates = {};
+  vector<BitExtent> out_deflates;
+  out_deflates = FindDeflates(extents, in_deflates);
+  EXPECT_EQ(out_deflates, expected_out_deflates);
+
+  extents = {};
+  in_deflates = {};
+  out_deflates = FindDeflates(extents, in_deflates);
+  EXPECT_EQ(out_deflates, expected_out_deflates);
+}
+
+TEST(DeflateUtilsTest, CompactTest) {
+  vector<Extent> extents = {
+      ExtentForRange(1, 1), ExtentForRange(5, 1), ExtentForRange(3, 1)};
+  vector<BitExtent> in_deflates =
+      ByteToBitExtent({{4096, 0}, {12288, 4096}, {4096, 100}, {20480, 100}});
+  vector<BitExtent> expected_out_deflates =
+      ByteToBitExtent({{0, 0}, {0, 100}, {4096, 100}, {8192, 4096}});
+  vector<BitExtent> out_deflates;
+  ASSERT_TRUE(CompactDeflates(extents, in_deflates, &out_deflates));
+  EXPECT_EQ(out_deflates, expected_out_deflates);
+}
+
+TEST(DeflateUtilsTest, CompactBoundaryTest) {
+  vector<Extent> extents = {};
+  vector<BitExtent> in_deflates = ByteToBitExtent({{4096, 0}});
+  vector<BitExtent> expected_out_deflates = {};
+  vector<BitExtent> out_deflates;
+  EXPECT_FALSE(CompactDeflates(extents, in_deflates, &out_deflates));
+  EXPECT_EQ(out_deflates, expected_out_deflates);
+
+  extents = {};
+  in_deflates = {};
+  ASSERT_TRUE(CompactDeflates(extents, in_deflates, &out_deflates));
+  EXPECT_EQ(out_deflates, expected_out_deflates);
+
+  extents = {ExtentForRange(1, 1)};
+  in_deflates = {};
+  ASSERT_TRUE(CompactDeflates(extents, in_deflates, &out_deflates));
+  EXPECT_EQ(out_deflates, expected_out_deflates);
+}
+
 }  // namespace deflate_utils
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/delta_diff_utils.cc b/payload_generator/delta_diff_utils.cc
index 39438dd..b2b6b51 100644
--- a/payload_generator/delta_diff_utils.cc
+++ b/payload_generator/delta_diff_utils.cc
@@ -18,7 +18,7 @@
 
 #include <endian.h>
 #if defined(__clang__)
-// TODO: Remove these pragmas when b/35721782 is fixed.
+// TODO(*): Remove these pragmas when b/35721782 is fixed.
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wmacro-redefined"
 #endif
@@ -30,9 +30,12 @@
 
 #include <algorithm>
 #include <map>
+#include <memory>
+#include <utility>
 
 #include <base/files/file_util.h>
 #include <base/format_macros.h>
+#include <base/strings/string_util.h>
 #include <base/strings/stringprintf.h>
 #include <base/threading/simple_thread.h>
 #include <bsdiff/bsdiff.h>
@@ -46,6 +49,7 @@
 #include "update_engine/payload_generator/delta_diff_generator.h"
 #include "update_engine/payload_generator/extent_ranges.h"
 #include "update_engine/payload_generator/extent_utils.h"
+#include "update_engine/payload_generator/squashfs_filesystem.h"
 #include "update_engine/payload_generator/xz.h"
 
 using std::map;
@@ -64,8 +68,8 @@
 
 // The maximum destination size allowed for puffdiff. In general, puffdiff
 // should work for arbitrary big files, but the payload application is quite
-// memory intensive, so we limit these operations to 50 MiB.
-const uint64_t kMaxPuffdiffDestinationSize = 50 * 1024 * 1024;  // bytes
+// memory intensive, so we limit these operations to 150 MiB.
+const uint64_t kMaxPuffdiffDestinationSize = 150 * 1024 * 1024;  // bytes
 
 // Process a range of blocks from |range_start| to |range_end| in the extent at
 // position |*idx_p| of |extents|. If |do_remove| is true, this range will be
@@ -176,6 +180,8 @@
                      const PayloadVersion& version,
                      const vector<Extent>& old_extents,
                      const vector<Extent>& new_extents,
+                     const vector<puffin::BitExtent>& old_deflates,
+                     const vector<puffin::BitExtent>& new_deflates,
                      const string& name,
                      ssize_t chunk_blocks,
                      BlobFileWriter* blob_file)
@@ -184,6 +190,8 @@
         version_(version),
         old_extents_(old_extents),
         new_extents_(new_extents),
+        old_deflates_(old_deflates),
+        new_deflates_(new_deflates),
         name_(name),
         chunk_blocks_(chunk_blocks),
         blob_file_(blob_file) {}
@@ -208,6 +216,8 @@
   // The block ranges of the old/new file within the src/tgt image
   const vector<Extent> old_extents_;
   const vector<Extent> new_extents_;
+  const vector<puffin::BitExtent> old_deflates_;
+  const vector<puffin::BitExtent> new_deflates_;
   const string name_;
   // Block limit of one aop.
   ssize_t chunk_blocks_;
@@ -230,6 +240,8 @@
                      new_part_,
                      old_extents_,
                      new_extents_,
+                     old_deflates_,
+                     new_deflates_,
                      name_,
                      chunk_blocks_,
                      version_,
@@ -266,19 +278,20 @@
       &old_visited_blocks,
       &new_visited_blocks));
 
-  map<string, vector<Extent>> old_files_map;
+  bool puffdiff_allowed = version.OperationAllowed(InstallOperation::PUFFDIFF);
+  map<string, FilesystemInterface::File> old_files_map;
   if (old_part.fs_interface) {
     vector<FilesystemInterface::File> old_files;
-    TEST_AND_RETURN_FALSE(
-        deflate_utils::PreprocessParitionFiles(old_part, &old_files));
+    TEST_AND_RETURN_FALSE(deflate_utils::PreprocessParitionFiles(
+        old_part, &old_files, puffdiff_allowed));
     for (const FilesystemInterface::File& file : old_files)
-      old_files_map[file.name] = file.extents;
+      old_files_map[file.name] = file;
   }
 
   TEST_AND_RETURN_FALSE(new_part.fs_interface);
   vector<FilesystemInterface::File> new_files;
-  TEST_AND_RETURN_FALSE(
-      deflate_utils::PreprocessParitionFiles(new_part, &new_files));
+  TEST_AND_RETURN_FALSE(deflate_utils::PreprocessParitionFiles(
+      new_part, &new_files, puffdiff_allowed));
 
   vector<FileDeltaProcessor> file_delta_processors;
 
@@ -309,8 +322,9 @@
     // from using a graph/cycle detection/etc to generate diffs, and at that
     // time, it will be easy (non-complex) to have many operations read
     // from the same source blocks. At that time, this code can die. -adlr
-    vector<Extent> old_file_extents = FilterExtentRanges(
-        old_files_map[new_file.name], old_visited_blocks);
+    auto old_file = old_files_map[new_file.name];
+    vector<Extent> old_file_extents =
+        FilterExtentRanges(old_file.extents, old_visited_blocks);
     old_visited_blocks.AddExtents(old_file_extents);
 
     file_delta_processors.emplace_back(old_part.path,
@@ -318,6 +332,8 @@
                                        version,
                                        std::move(old_file_extents),
                                        std::move(new_file_extents),
+                                       old_file.deflates,
+                                       new_file.deflates,
                                        new_file.name,  // operation name
                                        hard_chunk_blocks,
                                        blob_file);
@@ -361,6 +377,8 @@
                                       new_part.path,
                                       old_unvisited,
                                       new_unvisited,
+                                      {},                 // old_deflates,
+                                      {},                 // new_deflates
                                       "<non-file-data>",  // operation name
                                       soft_chunk_blocks,
                                       version,
@@ -465,6 +483,8 @@
                                         new_part,
                                         vector<Extent>(),        // old_extents
                                         vector<Extent>{extent},  // new_extents
+                                        {},                      // old_deflates
+                                        {},                      // new_deflates
                                         "<zeros>",
                                         chunk_blocks,
                                         version,
@@ -525,6 +545,8 @@
                    const string& new_part,
                    const vector<Extent>& old_extents,
                    const vector<Extent>& new_extents,
+                   const vector<puffin::BitExtent>& old_deflates,
+                   const vector<puffin::BitExtent>& new_deflates,
                    const string& name,
                    ssize_t chunk_blocks,
                    const PayloadVersion& version,
@@ -553,6 +575,8 @@
                                             new_part,
                                             old_extents_chunk,
                                             new_extents_chunk,
+                                            old_deflates,
+                                            new_deflates,
                                             version,
                                             &data,
                                             &operation));
@@ -643,6 +667,8 @@
                        const string& new_part,
                        const vector<Extent>& old_extents,
                        const vector<Extent>& new_extents,
+                       const vector<puffin::BitExtent>& old_deflates,
+                       const vector<puffin::BitExtent>& new_deflates,
                        const PayloadVersion& version,
                        brillo::Blob* out_data,
                        InstallOperation* out_op) {
@@ -731,8 +757,58 @@
         }
       }
       if (puffdiff_allowed) {
-        LOG(ERROR) << "puffdiff is not supported yet!";
-        return false;
+        // Find all deflate positions inside the given extents and then put all
+        // deflates together because we have already read all the extents into
+        // one buffer.
+        vector<puffin::BitExtent> src_deflates;
+        TEST_AND_RETURN_FALSE(deflate_utils::FindAndCompactDeflates(
+            src_extents, old_deflates, &src_deflates));
+
+        vector<puffin::BitExtent> dst_deflates;
+        TEST_AND_RETURN_FALSE(deflate_utils::FindAndCompactDeflates(
+            dst_extents, new_deflates, &dst_deflates));
+
+        // Remove equal deflates. TODO(*): We can do a N*N check using
+        // hashing. It will not reduce the payload size, but it will speeds up
+        // the puffing on the client device.
+        auto src = src_deflates.begin();
+        auto dst = dst_deflates.begin();
+        for (; src != src_deflates.end() && dst != dst_deflates.end();) {
+          auto src_in_bytes = deflate_utils::ExpandToByteExtent(*src);
+          auto dst_in_bytes = deflate_utils::ExpandToByteExtent(*dst);
+          if (src_in_bytes.length == dst_in_bytes.length &&
+              !memcmp(old_data.data() + src_in_bytes.offset,
+                      new_data.data() + dst_in_bytes.offset,
+                      src_in_bytes.length)) {
+            src = src_deflates.erase(src);
+            dst = dst_deflates.erase(dst);
+          } else {
+            src++;
+            dst++;
+          }
+        }
+
+        // Only Puffdiff if both files have at least one deflate left.
+        if (!src_deflates.empty() && !dst_deflates.empty()) {
+          brillo::Blob puffdiff_delta;
+          string temp_file_path;
+          TEST_AND_RETURN_FALSE(utils::MakeTempFile(
+              "puffdiff-delta.XXXXXX", &temp_file_path, nullptr));
+          ScopedPathUnlinker temp_file_unlinker(temp_file_path);
+
+          // Perform PuffDiff operation.
+          TEST_AND_RETURN_FALSE(puffin::PuffDiff(old_data,
+                                                 new_data,
+                                                 src_deflates,
+                                                 dst_deflates,
+                                                 temp_file_path,
+                                                 &puffdiff_delta));
+          TEST_AND_RETURN_FALSE(puffdiff_delta.size() > 0);
+          if (puffdiff_delta.size() < data_blob.size()) {
+            operation.set_type(InstallOperation::PUFFDIFF);
+            data_blob = std::move(puffdiff_delta);
+          }
+        }
       }
     }
   }
diff --git a/payload_generator/delta_diff_utils.h b/payload_generator/delta_diff_utils.h
index d21468f..dea8535 100644
--- a/payload_generator/delta_diff_utils.h
+++ b/payload_generator/delta_diff_utils.h
@@ -21,6 +21,7 @@
 #include <vector>
 
 #include <brillo/secure_blob.h>
+#include <puffin/puffdiff.h>
 
 #include "update_engine/payload_generator/annotated_operation.h"
 #include "update_engine/payload_generator/extent_ranges.h"
@@ -76,12 +77,15 @@
 // stored in |new_part| in the blocks described by |new_extents| and, if it
 // exists, the old version exists in |old_part| in the blocks described by
 // |old_extents|. The operations added to |aops| reference the data blob
-// in the |blob_file|. Returns true on success.
+// in the |blob_file|. |old_deflates| and |new_deflates| are all deflate
+// locations in |old_part| and |new_part|. Returns true on success.
 bool DeltaReadFile(std::vector<AnnotatedOperation>* aops,
                    const std::string& old_part,
                    const std::string& new_part,
                    const std::vector<Extent>& old_extents,
                    const std::vector<Extent>& new_extents,
+                   const std::vector<puffin::BitExtent>& old_deflates,
+                   const std::vector<puffin::BitExtent>& new_deflates,
                    const std::string& name,
                    ssize_t chunk_blocks,
                    const PayloadVersion& version,
@@ -94,11 +98,14 @@
 // MOVE or SOURCE_COPY operation. If there is a change, the smallest of the
 // operations allowed in the given |version| (REPLACE, REPLACE_BZ, BSDIFF,
 // SOURCE_BSDIFF, or PUFFDIFF) wins.
-// |new_extents| must not be empty. Returns true on success.
+// |new_extents| must not be empty. |old_deflates| and |new_deflates| are all
+// the deflate locations in |old_part| and |new_part|. Returns true on success.
 bool ReadExtentsToDiff(const std::string& old_part,
                        const std::string& new_part,
                        const std::vector<Extent>& old_extents,
                        const std::vector<Extent>& new_extents,
+                       const std::vector<puffin::BitExtent>& old_deflates,
+                       const std::vector<puffin::BitExtent>& new_deflates,
                        const PayloadVersion& version,
                        brillo::Blob* out_data,
                        InstallOperation* out_op);
diff --git a/payload_generator/delta_diff_utils_unittest.cc b/payload_generator/delta_diff_utils_unittest.cc
index 34e33ad..46e68c5 100644
--- a/payload_generator/delta_diff_utils_unittest.cc
+++ b/payload_generator/delta_diff_utils_unittest.cc
@@ -179,6 +179,8 @@
       new_part_.path,
       old_extents,
       new_extents,
+      {},  // old_deflates
+      {},  // new_deflates
       PayloadVersion(kChromeOSMajorPayloadVersion, kInPlaceMinorPayloadVersion),
       &data,
       &op));
@@ -238,6 +240,8 @@
       new_part_.path,
       old_extents,
       new_extents,
+      {},  // old_deflates
+      {},  // new_deflates
       PayloadVersion(kChromeOSMajorPayloadVersion, kInPlaceMinorPayloadVersion),
       &data,
       &op));
@@ -301,6 +305,8 @@
       new_part_.path,
       old_extents,
       new_extents,
+      {},  // old_deflates
+      {},  // new_deflates
       PayloadVersion(kChromeOSMajorPayloadVersion, kInPlaceMinorPayloadVersion),
       &data,
       &op));
@@ -349,6 +355,8 @@
         new_part_.path,
         old_extents,
         new_extents,
+        {},  // old_deflates
+        {},  // new_deflates
         PayloadVersion(kChromeOSMajorPayloadVersion,
                        kInPlaceMinorPayloadVersion),
         &data,
@@ -390,6 +398,8 @@
       new_part_.path,
       old_extents,
       new_extents,
+      {},  // old_deflates
+      {},  // new_deflates
       PayloadVersion(kChromeOSMajorPayloadVersion, kSourceMinorPayloadVersion),
       &data,
       &op));
@@ -422,6 +432,8 @@
       new_part_.path,
       old_extents,
       new_extents,
+      {},  // old_deflates
+      {},  // new_deflates
       PayloadVersion(kChromeOSMajorPayloadVersion, kSourceMinorPayloadVersion),
       &data,
       &op));
diff --git a/payload_generator/filesystem_interface.h b/payload_generator/filesystem_interface.h
index 866c46b..b1506e4 100644
--- a/payload_generator/filesystem_interface.h
+++ b/payload_generator/filesystem_interface.h
@@ -33,6 +33,7 @@
 
 #include <base/macros.h>
 #include <brillo/key_value_store.h>
+#include <puffin/utils.h>
 
 #include "update_engine/update_metadata.pb.h"
 
@@ -62,6 +63,10 @@
     // between 0 and GetBlockCount() - 1. The blocks are encoded in extents,
     // indicating the starting block, and the number of consecutive blocks.
     std::vector<Extent> extents;
+
+    // All the deflate locations in the file. These locations are not relative
+    // to the extents. They are relative to the file system itself.
+    std::vector<puffin::BitExtent> deflates;
   };
 
   virtual ~FilesystemInterface() = default;
diff --git a/payload_generator/inplace_generator.cc b/payload_generator/inplace_generator.cc
index bc140e8..cf02c91 100644
--- a/payload_generator/inplace_generator.cc
+++ b/payload_generator/inplace_generator.cc
@@ -573,6 +573,8 @@
         new_part,
         vector<Extent>(),  // old_extents
         new_extents,
+        {},  // old_deflates
+        {},  // new_deflates
         (*graph)[cut.old_dst].aop.name,
         -1,  // chunk_blocks, forces to have a single operation.
         kInPlacePayloadVersion,
diff --git a/payload_generator/squashfs_filesystem.cc b/payload_generator/squashfs_filesystem.cc
index 54a2c1a..c98ad12 100644
--- a/payload_generator/squashfs_filesystem.cc
+++ b/payload_generator/squashfs_filesystem.cc
@@ -20,6 +20,7 @@
 
 #include <algorithm>
 #include <string>
+#include <utility>
 
 #include <base/files/file_util.h>
 #include <base/logging.h>
@@ -29,6 +30,7 @@
 
 #include "update_engine/common/subprocess.h"
 #include "update_engine/common/utils.h"
+#include "update_engine/payload_generator/deflate_utils.h"
 #include "update_engine/payload_generator/delta_diff_generator.h"
 #include "update_engine/payload_generator/extent_ranges.h"
 #include "update_engine/payload_generator/extent_utils.h"
@@ -53,6 +55,7 @@
 // The size of the squashfs super block.
 constexpr size_t kSquashfsSuperBlockSize = 96;
 constexpr uint64_t kSquashfsCompressedBit = 1 << 24;
+constexpr uint32_t kSquashfsZlibCompression = 1;
 
 bool ReadSquashfsHeader(const brillo::Blob blob,
                         SquashfsFilesystem::SquashfsHeader* header) {
@@ -96,9 +99,18 @@
 }  // namespace
 
 bool SquashfsFilesystem::Init(const string& map,
+                              const string& sqfs_path,
                               size_t size,
-                              const SquashfsHeader& header) {
+                              const SquashfsHeader& header,
+                              bool extract_deflates) {
   size_ = size;
+
+  bool is_zlib = header.compression_type == kSquashfsZlibCompression;
+  if (!is_zlib) {
+    LOG(WARNING) << "Filesystem is not Gzipped. Not filling deflates!";
+  }
+  vector<puffin::ByteExtent> zlib_blks;
+
   // Reading files map. For the format of the file map look at the comments for
   // |CreateFromFileMap()|.
   auto lines = base::SplitStringPiece(map,
@@ -122,6 +134,12 @@
       // TODO(ahassani): For puffin push it into a proper list if uncompressed.
       auto new_blk_size = blk_size & ~kSquashfsCompressedBit;
       TEST_AND_RETURN_FALSE(new_blk_size <= header.block_size);
+      if (new_blk_size > 0 && !(blk_size & kSquashfsCompressedBit)) {
+        // Compressed block
+        if (is_zlib && extract_deflates) {
+          zlib_blks.emplace_back(cur_offset, new_blk_size);
+        }
+      }
       cur_offset += new_blk_size;
     }
 
@@ -192,11 +210,45 @@
   std::sort(files_.begin(), files_.end(), [](const File& a, const File& b) {
     return a.extents[0].start_block() < b.extents[0].start_block();
   });
+
+  if (is_zlib && extract_deflates) {
+    // If it is infact gzipped, then the sqfs_path should be valid to read its
+    // content.
+    TEST_AND_RETURN_FALSE(!sqfs_path.empty());
+    if (zlib_blks.empty()) {
+      return true;
+    }
+
+    // Sort zlib blocks.
+    std::sort(zlib_blks.begin(),
+              zlib_blks.end(),
+              [](const puffin::ByteExtent& a, const puffin::ByteExtent& b) {
+                return a.offset < b.offset;
+              });
+
+    // Sanity check. Make sure zlib blocks are not overlapping.
+    auto result = std::adjacent_find(
+        zlib_blks.begin(),
+        zlib_blks.end(),
+        [](const puffin::ByteExtent& a, const puffin::ByteExtent& b) {
+          return (a.offset + a.length) > b.offset;
+        });
+    TEST_AND_RETURN_FALSE(result == zlib_blks.end());
+
+    vector<puffin::BitExtent> deflates;
+    TEST_AND_RETURN_FALSE(
+        puffin::LocateDeflatesInZlibBlocks(sqfs_path, zlib_blks, &deflates));
+
+    // Add deflates for each file.
+    for (auto& file : files_) {
+      file.deflates = deflate_utils::FindDeflates(file.extents, deflates);
+    }
+  }
   return true;
 }
 
 unique_ptr<SquashfsFilesystem> SquashfsFilesystem::CreateFromFile(
-    const string& sqfs_path) {
+    const string& sqfs_path, bool extract_deflates) {
   if (sqfs_path.empty())
     return nullptr;
 
@@ -229,11 +281,12 @@
   }
 
   unique_ptr<SquashfsFilesystem> sqfs(new SquashfsFilesystem());
-  if (!sqfs->Init(filemap, sqfs_file->GetSize(), header)) {
+  if (!sqfs->Init(
+          filemap, sqfs_path, sqfs_file->GetSize(), header, extract_deflates)) {
     LOG(ERROR) << "Failed to initialized the Squashfs file system";
     return nullptr;
   }
-  // TODO(ahassani): Add a function that initializes the puffin related extents.
+
   return sqfs;
 }
 
@@ -245,7 +298,7 @@
   }
 
   unique_ptr<SquashfsFilesystem> sqfs(new SquashfsFilesystem());
-  if (!sqfs->Init(filemap, size, header)) {
+  if (!sqfs->Init(filemap, "", size, header, false)) {
     LOG(ERROR) << "Failed to initialize the Squashfs file system using filemap";
     return nullptr;
   }
@@ -276,5 +329,4 @@
   SquashfsHeader header;
   return ReadSquashfsHeader(blob, &header) && CheckHeader(header);
 }
-
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/squashfs_filesystem.h b/payload_generator/squashfs_filesystem.h
index aceebb8..b79f8c7 100644
--- a/payload_generator/squashfs_filesystem.h
+++ b/payload_generator/squashfs_filesystem.h
@@ -55,9 +55,11 @@
 
   ~SquashfsFilesystem() override = default;
 
-  // Creates the file system from the Squashfs file itself.
+  // Creates the file system from the Squashfs file itself. If
+  // |extract_deflates| is true, it will process files to find location of all
+  // deflate streams.
   static std::unique_ptr<SquashfsFilesystem> CreateFromFile(
-      const std::string& sqfs_path);
+      const std::string& sqfs_path, bool extract_deflates);
 
   // Creates the file system from a file map |filemap| which is a multi-line
   // string with each line with the following format:
@@ -99,7 +101,11 @@
   SquashfsFilesystem() = default;
 
   // Initialize and populates the files in the file system.
-  bool Init(const std::string& map, size_t size, const SquashfsHeader& header);
+  bool Init(const std::string& map,
+            const std::string& sqfs_path,
+            size_t size,
+            const SquashfsHeader& header,
+            bool extract_deflates);
 
   // The size of the image in bytes.
   size_t size_;
diff --git a/payload_generator/squashfs_filesystem_unittest.cc b/payload_generator/squashfs_filesystem_unittest.cc
index 77f4ec9..13b4dd3 100644
--- a/payload_generator/squashfs_filesystem_unittest.cc
+++ b/payload_generator/squashfs_filesystem_unittest.cc
@@ -110,7 +110,7 @@
 
 TEST_F(SquashfsFilesystemTest, EmptyFilesystemTest) {
   unique_ptr<SquashfsFilesystem> fs = SquashfsFilesystem::CreateFromFile(
-      GetBuildArtifactsPath("gen/disk_sqfs_empty.img"));
+      GetBuildArtifactsPath("gen/disk_sqfs_empty.img"), true);
   CheckSquashfs(fs);
 
   // Even an empty squashfs filesystem is rounded up to 4K.
@@ -131,7 +131,7 @@
 
 TEST_F(SquashfsFilesystemTest, DefaultFilesystemTest) {
   unique_ptr<SquashfsFilesystem> fs = SquashfsFilesystem::CreateFromFile(
-      GetBuildArtifactsPath("gen/disk_sqfs_default.img"));
+      GetBuildArtifactsPath("gen/disk_sqfs_default.img"), true);
   CheckSquashfs(fs);
 
   vector<FilesystemInterface::File> files;
diff --git a/update_engine.gyp b/update_engine.gyp
index 6f0a140..9f68810 100644
--- a/update_engine.gyp
+++ b/update_engine.gyp
@@ -353,6 +353,7 @@
       'variables': {
         'exported_deps': [
           'ext2fs',
+          'libpuffdiff',
         ],
         'deps': ['<@(exported_deps)'],
       },