update_engine: Apply PUFFDIFF operation

This patch adds operation for applying PUFFDIFF operation. As long as
the minor version is not increased, this patch will not be directly
used.

This patch sets the cache requirement for PUFFDIFF operation to 5MB as it has
been selected for now from sample data below:

Update from ToT(R64) to veyron-minnie-R63-9928.0.0
Cache Size (MB)	- Time to Update (minutes)
no-puffdiff	0:04:20
50	0:05:30
20	0:06:40
10	0:09:10
5	0:11:00
3	0:13:40
2	0:15:45
1	0:19:30

BUG=chromium:717785
TEST=brillo_update_payload verify passes; unittests pass;
CQ-DEPEND=CL:722452

Change-Id: I481481d40ef9df6412782198e3c97deb3997cc2e
Reviewed-on: https://chromium-review.googlesource.com/658297
Commit-Ready: Amin Hassani <ahassani@chromium.org>
Tested-by: Amin Hassani <ahassani@chromium.org>
Reviewed-by: Ben Chan <benchan@chromium.org>
Reviewed-by: Sen Jiang <senj@chromium.org>
diff --git a/Android.mk b/Android.mk
index 08798aa..523ff6f 100644
--- a/Android.mk
+++ b/Android.mk
@@ -101,6 +101,7 @@
     libxz \
     libbz \
     libbspatch \
+    libpuffpatch \
     $(ue_update_metadata_protos_exported_static_libraries)
 ue_libpayload_consumer_exported_shared_libraries := \
     libcrypto \
diff --git a/payload_consumer/delta_performer.cc b/payload_consumer/delta_performer.cc
index 78a8a42..c3fd411 100644
--- a/payload_consumer/delta_performer.cc
+++ b/payload_consumer/delta_performer.cc
@@ -37,6 +37,7 @@
 #include <brillo/data_encoding.h>
 #include <bsdiff/bspatch.h>
 #include <google/protobuf/repeated_field.h>
+#include <puffin/puffpatch.h>
 
 #include "update_engine/common/constants.h"
 #include "update_engine/common/hardware_interface.h"
@@ -771,8 +772,8 @@
         OP_DURATION_HISTOGRAM("SOURCE_BSDIFF", op_start_time);
         break;
       case InstallOperation::PUFFDIFF:
-        // TODO(ahassani): Later add PerformPuffdiffOperation(op, error);
-        op_result = false;
+        op_result = PerformPuffDiffOperation(op, error);
+        OP_DURATION_HISTOGRAM("PUFFDIFF", op_start_time);
         break;
       default:
         op_result = false;
@@ -1245,6 +1246,121 @@
   return true;
 }
 
