Add a brotli compressor for zucchini patches am: d960ef39d7

Original change: https://android-review.googlesource.com/c/platform/external/puffin/+/1804858

Change-Id: Ib6bd1d825681d9d6cc9331f178e18ea9d004047d

6 files changed, 192 insertions, 5 deletions

diff --git a/Android.bp b/Android.bp
index decde42..022be2e 100644
--- a/Android.bp
+++ b/Android.bp
@@ -50,8 +50,10 @@ cc_library_static {
         "puffin/src/puffin.proto",
         "src/bit_reader.cc",
         "src/bit_writer.cc",
+        "src/brotli_util.cc",
         "src/huffer.cc",
         "src/huffman_table.cc",
+        "src/memory_stream.cc",
         "src/puff_reader.cc",
         "src/puff_writer.cc",
         "src/puffer.cc",
@@ -75,7 +77,6 @@ cc_library_static {
     defaults: ["puffin_defaults"],
     srcs: [
         "src/file_stream.cc",
-        "src/memory_stream.cc",
         "src/puffdiff.cc",
         "src/utils.cc",
     ],
@@ -114,6 +115,7 @@ cc_test {
     cflags: ["-Wno-sign-compare"],
     srcs: [
         "src/bit_io_unittest.cc",
+        "src/brotli_util_unittest.cc",
         "src/extent_stream.cc",
         "src/integration_test.cc",
         "src/patching_unittest.cc",
diff --git a/src/brotli_util.cc b/src/brotli_util.cc
new file mode 100644
index 0000000..e36e98d
--- /dev/null
+++ b/src/brotli_util.cc
@@ -0,0 +1,104 @@
+// Copyright 2021 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/include/puffin/brotli_util.h"
+
+#include "brotli/decode.h"
+#include "brotli/encode.h"
+#include "puffin/src/logging.h"
+#include "puffin/src/memory_stream.h"
+
+namespace puffin {
+
+namespace {
+
+constexpr auto kBufferSize = 32768;
+constexpr auto kDefaultParamQuality = 9;
+constexpr auto kDefaultParamLgwin = 20;
+}  // namespace
+
+bool BrotliEncode(const uint8_t* input,
+                  size_t input_size,
+                  UniqueStreamPtr output_stream) {
+  std::unique_ptr<BrotliEncoderState, decltype(&BrotliEncoderDestroyInstance)>
+      encoder(BrotliEncoderCreateInstance(nullptr, nullptr, nullptr),
+              BrotliEncoderDestroyInstance);
+  TEST_AND_RETURN_FALSE(encoder != nullptr);
+
+  BrotliEncoderSetParameter(encoder.get(), BROTLI_PARAM_QUALITY,
+                            kDefaultParamQuality);
+  BrotliEncoderSetParameter(encoder.get(), BROTLI_PARAM_LGWIN,
+                            kDefaultParamLgwin);
+
+  size_t available_in = input_size;
+  while (available_in != 0 || !BrotliEncoderIsFinished(encoder.get())) {
+    const uint8_t* next_in = input + input_size - available_in;
+    // Set up the output buffer
+    uint8_t buffer[kBufferSize];
+    uint8_t* next_out = buffer;
+    size_t available_out = kBufferSize;
+
+    BrotliEncoderOperation op =
+        available_in == 0 ? BROTLI_OPERATION_FINISH : BROTLI_OPERATION_PROCESS;
+
+    if (!BrotliEncoderCompressStream(encoder.get(), op, &available_in, &next_in,
+                                     &available_out, &next_out, nullptr)) {
+      LOG(ERROR) << "Failed to compress " << input_size << " bytes with brotli";
+      return false;
+    }
+
+    size_t bytes_consumed = kBufferSize - available_out;
+    output_stream->Write(buffer, bytes_consumed);
+  }
+
+  return true;
+}
+
+bool BrotliEncode(const uint8_t* input,
+                  size_t input_size,
+                  std::vector<uint8_t>* output) {
+  TEST_AND_RETURN_FALSE(output != nullptr);
+  return BrotliEncode(input, input_size, MemoryStream::CreateForWrite(output));
+}
+
+bool BrotliDecode(const uint8_t* input,
+                  size_t input_size,
+                  UniqueStreamPtr output_stream) {
+  std::unique_ptr<BrotliDecoderState, decltype(&BrotliDecoderDestroyInstance)>
+      decoder(BrotliDecoderCreateInstance(nullptr, nullptr, nullptr),
+              BrotliDecoderDestroyInstance);
+  TEST_AND_RETURN_FALSE(decoder != nullptr);
+
+  size_t available_in = input_size;
+  while (available_in != 0 || !BrotliDecoderIsFinished(decoder.get())) {
+    const uint8_t* next_in = input + input_size - available_in;
+    // Set up the output buffer
+    uint8_t buffer[kBufferSize];
+    uint8_t* next_out = buffer;
+    size_t available_out = kBufferSize;
+
+    BrotliDecoderResult result =
+        BrotliDecoderDecompressStream(decoder.get(), &available_in, &next_in,
+                                      &available_out, &next_out, nullptr);
+    if (result == BROTLI_DECODER_RESULT_ERROR ||
+        result == BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT) {
+      LOG(ERROR) << "Failed to decompress " << input_size
+                 << " bytes with brotli, result " << result;
+      return false;
+    }
+
+    size_t bytes_consumed = kBufferSize - available_out;
+    output_stream->Write(buffer, bytes_consumed);
+  }
+  return true;
+}
+
+bool BrotliDecode(const uint8_t* input,
+                  size_t input_size,
+                  std::vector<uint8_t>* output) {
+  TEST_AND_RETURN_FALSE(output != nullptr);
+  return BrotliDecode(input, input_size, MemoryStream::CreateForWrite(output));
+}
+
+}  // namespace puffin
diff --git a/src/brotli_util_unittest.cc b/src/brotli_util_unittest.cc
new file mode 100644
index 0000000..3f4685a
--- /dev/null
+++ b/src/brotli_util_unittest.cc
@@ -0,0 +1,32 @@
+// Copyright 2021 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 "gtest/gtest.h"
+
+#include "puffin/src/include/puffin/brotli_util.h"
+#include "puffin/src/memory_stream.h"
+#include "puffin/src/puffin_stream.h"
+
+namespace puffin {
+
+namespace {
+
+// echo "puffin test" | xxd -i
+const Buffer kTestString = {0x70, 0x75, 0x66, 0x66, 0x69, 0x6e,
+                            0x20, 0x74, 0x65, 0x73, 0x74, 0x0a};
+}  // namespace
+
+TEST(BrotliUtilTest, CompressAndDecompressTest) {
+  Buffer compressed;
+  ASSERT_TRUE(BrotliEncode(kTestString.data(), kTestString.size(),
+                           MemoryStream::CreateForWrite(&compressed)));
+  ASSERT_FALSE(compressed.empty());
+
+  Buffer decompressed;
+  ASSERT_TRUE(BrotliDecode(compressed.data(), compressed.size(),
+                           MemoryStream::CreateForWrite(&decompressed)));
+  ASSERT_EQ(kTestString, decompressed);
+}
+
+}  // namespace puffin
diff --git a/src/include/puffin/brotli_util.h b/src/include/puffin/brotli_util.h
new file mode 100644
index 0000000..07240c3
--- /dev/null
+++ b/src/include/puffin/brotli_util.h
@@ -0,0 +1,37 @@
+// Copyright 2021 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.
+
+#ifndef SRC_INCLUDE_PUFFIN_BROTLI_UTIL_H_
+#define SRC_INCLUDE_PUFFIN_BROTLI_UTIL_H_
+
+#include <stdint.h>
+
+#include <vector>
+
+#include "puffin/stream.h"
+
+namespace puffin {
+
+// Use brotli to compress |input_size| bytes of data, and write the result to
+// |output_stream|.
+bool BrotliEncode(const uint8_t* input,
+                  size_t input_size,
+                  UniqueStreamPtr output_stream);
+// Similar to above function, and writes to a output buffer.
+bool BrotliEncode(const uint8_t* input,
+                  size_t input_size,
+                  std::vector<uint8_t>* output);
+
+// Decompress |input_size| bytes of data with brotli, and write the result to
+// |output_stream|.
+bool BrotliDecode(const uint8_t* input,
+                  size_t input_size,
+                  UniqueStreamPtr output_stream);
+// Similar to above function, and writes to a output buffer.
+bool BrotliDecode(const uint8_t* input,
+                  size_t input_size,
+                  std::vector<uint8_t>* output);
+}  // namespace puffin
+
+#endif  // SRC_INCLUDE_PUFFIN_BROTLI_UTIL_H_
diff --git a/src/puffdiff.cc b/src/puffdiff.cc
index 6d9bc27..7f7d9ce 100644
--- a/src/puffdiff.cc
+++ b/src/puffdiff.cc
@@ -18,6 +18,7 @@
 #include "zucchini/zucchini.h"
 
 #include "puffin/src/file_stream.h"
+#include "puffin/src/include/puffin/brotli_util.h"
 #include "puffin/src/include/puffin/common.h"
 #include "puffin/src/include/puffin/puffer.h"
 #include "puffin/src/include/puffin/puffpatch.h"
@@ -165,14 +166,21 @@ bool PuffDiff(UniqueStreamPtr src,
 
     zucchini::EnsemblePatchWriter patch_writer(src_bytes, dst_bytes);
     auto status = zucchini::GenerateBuffer(src_bytes, dst_bytes, &patch_writer);
-    CHECK_EQ(zucchini::status::kStatusSuccess, status);
+    TEST_AND_RETURN_FALSE(status == zucchini::status::kStatusSuccess);
 
     Buffer zucchini_patch_buf(patch_writer.SerializedSize());
     patch_writer.SerializeInto(
         {zucchini_patch_buf.data(), zucchini_patch_buf.size()});
 
+    // Use brotli to compress the zucchini patch.
+    // TODO(197361113) respect the CompressorType parameter for zucchini.
+    Buffer compressed_patch;
+    TEST_AND_RETURN_FALSE(BrotliEncode(zucchini_patch_buf.data(),
+                                       zucchini_patch_buf.size(),
+                                       &compressed_patch));
+
     TEST_AND_RETURN_FALSE(CreatePatch(
-        zucchini_patch_buf, src_deflates, dst_deflates, src_puffs, dst_puffs,
+        compressed_patch, src_deflates, dst_deflates, src_puffs, dst_puffs,
         src_puff_buffer.size(), dst_puff_buffer.size(), patchAlgorithm, patch));
   } else {
     LOG(ERROR) << "unsupported type " << static_cast<int>(patchAlgorithm);
diff --git a/src/puffpatch.cc b/src/puffpatch.cc
index 09bfb98..3802b17 100644
--- a/src/puffpatch.cc
+++ b/src/puffpatch.cc
@@ -17,11 +17,13 @@
 #include "zucchini/patch_reader.h"
 #include "zucchini/zucchini.h"
 
+#include "puffin/src/include/puffin/brotli_util.h"
 #include "puffin/src/include/puffin/common.h"
 #include "puffin/src/include/puffin/huffer.h"
 #include "puffin/src/include/puffin/puffer.h"
 #include "puffin/src/include/puffin/stream.h"
 #include "puffin/src/logging.h"
+#include "puffin/src/memory_stream.h"
 #include "puffin/src/puffin.pb.h"
 #include "puffin/src/puffin_stream.h"
 
@@ -159,8 +161,10 @@ bool ApplyZucchiniPatch(UniqueStreamPtr src_stream,
     bytes_wrote += write_size;
   }
   // Read the patch
-  auto patch_reader =
-      zucchini::EnsemblePatchReader::Create({patch_start, patch_size});
+  Buffer zucchini_patch;
+  TEST_AND_RETURN_FALSE(BrotliDecode(patch_start, patch_size, &zucchini_patch));
+  auto patch_reader = zucchini::EnsemblePatchReader::Create(
+      {zucchini_patch.data(), zucchini_patch.size()});
   if (!patch_reader.has_value()) {
     LOG(ERROR) << "Failed to parse the zucchini patch.";
     return false;