libsnapshot: Add support for brotli compression.

Bug: 162274240
Test: cow_api_test
Change-Id: I0b0ceec3c3041a6aea4b1e6c4d01ed0a8860d7e8

fs_mgr/libsnapshot/Android.bp

@@ -152,6 +152,7 @@ cc_library_static {
         "liblog",
     ],
     static_libs: [
+        "libbrotli",
         "libz",
     ],
     ramdisk_available: true,
@@ -362,10 +363,11 @@ cc_defaults {
 
     static_libs: [
         "libbase",
+        "libbrotli",
         "liblog",
         "libdm",
-	"libz",
-	"libsnapshot_cow",
+        "libz",
+        "libsnapshot_cow",
     ],
 }
 
@@ -403,6 +405,7 @@ cc_test {
         "libz",
     ],
     static_libs: [
+        "libbrotli",
         "libgtest",
         "libsnapshot_cow",
     ],
@@ -494,11 +497,12 @@ cc_test {
     shared_libs: [
         "libbase",
         "liblog",
-        "libz",
     ],
     static_libs: [
+        "libbrotli",
         "libgtest",
         "libsnapshot_cow",
+        "libz",
     ],
     header_libs: [
         "libstorage_literals_headers",

fs_mgr/libsnapshot/cow_api_test.cpp

@@ -30,12 +30,12 @@ namespace snapshot {
 
 class CowTest : public ::testing::Test {
   protected:
-    void SetUp() override {
+    virtual void SetUp() override {
         cow_ = std::make_unique<TemporaryFile>();
         ASSERT_GE(cow_->fd, 0) << strerror(errno);
     }
 
-    void TearDown() override { cow_ = nullptr; }
+    virtual void TearDown() override { cow_ = nullptr; }
 
     std::unique_ptr<TemporaryFile> cow_;
 };
@@ -211,9 +211,11 @@ class HorribleStringSink : public StringSink {
     void* GetBuffer(size_t, size_t* actual) override { return StringSink::GetBuffer(1, actual); }
 };
 
-TEST_F(CowTest, HorribleSink) {
+class CompressionTest : public CowTest, public testing::WithParamInterface<const char*> {};
+
+TEST_P(CompressionTest, HorribleSink) {
     CowOptions options;
-    options.compression = "gz";
+    options.compression = GetParam();
     CowWriter writer(options);
 
     ASSERT_TRUE(writer.Initialize(cow_->fd));
@@ -239,6 +241,8 @@ TEST_F(CowTest, HorribleSink) {
     ASSERT_EQ(sink.stream(), data);
 }
 
+INSTANTIATE_TEST_SUITE_P(CowApi, CompressionTest, testing::Values("none", "gz", "brotli"));
+
 TEST_F(CowTest, GetSize) {
     CowOptions options;
     CowWriter writer(options);

fs_mgr/libsnapshot/cow_decompress.cpp

@@ -19,6 +19,7 @@
 #include <utility>
 
 #include <android-base/logging.h>
+#include <brotli/decode.h>
 #include <zlib.h>
 
 namespace android {
@@ -207,5 +208,57 @@ std::unique_ptr<IDecompressor> IDecompressor::Gz() {
     return std::unique_ptr<IDecompressor>(new GzDecompressor());
 }
 
+class BrotliDecompressor final : public StreamDecompressor {
+  public:
+    ~BrotliDecompressor();
+
+    bool Init() override;
+    bool DecompressInput(const uint8_t* data, size_t length) override;
+    bool Done() override { return BrotliDecoderIsFinished(decoder_); }
+
+  private:
+    BrotliDecoderState* decoder_ = nullptr;
+};
+
+bool BrotliDecompressor::Init() {
+    decoder_ = BrotliDecoderCreateInstance(nullptr, nullptr, nullptr);
+    return true;
+}
+
+BrotliDecompressor::~BrotliDecompressor() {
+    if (decoder_) {
+        BrotliDecoderDestroyInstance(decoder_);
+    }
+}
+
+bool BrotliDecompressor::DecompressInput(const uint8_t* data, size_t length) {
+    size_t available_in = length;
+    const uint8_t* next_in = data;
+
+    bool needs_more_output = false;
+    while (available_in || needs_more_output) {
+        if (!output_buffer_remaining_ && !GetFreshBuffer()) {
+            return false;
+        }
+
+        auto output_buffer = output_buffer_;
+        auto r = BrotliDecoderDecompressStream(decoder_, &available_in, &next_in,
+                                               &output_buffer_remaining_, &output_buffer_, nullptr);
+        if (r == BROTLI_DECODER_RESULT_ERROR) {
+            LOG(ERROR) << "brotli decode failed";
+            return false;
+        }
+        if (!sink_->ReturnData(output_buffer, output_buffer_ - output_buffer)) {
+            return false;
+        }
+        needs_more_output = (r == BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT);
+    }
+    return true;
+}
+
+std::unique_ptr<IDecompressor> IDecompressor::Brotli() {
+    return std::unique_ptr<IDecompressor>(new BrotliDecompressor());
+}
+
 }  // namespace snapshot
 }  // namespace android

fs_mgr/libsnapshot/cow_decompress.h

@@ -40,6 +40,7 @@ class IDecompressor {
     // Factory methods for decompression methods.
     static std::unique_ptr<IDecompressor> Uncompressed();
     static std::unique_ptr<IDecompressor> Gz();
+    static std::unique_ptr<IDecompressor> Brotli();
 
     // |output_bytes| is the expected total number of bytes to sink.
     virtual bool Decompress(size_t output_bytes) = 0;

fs_mgr/libsnapshot/cow_reader.cpp

@@ -233,6 +233,9 @@ bool CowReader::ReadData(const CowOperation& op, IByteSink* sink) {
         case kCowCompressGz:
             decompressor = IDecompressor::Gz();
             break;
+        case kCowCompressBrotli:
+            decompressor = IDecompressor::Brotli();
+            break;
         default:
             LOG(ERROR) << "Unknown compression type: " << op.compression;
             return false;

fs_mgr/libsnapshot/cow_writer.cpp

@@ -22,6 +22,7 @@
 #include <android-base/file.h>
 #include <android-base/logging.h>
 #include <android-base/unique_fd.h>
+#include <brotli/encode.h>
 #include <libsnapshot/cow_writer.h>
 #include <zlib.h>
 
@@ -63,6 +64,10 @@ bool CowWriter::Initialize(android::base::borrowed_fd fd) {
 
     if (options_.compression == "gz") {
         compression_ = kCowCompressGz;
+    } else if (options_.compression == "brotli") {
+        compression_ = kCowCompressBrotli;
+    } else if (options_.compression == "none") {
+        compression_ = kCowCompressNone;
     } else if (!options_.compression.empty()) {
         LOG(ERROR) << "unrecognized compression: " << options_.compression;
         return false;
@@ -171,6 +176,24 @@ std::basic_string<uint8_t> CowWriter::Compress(const void* data, size_t length)
             }
             return std::basic_string<uint8_t>(buffer.get(), dest_len);
         }
+        case kCowCompressBrotli: {
+            auto bound = BrotliEncoderMaxCompressedSize(length);
+            if (!bound) {
+                LOG(ERROR) << "BrotliEncoderMaxCompressedSize returned 0";
+                return {};
+            }
+            auto buffer = std::make_unique<uint8_t[]>(bound);
+
+            size_t encoded_size = bound;
+            auto rv = BrotliEncoderCompress(
+                    BROTLI_DEFAULT_QUALITY, BROTLI_DEFAULT_WINDOW, BROTLI_DEFAULT_MODE, length,
+                    reinterpret_cast<const uint8_t*>(data), &encoded_size, buffer.get());
+            if (!rv) {
+                LOG(ERROR) << "BrotliEncoderCompress failed";
+                return {};
+            }
+            return std::basic_string<uint8_t>(buffer.get(), encoded_size);
+        }
         default:
             LOG(ERROR) << "unhandled compression type: " << compression_;
             break;

fs_mgr/libsnapshot/include/libsnapshot/cow_format.h

@@ -98,6 +98,7 @@ static constexpr uint8_t kCowZeroOp = 3;
 
 static constexpr uint8_t kCowCompressNone = 0;
 static constexpr uint8_t kCowCompressGz = 1;
+static constexpr uint8_t kCowCompressBrotli = 2;
 
 }  // namespace snapshot
 }  // namespace android