|  | /* | 
|  | * Copyright (C) 2020 The Android Open Source Project | 
|  | * | 
|  | * Licensed under the Apache License, Version 2.0 (the "License"); | 
|  | * you may not use this file except in compliance with the License. | 
|  | * You may obtain a copy of the License at | 
|  | * | 
|  | *      http://www.apache.org/licenses/LICENSE-2.0 | 
|  | * | 
|  | * Unless required by applicable law or agreed to in writing, software | 
|  | * distributed under the License is distributed on an "AS IS" BASIS, | 
|  | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | 
|  | * See the License for the specific language governing permissions and | 
|  | * limitations under the License. | 
|  | */ | 
|  |  | 
|  | #pragma once | 
|  |  | 
|  | #include <algorithm> | 
|  | #include <memory> | 
|  | #include <span> | 
|  |  | 
|  | #include <android-base/logging.h> | 
|  |  | 
|  | #include <brotli/decode.h> | 
|  | #include <brotli/encode.h> | 
|  | #include <lz4frame.h> | 
|  |  | 
|  | #include "types.h" | 
|  |  | 
|  | enum class DecodeResult { | 
|  | Error, | 
|  | Done, | 
|  | NeedInput, | 
|  | MoreOutput, | 
|  | }; | 
|  |  | 
|  | enum class EncodeResult { | 
|  | Error, | 
|  | Done, | 
|  | NeedInput, | 
|  | MoreOutput, | 
|  | }; | 
|  |  | 
|  | struct Decoder { | 
|  | void Append(Block&& block) { input_buffer_.append(std::move(block)); } | 
|  | bool Finish() { | 
|  | bool old = std::exchange(finished_, true); | 
|  | if (old) { | 
|  | LOG(FATAL) << "Decoder::Finish called while already finished?"; | 
|  | return false; | 
|  | } | 
|  | return true; | 
|  | } | 
|  |  | 
|  | virtual DecodeResult Decode(std::span<char>* output) = 0; | 
|  |  | 
|  | protected: | 
|  | Decoder(std::span<char> output_buffer) : output_buffer_(output_buffer) {} | 
|  | ~Decoder() = default; | 
|  |  | 
|  | bool finished_ = false; | 
|  | IOVector input_buffer_; | 
|  | std::span<char> output_buffer_; | 
|  | }; | 
|  |  | 
|  | struct Encoder { | 
|  | void Append(Block input) { input_buffer_.append(std::move(input)); } | 
|  | bool Finish() { | 
|  | bool old = std::exchange(finished_, true); | 
|  | if (old) { | 
|  | LOG(FATAL) << "Decoder::Finish called while already finished?"; | 
|  | return false; | 
|  | } | 
|  | return true; | 
|  | } | 
|  |  | 
|  | virtual EncodeResult Encode(Block* output) = 0; | 
|  |  | 
|  | protected: | 
|  | explicit Encoder(size_t output_block_size) : output_block_size_(output_block_size) {} | 
|  | ~Encoder() = default; | 
|  |  | 
|  | const size_t output_block_size_; | 
|  | bool finished_ = false; | 
|  | IOVector input_buffer_; | 
|  | }; | 
|  |  | 
|  | struct NullDecoder final : public Decoder { | 
|  | explicit NullDecoder(std::span<char> output_buffer) : Decoder(output_buffer) {} | 
|  |  | 
|  | DecodeResult Decode(std::span<char>* output) final { | 
|  | size_t available_out = output_buffer_.size(); | 
|  | void* p = output_buffer_.data(); | 
|  | while (available_out > 0 && !input_buffer_.empty()) { | 
|  | size_t len = std::min(available_out, input_buffer_.front_size()); | 
|  | p = mempcpy(p, input_buffer_.front_data(), len); | 
|  | available_out -= len; | 
|  | input_buffer_.drop_front(len); | 
|  | } | 
|  | *output = std::span(output_buffer_.data(), static_cast<char*>(p)); | 
|  | if (input_buffer_.empty()) { | 
|  | return finished_ ? DecodeResult::Done : DecodeResult::NeedInput; | 
|  | } | 
|  | return DecodeResult::MoreOutput; | 
|  | } | 
|  | }; | 
|  |  | 
|  | struct NullEncoder final : public Encoder { | 
|  | explicit NullEncoder(size_t output_block_size) : Encoder(output_block_size) {} | 
|  |  | 
|  | EncodeResult Encode(Block* output) final { | 
|  | output->clear(); | 
|  | output->resize(output_block_size_); | 
|  |  | 
|  | size_t available_out = output->size(); | 
|  | void* p = output->data(); | 
|  |  | 
|  | while (available_out > 0 && !input_buffer_.empty()) { | 
|  | size_t len = std::min(available_out, input_buffer_.front_size()); | 
|  | p = mempcpy(p, input_buffer_.front_data(), len); | 
|  | available_out -= len; | 
|  | input_buffer_.drop_front(len); | 
|  | } | 
|  |  | 
|  | output->resize(output->size() - available_out); | 
|  |  | 
|  | if (input_buffer_.empty()) { | 
|  | return finished_ ? EncodeResult::Done : EncodeResult::NeedInput; | 
|  | } | 
|  | return EncodeResult::MoreOutput; | 
|  | } | 
|  | }; | 
|  |  | 
|  | struct BrotliDecoder final : public Decoder { | 
|  | explicit BrotliDecoder(std::span<char> output_buffer) | 
|  | : Decoder(output_buffer), | 
|  | decoder_(BrotliDecoderCreateInstance(nullptr, nullptr, nullptr), | 
|  | BrotliDecoderDestroyInstance) {} | 
|  |  | 
|  | DecodeResult Decode(std::span<char>* output) final { | 
|  | size_t available_in = input_buffer_.front_size(); | 
|  | const uint8_t* next_in = reinterpret_cast<const uint8_t*>(input_buffer_.front_data()); | 
|  |  | 
|  | size_t available_out = output_buffer_.size(); | 
|  | uint8_t* next_out = reinterpret_cast<uint8_t*>(output_buffer_.data()); | 
|  |  | 
|  | BrotliDecoderResult r = BrotliDecoderDecompressStream( | 
|  | decoder_.get(), &available_in, &next_in, &available_out, &next_out, nullptr); | 
|  |  | 
|  | size_t bytes_consumed = input_buffer_.front_size() - available_in; | 
|  | input_buffer_.drop_front(bytes_consumed); | 
|  |  | 
|  | size_t bytes_emitted = output_buffer_.size() - available_out; | 
|  | *output = std::span<char>(output_buffer_.data(), bytes_emitted); | 
|  |  | 
|  | switch (r) { | 
|  | case BROTLI_DECODER_RESULT_SUCCESS: | 
|  | // We need to wait for ID_DONE from the other end. | 
|  | return finished_ ? DecodeResult::Done : DecodeResult::NeedInput; | 
|  | case BROTLI_DECODER_RESULT_ERROR: | 
|  | return DecodeResult::Error; | 
|  | case BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT: | 
|  | // Brotli guarantees as one of its invariants that if it returns NEEDS_MORE_INPUT, | 
|  | // it will consume the entire input buffer passed in, so we don't have to worry | 
|  | // about bytes left over in the front block with more input remaining. | 
|  | return DecodeResult::NeedInput; | 
|  | case BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT: | 
|  | return DecodeResult::MoreOutput; | 
|  | } | 
|  | } | 
|  |  | 
|  | private: | 
|  | std::unique_ptr<BrotliDecoderState, void (*)(BrotliDecoderState*)> decoder_; | 
|  | }; | 
|  |  | 
|  | struct BrotliEncoder final : public Encoder { | 
|  | explicit BrotliEncoder(size_t output_block_size) | 
|  | : Encoder(output_block_size), | 
|  | output_block_(output_block_size_), | 
|  | output_bytes_left_(output_block_size_), | 
|  | encoder_(BrotliEncoderCreateInstance(nullptr, nullptr, nullptr), | 
|  | BrotliEncoderDestroyInstance) { | 
|  | BrotliEncoderSetParameter(encoder_.get(), BROTLI_PARAM_QUALITY, 1); | 
|  | } | 
|  |  | 
|  | EncodeResult Encode(Block* output) final { | 
|  | output->clear(); | 
|  |  | 
|  | while (true) { | 
|  | size_t available_in = input_buffer_.front_size(); | 
|  | const uint8_t* next_in = reinterpret_cast<const uint8_t*>(input_buffer_.front_data()); | 
|  |  | 
|  | size_t available_out = output_bytes_left_; | 
|  | uint8_t* next_out = reinterpret_cast<uint8_t*>( | 
|  | output_block_.data() + (output_block_size_ - output_bytes_left_)); | 
|  |  | 
|  | BrotliEncoderOperation op = BROTLI_OPERATION_PROCESS; | 
|  | if (finished_) { | 
|  | op = BROTLI_OPERATION_FINISH; | 
|  | } | 
|  |  | 
|  | if (!BrotliEncoderCompressStream(encoder_.get(), op, &available_in, &next_in, | 
|  | &available_out, &next_out, nullptr)) { | 
|  | return EncodeResult::Error; | 
|  | } | 
|  |  | 
|  | size_t bytes_consumed = input_buffer_.front_size() - available_in; | 
|  | input_buffer_.drop_front(bytes_consumed); | 
|  |  | 
|  | output_bytes_left_ = available_out; | 
|  |  | 
|  | if (BrotliEncoderIsFinished(encoder_.get())) { | 
|  | output_block_.resize(output_block_size_ - output_bytes_left_); | 
|  | *output = std::move(output_block_); | 
|  | return EncodeResult::Done; | 
|  | } else if (output_bytes_left_ == 0) { | 
|  | *output = std::move(output_block_); | 
|  | output_block_.resize(output_block_size_); | 
|  | output_bytes_left_ = output_block_size_; | 
|  | return EncodeResult::MoreOutput; | 
|  | } else if (input_buffer_.empty()) { | 
|  | return EncodeResult::NeedInput; | 
|  | } | 
|  | } | 
|  | } | 
|  |  | 
|  | private: | 
|  | Block output_block_; | 
|  | size_t output_bytes_left_; | 
|  | std::unique_ptr<BrotliEncoderState, void (*)(BrotliEncoderState*)> encoder_; | 
|  | }; | 
|  |  | 
|  | struct LZ4Decoder final : public Decoder { | 
|  | explicit LZ4Decoder(std::span<char> output_buffer) | 
|  | : Decoder(output_buffer), decoder_(nullptr, nullptr) { | 
|  | LZ4F_dctx* dctx; | 
|  | if (LZ4F_createDecompressionContext(&dctx, LZ4F_VERSION) != 0) { | 
|  | LOG(FATAL) << "failed to initialize LZ4 decompression context"; | 
|  | } | 
|  | decoder_ = std::unique_ptr<LZ4F_dctx, decltype(&LZ4F_freeDecompressionContext)>( | 
|  | dctx, LZ4F_freeDecompressionContext); | 
|  | } | 
|  |  | 
|  | DecodeResult Decode(std::span<char>* output) final { | 
|  | size_t available_in = input_buffer_.front_size(); | 
|  | const char* next_in = input_buffer_.front_data(); | 
|  |  | 
|  | size_t available_out = output_buffer_.size(); | 
|  | char* next_out = output_buffer_.data(); | 
|  |  | 
|  | size_t rc = LZ4F_decompress(decoder_.get(), next_out, &available_out, next_in, | 
|  | &available_in, nullptr); | 
|  | if (LZ4F_isError(rc)) { | 
|  | LOG(ERROR) << "LZ4F_decompress failed: " << LZ4F_getErrorName(rc); | 
|  | return DecodeResult::Error; | 
|  | } | 
|  |  | 
|  | input_buffer_.drop_front(available_in); | 
|  |  | 
|  | if (rc == 0) { | 
|  | if (!input_buffer_.empty()) { | 
|  | LOG(ERROR) << "LZ4 stream hit end before reading all data"; | 
|  | return DecodeResult::Error; | 
|  | } | 
|  | lz4_done_ = true; | 
|  | } | 
|  |  | 
|  | *output = std::span<char>(output_buffer_.data(), available_out); | 
|  |  | 
|  | if (finished_) { | 
|  | return input_buffer_.empty() && lz4_done_ ? DecodeResult::Done | 
|  | : DecodeResult::MoreOutput; | 
|  | } | 
|  |  | 
|  | return DecodeResult::NeedInput; | 
|  | } | 
|  |  | 
|  | private: | 
|  | bool lz4_done_ = false; | 
|  | std::unique_ptr<LZ4F_dctx, LZ4F_errorCode_t (*)(LZ4F_dctx*)> decoder_; | 
|  | }; | 
|  |  | 
|  | struct LZ4Encoder final : public Encoder { | 
|  | explicit LZ4Encoder(size_t output_block_size) | 
|  | : Encoder(output_block_size), encoder_(nullptr, nullptr) { | 
|  | LZ4F_cctx* cctx; | 
|  | if (LZ4F_createCompressionContext(&cctx, LZ4F_VERSION) != 0) { | 
|  | LOG(FATAL) << "failed to initialize LZ4 compression context"; | 
|  | } | 
|  | encoder_ = std::unique_ptr<LZ4F_cctx, decltype(&LZ4F_freeCompressionContext)>( | 
|  | cctx, LZ4F_freeCompressionContext); | 
|  | Block header(LZ4F_HEADER_SIZE_MAX); | 
|  | size_t rc = LZ4F_compressBegin(encoder_.get(), header.data(), header.size(), nullptr); | 
|  | if (LZ4F_isError(rc)) { | 
|  | LOG(FATAL) << "LZ4F_compressBegin failed: %s", LZ4F_getErrorName(rc); | 
|  | } | 
|  | header.resize(rc); | 
|  | output_buffer_.append(std::move(header)); | 
|  | } | 
|  |  | 
|  | // As an optimization, only emit a block if we have an entire output block ready, or we're done. | 
|  | bool OutputReady() const { | 
|  | return output_buffer_.size() >= output_block_size_ || lz4_finalized_; | 
|  | } | 
|  |  | 
|  | // TODO: Switch the output type to IOVector to remove a copy? | 
|  | EncodeResult Encode(Block* output) final { | 
|  | size_t available_in = input_buffer_.front_size(); | 
|  | const char* next_in = input_buffer_.front_data(); | 
|  |  | 
|  | // LZ4 makes no guarantees about being able to recover from trying to compress with an | 
|  | // insufficiently large output buffer. LZ4F_compressBound tells us how much buffer we | 
|  | // need to compress a given number of bytes, but the smallest value seems to be bigger | 
|  | // than SYNC_DATA_MAX, so we need to buffer ourselves. | 
|  |  | 
|  | // Input size chosen to be a local maximum for LZ4F_compressBound (i.e. the block size). | 
|  | constexpr size_t max_input_size = 65536; | 
|  | const size_t encode_block_size = LZ4F_compressBound(max_input_size, nullptr); | 
|  |  | 
|  | if (available_in != 0) { | 
|  | if (lz4_finalized_) { | 
|  | LOG(ERROR) << "LZ4Encoder received data after Finish?"; | 
|  | return EncodeResult::Error; | 
|  | } | 
|  |  | 
|  | available_in = std::min(available_in, max_input_size); | 
|  |  | 
|  | Block encode_block(encode_block_size); | 
|  | size_t available_out = encode_block.capacity(); | 
|  | char* next_out = encode_block.data(); | 
|  |  | 
|  | size_t rc = LZ4F_compressUpdate(encoder_.get(), next_out, available_out, next_in, | 
|  | available_in, nullptr); | 
|  | if (LZ4F_isError(rc)) { | 
|  | LOG(ERROR) << "LZ4F_compressUpdate failed: " << LZ4F_getErrorName(rc); | 
|  | return EncodeResult::Error; | 
|  | } | 
|  |  | 
|  | input_buffer_.drop_front(available_in); | 
|  |  | 
|  | available_out -= rc; | 
|  | next_out += rc; | 
|  |  | 
|  | encode_block.resize(encode_block_size - available_out); | 
|  | output_buffer_.append(std::move(encode_block)); | 
|  | } | 
|  |  | 
|  | if (finished_ && !lz4_finalized_) { | 
|  | lz4_finalized_ = true; | 
|  |  | 
|  | Block final_block(encode_block_size + 4); | 
|  | size_t rc = LZ4F_compressEnd(encoder_.get(), final_block.data(), final_block.size(), | 
|  | nullptr); | 
|  | if (LZ4F_isError(rc)) { | 
|  | LOG(ERROR) << "LZ4F_compressEnd failed: " << LZ4F_getErrorName(rc); | 
|  | return EncodeResult::Error; | 
|  | } | 
|  |  | 
|  | final_block.resize(rc); | 
|  | output_buffer_.append(std::move(final_block)); | 
|  | } | 
|  |  | 
|  | if (OutputReady()) { | 
|  | size_t len = std::min(output_block_size_, output_buffer_.size()); | 
|  | *output = output_buffer_.take_front(len).coalesce(); | 
|  | } else { | 
|  | output->clear(); | 
|  | } | 
|  |  | 
|  | if (lz4_finalized_ && output_buffer_.empty()) { | 
|  | return EncodeResult::Done; | 
|  | } else if (OutputReady()) { | 
|  | return EncodeResult::MoreOutput; | 
|  | } | 
|  | return EncodeResult::NeedInput; | 
|  | } | 
|  |  | 
|  | private: | 
|  | bool lz4_finalized_ = false; | 
|  | std::unique_ptr<LZ4F_cctx, LZ4F_errorCode_t (*)(LZ4F_cctx*)> encoder_; | 
|  | IOVector output_buffer_; | 
|  | }; |