+namespace {
+
+// A class to be passed to |puffpatch| for reading from |source_fd_| and writing
+// into |target_fd_|.
+class PuffinExtentStream : public puffin::StreamInterface {
+ public:
+  // Constructor for creating a stream for reading from an |ExtentReader|.
+  PuffinExtentStream(std::unique_ptr<ExtentReader> reader, size_t size)
+      : PuffinExtentStream(std::move(reader), nullptr, size) {}
+
+  // Constructor for creating a stream for writing to an |ExtentWriter|.
+  PuffinExtentStream(std::unique_ptr<ExtentWriter> writer, size_t size)
+      : PuffinExtentStream(nullptr, std::move(writer), size) {}
+
+  ~PuffinExtentStream() override = default;
+
+  bool GetSize(size_t* size) const override {
+    *size = size_;
+    return true;
+  }
+
+  bool GetOffset(size_t* offset) const override {
+    *offset = offset_;
+    return true;
+  }
+
+  bool Seek(size_t offset) override {
+    if (is_read_) {
+      TEST_AND_RETURN_FALSE(reader_->Seek(offset));
+      offset_ = offset;
+    } else {
+      // For writes technically there should be no change of position, or it
+      // should equivalent of current offset.
+      TEST_AND_RETURN_FALSE(offset_ == offset);
+    }
+    return true;
+  }
+
+  bool Read(void* buffer, size_t count) override {
+    TEST_AND_RETURN_FALSE(is_read_);
+    TEST_AND_RETURN_FALSE(reader_->Read(buffer, count));
+    offset_ += count;
+    return true;
+  }
+
+  bool Write(const void* buffer, size_t count) override {
+    TEST_AND_RETURN_FALSE(!is_read_);
+    TEST_AND_RETURN_FALSE(writer_->Write(buffer, count));
+    offset_ += count;
+    return true;
+  }
+
+  bool Close() override {
+    if (!is_read_) {
+      TEST_AND_RETURN_FALSE(writer_->End());
+    }
+    return true;
+  }
+
+ private:
+  PuffinExtentStream(std::unique_ptr<ExtentReader> reader,
+                     std::unique_ptr<ExtentWriter> writer,
+                     size_t size)
+      : reader_(std::move(reader)),
+        writer_(std::move(writer)),
+        size_(size),
+        offset_(0),
+        is_read_(reader_ ? true : false) {}
+
+  std::unique_ptr<ExtentReader> reader_;
+  std::unique_ptr<ExtentWriter> writer_;
+  uint64_t size_;
+  uint64_t offset_;
+  bool is_read_;
+
+  DISALLOW_COPY_AND_ASSIGN(PuffinExtentStream);
+};
+
+}  // namespace
+
+bool DeltaPerformer::PerformPuffDiffOperation(const InstallOperation& operation,
+                                              ErrorCode* error) {
+  // Since we delete data off the beginning of the buffer as we use it,
+  // the data we need should be exactly at the beginning of the buffer.
+  TEST_AND_RETURN_FALSE(buffer_offset_ == operation.data_offset());
+  TEST_AND_RETURN_FALSE(buffer_.size() >= operation.data_length());
+
+  if (operation.has_src_sha256_hash()) {
+    TEST_AND_RETURN_FALSE(CalculateAndValidateSourceHash(operation, error));
+  }
+
+  auto reader = std::make_unique<DirectExtentReader>();
+  TEST_AND_RETURN_FALSE(
+      reader->Init(source_fd_, operation.src_extents(), block_size_));
+  puffin::UniqueStreamPtr src_stream(new PuffinExtentStream(
+      std::move(reader),
+      utils::BlocksInExtents(operation.src_extents()) * block_size_));
+
+  auto writer = std::make_unique<DirectExtentWriter>();
+  TEST_AND_RETURN_FALSE(
+      writer->Init(target_fd_, operation.dst_extents(), block_size_));
+  puffin::UniqueStreamPtr dst_stream(new PuffinExtentStream(
+      std::move(writer),
+      utils::BlocksInExtents(operation.dst_extents()) * block_size_));
+
+  const size_t kMaxCacheSize = 5 * 1024 * 1024;  // Total 5MB cache.
+  TEST_AND_RETURN_FALSE(puffin::PuffPatch(std::move(src_stream),
+                                          std::move(dst_stream),
+                                          buffer_.data(),
+                                          buffer_.size(),
+                                          kMaxCacheSize));
+  DiscardBuffer(true, buffer_.size());
+  return true;
+}
+
 bool DeltaPerformer::ExtractSignatureMessageFromOperation(
     const InstallOperation& operation) {
   if (operation.type() != InstallOperation::REPLACE ||
diff --git a/payload_consumer/delta_performer.h b/payload_consumer/delta_performer.h
index 5fc972b..8aa036e 100644
--- a/payload_consumer/delta_performer.h
+++ b/payload_consumer/delta_performer.h
@@ -261,6 +261,8 @@
                                   ErrorCode* error);
   bool PerformSourceBsdiffOperation(const InstallOperation& operation,
                                     ErrorCode* error);
+  bool PerformPuffDiffOperation(const InstallOperation& operation,
+                                ErrorCode* error);
 
   // Extracts the payload signature message from the blob on the |operation| if
   // the offset matches the one specified by the manifest. Returns whether the
diff --git a/payload_consumer/delta_performer_unittest.cc b/payload_consumer/delta_performer_unittest.cc
index fc6a4b5..ea35c47 100644
--- a/payload_consumer/delta_performer_unittest.cc
+++ b/payload_consumer/delta_performer_unittest.cc
@@ -87,6 +87,61 @@
     0x00, 0x00, 0x59, 0x5a,
 };
 
