fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.cpp - third_party/android/platform/system/core - Git at Google

 //
 // 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.
 //

 #include "cow_decompress.h"

 #include <utility>

 #include <android-base/logging.h>
 #include <brotli/decode.h>
 #include <lz4.h>
 #include <zlib.h>

 namespace android {
 namespace snapshot {

 class NoDecompressor final : public IDecompressor {
   public:
     bool Decompress(size_t) override;
 };

 bool NoDecompressor::Decompress(size_t) {
     size_t stream_remaining = stream_->Size();
     while (stream_remaining) {
         size_t buffer_size = stream_remaining;
         uint8_t* buffer = reinterpret_cast<uint8_t*>(sink_->GetBuffer(buffer_size, &buffer_size));
         if (!buffer) {
             LOG(ERROR) << "Could not acquire buffer from sink";
             return false;
         }

         // Read until we can fill the buffer.
         uint8_t* buffer_pos = buffer;
         size_t bytes_to_read = std::min(buffer_size, stream_remaining);
         while (bytes_to_read) {
             size_t read;
             if (!stream_->Read(buffer_pos, bytes_to_read, &read)) {
                 return false;
             }
             if (!read) {
                 LOG(ERROR) << "Stream ended prematurely";
                 return false;
             }
             if (!sink_->ReturnData(buffer_pos, read)) {
                 LOG(ERROR) << "Could not return buffer to sink";
                 return false;
             }
             buffer_pos += read;
             bytes_to_read -= read;
             stream_remaining -= read;
         }
     }
     return true;
 }

 std::unique_ptr<IDecompressor> IDecompressor::Uncompressed() {
     return std::unique_ptr<IDecompressor>(new NoDecompressor());
 }

 // Read chunks of the COW and incrementally stream them to the decoder.
 class StreamDecompressor : public IDecompressor {
   public:
     bool Decompress(size_t output_bytes) override;

     virtual bool Init() = 0;
     virtual bool DecompressInput(const uint8_t* data, size_t length) = 0;
     virtual bool Done() = 0;

   protected:
     bool GetFreshBuffer();

     size_t output_bytes_;
     size_t stream_remaining_;
     uint8_t* output_buffer_ = nullptr;
     size_t output_buffer_remaining_ = 0;
 };

 static constexpr size_t kChunkSize = 4096;

 bool StreamDecompressor::Decompress(size_t output_bytes) {
     if (!Init()) {
         return false;
     }

     stream_remaining_ = stream_->Size();
     output_bytes_ = output_bytes;

     uint8_t chunk[kChunkSize];
     while (stream_remaining_) {
         size_t read = std::min(stream_remaining_, sizeof(chunk));
         if (!stream_->Read(chunk, read, &read)) {
             return false;
         }
         if (!read) {
             LOG(ERROR) << "Stream ended prematurely";
             return false;
         }
         if (!DecompressInput(chunk, read)) {
             return false;
         }

         stream_remaining_ -= read;

         if (stream_remaining_ && Done()) {
             LOG(ERROR) << "Decompressor terminated early";
             return false;
         }
     }
     if (!Done()) {
         LOG(ERROR) << "Decompressor expected more bytes";
         return false;
     }
     return true;
 }

 bool StreamDecompressor::GetFreshBuffer() {
     size_t request_size = std::min(output_bytes_, kChunkSize);
     output_buffer_ =
             reinterpret_cast<uint8_t*>(sink_->GetBuffer(request_size, &output_buffer_remaining_));
     if (!output_buffer_) {
         LOG(ERROR) << "Could not acquire buffer from sink";
         return false;
     }
     return true;
 }

 class GzDecompressor final : public StreamDecompressor {
   public:
     ~GzDecompressor();

     bool Init() override;
     bool DecompressInput(const uint8_t* data, size_t length) override;
     bool Done() override { return ended_; }

   private:
     z_stream z_ = {};
     bool ended_ = false;
 };

 bool GzDecompressor::Init() {
     if (int rv = inflateInit(&z_); rv != Z_OK) {
         LOG(ERROR) << "inflateInit returned error code " << rv;
         return false;
     }
     return true;
 }

 GzDecompressor::~GzDecompressor() {
     inflateEnd(&z_);
 }

 bool GzDecompressor::DecompressInput(const uint8_t* data, size_t length) {
     z_.next_in = reinterpret_cast<Bytef*>(const_cast<uint8_t*>(data));
     z_.avail_in = length;

     while (z_.avail_in) {
         // If no more output buffer, grab a new buffer.
         if (z_.avail_out == 0) {
             if (!GetFreshBuffer()) {
                 return false;
             }
             z_.next_out = reinterpret_cast<Bytef*>(output_buffer_);
             z_.avail_out = output_buffer_remaining_;
         }

         // Remember the position of the output buffer so we can call ReturnData.
         auto avail_out = z_.avail_out;

         // Decompress.
         int rv = inflate(&z_, Z_NO_FLUSH);
         if (rv != Z_OK && rv != Z_STREAM_END) {
             LOG(ERROR) << "inflate returned error code " << rv;
             return false;
         }

         size_t returned = avail_out - z_.avail_out;
         if (!sink_->ReturnData(output_buffer_, returned)) {
             LOG(ERROR) << "Could not return buffer to sink";
             return false;
         }
         output_buffer_ += returned;
         output_buffer_remaining_ -= returned;

         if (rv == Z_STREAM_END) {
             if (z_.avail_in) {
                 LOG(ERROR) << "Gz stream ended prematurely";
                 return false;
             }
             ended_ = true;
             return true;
         }
     }
     return true;
 }

 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());
 }

 class Lz4Decompressor final : public IDecompressor {
   public:
     ~Lz4Decompressor() override = default;

     bool Decompress(const size_t output_size) override {
         size_t actual_buffer_size = 0;
         auto&& output_buffer = sink_->GetBuffer(output_size, &actual_buffer_size);
         if (actual_buffer_size != output_size) {
             LOG(ERROR) << "Failed to allocate buffer of size " << output_size << " only got "
                        << actual_buffer_size << " bytes";
             return false;
         }
         // If input size is same as output size, then input is uncompressed.
         if (stream_->Size() == output_size) {
             size_t bytes_read = 0;
             stream_->Read(output_buffer, output_size, &bytes_read);
             if (bytes_read != output_size) {
                 LOG(ERROR) << "Failed to read all input at once. Expected: " << output_size
                            << " actual: " << bytes_read;
                 return false;
             }
             sink_->ReturnData(output_buffer, output_size);
             return true;
         }
         std::string input_buffer;
         input_buffer.resize(stream_->Size());
         size_t bytes_read = 0;
         stream_->Read(input_buffer.data(), input_buffer.size(), &bytes_read);
         if (bytes_read != input_buffer.size()) {
             LOG(ERROR) << "Failed to read all input at once. Expected: " << input_buffer.size()
                        << " actual: " << bytes_read;
             return false;
         }
         const int bytes_decompressed =
                 LZ4_decompress_safe(input_buffer.data(), static_cast<char*>(output_buffer),
                                     input_buffer.size(), output_size);
         if (bytes_decompressed != output_size) {
             LOG(ERROR) << "Failed to decompress LZ4 block, expected output size: " << output_size
                        << ", actual: " << bytes_decompressed;
             return false;
         }
         sink_->ReturnData(output_buffer, output_size);
         return true;
     }
 };

 std::unique_ptr<IDecompressor> IDecompressor::Lz4() {
     return std::make_unique<Lz4Decompressor>();
 }

 }  // namespace snapshot
 }  // namespace android
	//
	// 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.
	//

	#include "cow_decompress.h"

	#include <utility>

	#include <android-base/logging.h>
	#include <brotli/decode.h>
	#include <lz4.h>
	#include <zlib.h>

	namespace android {
	namespace snapshot {

	class NoDecompressor final : public IDecompressor {
	public:
	bool Decompress(size_t) override;
	};

	bool NoDecompressor::Decompress(size_t) {
	size_t stream_remaining = stream_->Size();
	while (stream_remaining) {
	size_t buffer_size = stream_remaining;
	uint8_t* buffer = reinterpret_cast<uint8_t*>(sink_->GetBuffer(buffer_size, &buffer_size));
	if (!buffer) {
	LOG(ERROR) << "Could not acquire buffer from sink";
	return false;
	}

	// Read until we can fill the buffer.
	uint8_t* buffer_pos = buffer;
	size_t bytes_to_read = std::min(buffer_size, stream_remaining);
	while (bytes_to_read) {
	size_t read;
	if (!stream_->Read(buffer_pos, bytes_to_read, &read)) {
	return false;
	}
	if (!read) {
	LOG(ERROR) << "Stream ended prematurely";
	return false;
	}
	if (!sink_->ReturnData(buffer_pos, read)) {
	LOG(ERROR) << "Could not return buffer to sink";
	return false;
	}
	buffer_pos += read;
	bytes_to_read -= read;
	stream_remaining -= read;
	}
	}
	return true;
	}

	std::unique_ptr<IDecompressor> IDecompressor::Uncompressed() {
	return std::unique_ptr<IDecompressor>(new NoDecompressor());
	}

	// Read chunks of the COW and incrementally stream them to the decoder.
	class StreamDecompressor : public IDecompressor {
	public:
	bool Decompress(size_t output_bytes) override;

	virtual bool Init() = 0;
	virtual bool DecompressInput(const uint8_t* data, size_t length) = 0;
	virtual bool Done() = 0;

	protected:
	bool GetFreshBuffer();

	size_t output_bytes_;
	size_t stream_remaining_;
	uint8_t* output_buffer_ = nullptr;
	size_t output_buffer_remaining_ = 0;
	};

	static constexpr size_t kChunkSize = 4096;

	bool StreamDecompressor::Decompress(size_t output_bytes) {
	if (!Init()) {
	return false;
	}

	stream_remaining_ = stream_->Size();
	output_bytes_ = output_bytes;

	uint8_t chunk[kChunkSize];
	while (stream_remaining_) {
	size_t read = std::min(stream_remaining_, sizeof(chunk));
	if (!stream_->Read(chunk, read, &read)) {
	return false;
	}
	if (!read) {
	LOG(ERROR) << "Stream ended prematurely";
	return false;
	}
	if (!DecompressInput(chunk, read)) {
	return false;
	}

	stream_remaining_ -= read;

	if (stream_remaining_ && Done()) {
	LOG(ERROR) << "Decompressor terminated early";
	return false;
	}
	}
	if (!Done()) {
	LOG(ERROR) << "Decompressor expected more bytes";
	return false;
	}
	return true;
	}

	bool StreamDecompressor::GetFreshBuffer() {
	size_t request_size = std::min(output_bytes_, kChunkSize);
	output_buffer_ =
	reinterpret_cast<uint8_t*>(sink_->GetBuffer(request_size, &output_buffer_remaining_));
	if (!output_buffer_) {
	LOG(ERROR) << "Could not acquire buffer from sink";
	return false;
	}
	return true;
	}

	class GzDecompressor final : public StreamDecompressor {
	public:
	~GzDecompressor();

	bool Init() override;
	bool DecompressInput(const uint8_t* data, size_t length) override;
	bool Done() override { return ended_; }

	private:
	z_stream z_ = {};
	bool ended_ = false;
	};

	bool GzDecompressor::Init() {
	if (int rv = inflateInit(&z_); rv != Z_OK) {
	LOG(ERROR) << "inflateInit returned error code " << rv;
	return false;
	}
	return true;
	}

	GzDecompressor::~GzDecompressor() {
	inflateEnd(&z_);
	}

	bool GzDecompressor::DecompressInput(const uint8_t* data, size_t length) {
	z_.next_in = reinterpret_cast<Bytef>(const_cast<uint8_t>(data));
	z_.avail_in = length;

	while (z_.avail_in) {
	// If no more output buffer, grab a new buffer.
	if (z_.avail_out == 0) {
	if (!GetFreshBuffer()) {
	return false;
	}
	z_.next_out = reinterpret_cast<Bytef*>(output_buffer_);
	z_.avail_out = output_buffer_remaining_;
	}

	// Remember the position of the output buffer so we can call ReturnData.
	auto avail_out = z_.avail_out;

	// Decompress.
	int rv = inflate(&z_, Z_NO_FLUSH);
	if (rv != Z_OK && rv != Z_STREAM_END) {
	LOG(ERROR) << "inflate returned error code " << rv;
	return false;
	}

	size_t returned = avail_out - z_.avail_out;
	if (!sink_->ReturnData(output_buffer_, returned)) {
	LOG(ERROR) << "Could not return buffer to sink";
	return false;
	}
	output_buffer_ += returned;
	output_buffer_remaining_ -= returned;

	if (rv == Z_STREAM_END) {
	if (z_.avail_in) {
	LOG(ERROR) << "Gz stream ended prematurely";
	return false;
	}
	ended_ = true;
	return true;
	}
	}
	return true;
	}

	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());
	}

	class Lz4Decompressor final : public IDecompressor {
	public:
	~Lz4Decompressor() override = default;

	bool Decompress(const size_t output_size) override {
	size_t actual_buffer_size = 0;
	auto&& output_buffer = sink_->GetBuffer(output_size, &actual_buffer_size);
	if (actual_buffer_size != output_size) {
	LOG(ERROR) << "Failed to allocate buffer of size " << output_size << " only got "
	<< actual_buffer_size << " bytes";
	return false;
	}
	// If input size is same as output size, then input is uncompressed.
	if (stream_->Size() == output_size) {
	size_t bytes_read = 0;
	stream_->Read(output_buffer, output_size, &bytes_read);
	if (bytes_read != output_size) {
	LOG(ERROR) << "Failed to read all input at once. Expected: " << output_size
	<< " actual: " << bytes_read;
	return false;
	}
	sink_->ReturnData(output_buffer, output_size);
	return true;
	}
	std::string input_buffer;
	input_buffer.resize(stream_->Size());
	size_t bytes_read = 0;
	stream_->Read(input_buffer.data(), input_buffer.size(), &bytes_read);
	if (bytes_read != input_buffer.size()) {
	LOG(ERROR) << "Failed to read all input at once. Expected: " << input_buffer.size()
	<< " actual: " << bytes_read;
	return false;
	}
	const int bytes_decompressed =
	LZ4_decompress_safe(input_buffer.data(), static_cast<char*>(output_buffer),
	input_buffer.size(), output_size);
	if (bytes_decompressed != output_size) {
	LOG(ERROR) << "Failed to decompress LZ4 block, expected output size: " << output_size
	<< ", actual: " << bytes_decompressed;
	return false;
	}
	sink_->ReturnData(output_buffer, output_size);
	return true;
	}
	};

	std::unique_ptr<IDecompressor> IDecompressor::Lz4() {
	return std::make_unique<Lz4Decompressor>();
	}

	} // namespace snapshot
	} // namespace android