+const uint8_t src_deflates[] = {
+  /* raw      0  */ 0x11, 0x22,
+  /* deflate  2  */ 0x63, 0x64, 0x62, 0x66, 0x61, 0x05, 0x00,
+  /* raw      9  */ 0x33,
+  /* deflate  10 */ 0x03, 0x00,
+  /* raw      12 */
+  /* deflate  12 */ 0x63, 0x04, 0x00,
+  /* raw      15 */ 0x44, 0x55
+};
+
+const uint8_t dst_deflates[] = {
+  /* deflate  0  */ 0x63, 0x64, 0x62, 0x66, 0x61, 0x05, 0x00,
+  /* raw      7  */ 0x33, 0x66,
+  /* deflate  9  */ 0x01, 0x05, 0x00, 0xFA, 0xFF, 0x01, 0x02, 0x03, 0x04, 0x05,
+  /* deflate  19 */ 0x63, 0x04, 0x00
+};
+
+// To generate this patch either:
+// - Use puffin/src/patching_unittest.cc:TestPatching
+// Or
+// - Use the following approach:
+// * Make src_deflate a string of hex with only spaces. (e.g. "0XTE 0xST")
+// * echo "0XTE 0xST" | xxd -r -p > src.bin
+// * Find the location of deflates in src_deflates (in bytes) in the format of
+//   "offset:length,...". (e.g. "2:7,10:2,12:3")
+// * Do previous three steps for dst_deflates.
+// * puffin --operation=puffdiff --src_file=src.bin --dst_file=dst.bin \
+//   --src_deflates_byte="2:7,10:2,12:3" --dst_deflates_byte="0:7,9:10,19:3" \
+//   --patch_file=patch.bin
+// * hexdump -ve '"  " 12/1 "0x%02x, " "\n"' patch.bin
+const uint8_t puffdiff_patch[] = {
+  0x50, 0x55, 0x46, 0x31, 0x00, 0x00, 0x00, 0x51, 0x08, 0x01, 0x12, 0x27,
+  0x0A, 0x04, 0x08, 0x10, 0x10, 0x32, 0x0A, 0x04, 0x08, 0x50, 0x10, 0x0A,
+  0x0A, 0x04, 0x08, 0x60, 0x10, 0x12, 0x12, 0x04, 0x08, 0x10, 0x10, 0x58,
+  0x12, 0x04, 0x08, 0x78, 0x10, 0x28, 0x12, 0x05, 0x08, 0xA8, 0x01, 0x10,
+  0x38, 0x18, 0x1F, 0x1A, 0x24, 0x0A, 0x02, 0x10, 0x32, 0x0A, 0x04, 0x08,
+  0x48, 0x10, 0x50, 0x0A, 0x05, 0x08, 0x98, 0x01, 0x10, 0x12, 0x12, 0x02,
+  0x10, 0x58, 0x12, 0x04, 0x08, 0x70, 0x10, 0x58, 0x12, 0x05, 0x08, 0xC8,
+  0x01, 0x10, 0x38, 0x18, 0x21, 0x42, 0x53, 0x44, 0x49, 0x46, 0x46, 0x34,
+  0x30, 0x38, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x34, 0x00, 0x00,
+  0x00, 0x00, 0x00, 0x00, 0x00, 0x21, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+  0x00, 0x42, 0x5A, 0x68, 0x39, 0x31, 0x41, 0x59, 0x26, 0x53, 0x59, 0x65,
+  0x29, 0x8C, 0x9B, 0x00, 0x00, 0x03, 0x60, 0x40, 0x7A, 0x0E, 0x08, 0x00,
+  0x40, 0x00, 0x20, 0x00, 0x21, 0x22, 0x9A, 0x3D, 0x4F, 0x50, 0x40, 0x0C,
+  0x3B, 0xC7, 0x9B, 0xB2, 0x21, 0x0E, 0xE9, 0x15, 0x98, 0x7A, 0x7C, 0x5D,
+  0xC9, 0x14, 0xE1, 0x42, 0x41, 0x94, 0xA6, 0x32, 0x6C, 0x42, 0x5A, 0x68,
+  0x39, 0x31, 0x41, 0x59, 0x26, 0x53, 0x59, 0xF1, 0x20, 0x5F, 0x0D, 0x00,
+  0x00, 0x02, 0x41, 0x15, 0x42, 0x08, 0x20, 0x00, 0x40, 0x00, 0x00, 0x02,
+  0x40, 0x00, 0x20, 0x00, 0x22, 0x3D, 0x23, 0x10, 0x86, 0x03, 0x96, 0x54,
+  0x11, 0x16, 0x5F, 0x17, 0x72, 0x45, 0x38, 0x50, 0x90, 0xF1, 0x20, 0x5F,
+  0x0D, 0x42, 0x5A, 0x68, 0x39, 0x31, 0x41, 0x59, 0x26, 0x53, 0x59, 0x07,
+  0xD4, 0xCB, 0x6E, 0x00, 0x00, 0x00, 0x01, 0x00, 0x01, 0x00, 0x20, 0x00,
+  0x21, 0x18, 0x46, 0x82, 0xEE, 0x48, 0xA7, 0x0A, 0x12, 0x00, 0xFA, 0x99,
+  0x6D, 0xC0};
+
 }  // namespace
 
 class DeltaPerformerTest : public ::testing::Test {
@@ -484,6 +539,32 @@
   EXPECT_EQ(expected_data, ApplyPayload(payload_data, source_path, true));
 }
 
+TEST_F(DeltaPerformerTest, PuffdiffOperationTest) {
+  AnnotatedOperation aop;
+  *(aop.op.add_src_extents()) = ExtentForRange(0, 1);
+  *(aop.op.add_dst_extents()) = ExtentForRange(0, 1);
+  brillo::Blob puffdiff_payload(std::begin(puffdiff_patch),
+                                std::end(puffdiff_patch));
+  aop.op.set_data_offset(0);
+  aop.op.set_data_length(puffdiff_payload.size());
+  aop.op.set_type(InstallOperation::PUFFDIFF);
+  brillo::Blob src(std::begin(src_deflates), std::end(src_deflates));
+  src.resize(4096);  // block size
+  brillo::Blob src_hash;
+  EXPECT_TRUE(HashCalculator::RawHashOfData(src, &src_hash));
+  aop.op.set_src_sha256_hash(src_hash.data(), src_hash.size());
+
+  brillo::Blob payload_data = GeneratePayload(puffdiff_payload, {aop}, false);
+
+  string source_path;
+  EXPECT_TRUE(utils::MakeTempFile("Source-XXXXXX", &source_path, nullptr));
+  ScopedPathUnlinker path_unlinker(source_path);
+  EXPECT_TRUE(utils::WriteFile(source_path.c_str(), src.data(), src.size()));
+
+  brillo::Blob dst(std::begin(dst_deflates), std::end(dst_deflates));
+  EXPECT_EQ(dst, ApplyPayload(payload_data, source_path, true));
+}
+
 TEST_F(DeltaPerformerTest, SourceHashMismatchTest) {
   brillo::Blob expected_data = {'f', 'o', 'o'};
   brillo::Blob actual_data = {'b', 'a', 'r'};
diff --git a/update_engine.gyp b/update_engine.gyp
index 9f68810..4da04d4 100644
--- a/update_engine.gyp
+++ b/update_engine.gyp
@@ -130,6 +130,7 @@
         'exported_deps': [
           'libcrypto',
           'xz-embedded',
+          'libpuffpatch',
         ],
         'deps': ['<@(exported_deps)'],
       },