Donate to e Foundation | Murena handsets with /e/OS | Own a part of Murena! Learn more

Commit 918e6ea1 authored by Tianjie Xu's avatar Tianjie Xu Committed by Gerrit Code Review
Browse files

Merge "Add support to decompress brotli compressed new data"

parents b87a166a 107a34f9

tests/Android.mk

+1 −0
Original line number Diff line number Diff line
@@ -160,6 +160,7 @@ LOCAL_STATIC_LIBRARIES := \ 
    libfec_rs \

    libsquashfs_utils \

    libcutils \

    libbrotli \

    $(tune2fs_static_libraries)



testdata_files := $(call find-subdir-files, testdata/*)

tests/component/updater_test.cpp

+67 −0
Original line number Diff line number Diff line
@@ -15,10 +15,12 @@ 
 */



#include <stdio.h>

#include <stdlib.h>

#include <sys/stat.h>

#include <sys/types.h>

#include <unistd.h>



#include <algorithm>

#include <memory>

#include <string>

#include <vector>
@@ -29,6 +31,7 @@ 
#include <android-base/strings.h>

#include <android-base/test_utils.h>

#include <bootloader_message/bootloader_message.h>

#include <brotli/encode.h>

#include <bsdiff.h>

#include <gtest/gtest.h>

#include <ziparchive/zip_archive.h>
@@ -576,4 +579,68 @@ TEST_F(UpdaterTest, new_data_short_write) { 
  std::string script_exact_data = "block_image_update(\"" + std::string(update_file.path) +

      R"(", package_extract_file("transfer_list"), "exact_new_data", "patch_data"))";

  expect("t", script_exact_data.c_str(), kNoCause, &updater_info);

  CloseArchive(handle);

}



TEST_F(UpdaterTest, brotli_new_data) {

  // Create a zip file with new_data.

  TemporaryFile zip_file;

  FILE* zip_file_ptr = fdopen(zip_file.fd, "wb");

  ZipWriter zip_writer(zip_file_ptr);



  // Add a brotli compressed new data entry.

  ASSERT_EQ(0, zip_writer.StartEntry("new.dat.br", 0));



  auto generator = []() { return rand() % 128; };

  // Generate 2048 blocks of random data.

  std::string brotli_new_data;

  brotli_new_data.reserve(4096 * 2048);

  generate_n(back_inserter(brotli_new_data), 4096 * 2048, generator);



  size_t encoded_size = BrotliEncoderMaxCompressedSize(brotli_new_data.size());

  std::vector<uint8_t> encoded_data(encoded_size);

  ASSERT_TRUE(BrotliEncoderCompress(

      BROTLI_DEFAULT_QUALITY, BROTLI_DEFAULT_WINDOW, BROTLI_DEFAULT_MODE, brotli_new_data.size(),

      reinterpret_cast<const uint8_t*>(brotli_new_data.data()), &encoded_size, encoded_data.data()));



  ASSERT_EQ(0, zip_writer.WriteBytes(encoded_data.data(), encoded_size));

  ASSERT_EQ(0, zip_writer.FinishEntry());

  // Add a dummy patch data.

  ASSERT_EQ(0, zip_writer.StartEntry("patch_data", 0));

  ASSERT_EQ(0, zip_writer.FinishEntry());



  std::vector<std::string> transfer_list = {

    "4", "2048", "0", "0", "new 4,0,512,512,1024", "new 2,1024,2048",

  };

  ASSERT_EQ(0, zip_writer.StartEntry("transfer_list", 0));

  std::string commands = android::base::Join(transfer_list, '\n');

  ASSERT_EQ(0, zip_writer.WriteBytes(commands.data(), commands.size()));

  ASSERT_EQ(0, zip_writer.FinishEntry());

  ASSERT_EQ(0, zip_writer.Finish());

  ASSERT_EQ(0, fclose(zip_file_ptr));



  MemMapping map;

  ASSERT_TRUE(map.MapFile(zip_file.path));

  ZipArchiveHandle handle;

  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));



  // Set up the handler, command_pipe, patch offset & length.

  UpdaterInfo updater_info;

  updater_info.package_zip = handle;

  TemporaryFile temp_pipe;

  updater_info.cmd_pipe = fopen(temp_pipe.path, "wb");

  updater_info.package_zip_addr = map.addr;

  updater_info.package_zip_len = map.length;



  // Check if we can decompress the new data correctly.

  TemporaryFile update_file;

  std::string script_new_data =

      "block_image_update(\"" + std::string(update_file.path) +

      R"(", package_extract_file("transfer_list"), "new.dat.br", "patch_data"))";

  expect("t", script_new_data.c_str(), kNoCause, &updater_info);



  std::string updated_content;

  ASSERT_TRUE(android::base::ReadFileToString(update_file.path, &updated_content));

  ASSERT_EQ(brotli_new_data, updated_content);

  CloseArchive(handle);

}

updater/Android.mk

+1 −0
Original line number Diff line number Diff line
@@ -47,6 +47,7 @@ updater_common_static_libraries := \ 
    libcrypto_utils \

    libcutils \

    libtune2fs \

    libbrotli \

    $(tune2fs_static_libraries)



# libupdater (static library)

updater/blockimg.cpp

+143 −32
Original line number Diff line number Diff line
@@ -44,6 +44,7 @@ 
#include <android-base/strings.h>

#include <android-base/unique_fd.h>

#include <applypatch/applypatch.h>

#include <brotli/decode.h>

#include <openssl/sha.h>

#include <private/android_filesystem_config.h>

#include <ziparchive/zip_archive.h>
@@ -149,42 +150,34 @@ static void allocate(size_t size, std::vector<uint8_t>& buffer) { 
class RangeSinkWriter {

 public:

  RangeSinkWriter(int fd, const RangeSet& tgt)

      : fd_(fd), tgt_(tgt), next_range_(0), current_range_left_(0), bytes_written_(0) {

      : fd_(fd),

        tgt_(tgt),

        next_range_(0),

        current_range_left_(0),

        bytes_written_(0) {

    CHECK_NE(tgt.size(), static_cast<size_t>(0));

  };



  virtual ~RangeSinkWriter() {};



  bool Finished() const {

    return next_range_ == tgt_.size() && current_range_left_ == 0;

  }



  size_t Write(const uint8_t* data, size_t size) {

  // Return number of bytes consumed; and 0 indicates a writing failure.

  virtual size_t Write(const uint8_t* data, size_t size) {

    if (Finished()) {

      LOG(ERROR) << "range sink write overrun; can't write " << size << " bytes";

      return 0;

    }



    size_t written = 0;

    size_t consumed = 0;

    while (size > 0) {

      // Move to the next range as needed.

      if (current_range_left_ == 0) {

        if (next_range_ < tgt_.size()) {

          const Range& range = tgt_[next_range_];

          off64_t offset = static_cast<off64_t>(range.first) * BLOCKSIZE;

          current_range_left_ = (range.second - range.first) * BLOCKSIZE;

          next_range_++;

          if (!discard_blocks(fd_, offset, current_range_left_)) {

      if (!SeekToOutputRange()) {

        break;

      }



          if (!check_lseek(fd_, offset, SEEK_SET)) {

            break;

          }

        } else {

          // We can't write any more; return how many bytes have been written so far.

          break;

        }

      }



      size_t write_now = size;

      if (current_range_left_ < write_now) {

        write_now = current_range_left_;
@@ -198,21 +191,47 @@ class RangeSinkWriter { 
      size -= write_now;



      current_range_left_ -= write_now;

      written += write_now;

      consumed += write_now;

    }



    bytes_written_ += written;

    return written;

    bytes_written_ += consumed;

    return consumed;

  }



  size_t BytesWritten() const {

    return bytes_written_;

  }



 private:

  // The input data.

 protected:

  // Set up the output cursor, move to next range if needed.

  bool SeekToOutputRange() {

    // We haven't finished the current range yet.

    if (current_range_left_ != 0) {

      return true;

    }

    // We can't write any more; let the write function return how many bytes have been written

    // so far.

    if (next_range_ >= tgt_.size()) {

      return false;

    }



    const Range& range = tgt_[next_range_];

    off64_t offset = static_cast<off64_t>(range.first) * BLOCKSIZE;

    current_range_left_ = (range.second - range.first) * BLOCKSIZE;

    next_range_++;



    if (!discard_blocks(fd_, offset, current_range_left_)) {

      return false;

    }

    if (!check_lseek(fd_, offset, SEEK_SET)) {

      return false;

    }

    return true;

  }



  // The output file descriptor.

  int fd_;

  // The destination for the data.

  // The destination ranges for the data.

  const RangeSet& tgt_;

  // The next range that we should write to.

  size_t next_range_;
@@ -222,6 +241,75 @@ class RangeSinkWriter { 
  size_t bytes_written_;

};



class BrotliNewDataWriter : public RangeSinkWriter {

 public:

  BrotliNewDataWriter(int fd, const RangeSet& tgt, BrotliDecoderState* state)

      : RangeSinkWriter(fd, tgt), state_(state) {}



  size_t Write(const uint8_t* data, size_t size) override {

    if (Finished()) {

      LOG(ERROR) << "Brotli new data write overrun; can't write " << size << " bytes";

      return 0;

    }

    CHECK(state_ != nullptr);



    size_t consumed = 0;

    while (true) {

      // Move to the next range as needed.

      if (!SeekToOutputRange()) {

        break;

      }



      size_t available_in = size;

      size_t write_now = std::min<size_t>(32768, current_range_left_);

      uint8_t buffer[write_now];



      size_t available_out = write_now;

      uint8_t* next_out = buffer;



      // The brotli decoder will update |data|, |available_in|, |next_out| and |available_out|.

      BrotliDecoderResult result = BrotliDecoderDecompressStream(

          state_, &available_in, &data, &available_out, &next_out, nullptr);



      // We don't have a way to recover from the decode error; report the failure.

      if (result == BROTLI_DECODER_RESULT_ERROR) {

        LOG(ERROR) << "Decompression failed with "

                   << BrotliDecoderErrorString(BrotliDecoderGetErrorCode(state_));

        return 0;

      }



      if (write_all(fd_, buffer, write_now - available_out) == -1) {

        return 0;

      }



      LOG(DEBUG) << "bytes written: " << write_now - available_out << ", bytes consumed "

                 << size - available_in << ", decoder status " << result;



      // Update the total bytes written to output by the current writer; this is different from the

      // consumed input bytes.

      bytes_written_ += write_now - available_out;

      current_range_left_ -= (write_now - available_out);

      consumed += (size - available_in);



      // Update the remaining size. The input data ptr is already updated by brotli decoder

      // function.

      size = available_in;



      // Continue if we have more output to write, or more input to consume.

      if (result == BROTLI_DECODER_RESULT_SUCCESS ||

          (result == BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT && size == 0)) {

        break;

      }

    }



    return consumed;

  }



 private:

  // Pointer to the decoder state. (initialized by PerformBlockImageUpdate)

  BrotliDecoderState* state_;

};



/**

 * All of the data for all the 'new' transfers is contained in one file in the update package,

 * concatenated together in the order in which transfers.list will need it. We want to stream it out
@@ -243,8 +331,10 @@ class RangeSinkWriter { 
struct NewThreadInfo {

  ZipArchiveHandle za;

  ZipEntry entry;

  bool brotli_compressed;



  RangeSinkWriter* writer;

  std::unique_ptr<RangeSinkWriter> writer;

  BrotliDecoderState* brotli_decoder_state;

  bool receiver_available;



  pthread_mutex_t mu;
@@ -264,9 +354,16 @@ static bool receive_new_data(const uint8_t* data, size_t size, void* cookie) { 


    // At this point nti->writer is set, and we own it. The main thread is waiting for it to

    // disappear from nti.

    size_t written = nti->writer->Write(data, size);

    data += written;

    size -= written;

    size_t consumed = nti->writer->Write(data, size);



    // We encounter a fatal error if we fail to consume any input bytes. If this happens, abort the

    // extraction.

    if (consumed == 0) {

      LOG(ERROR) << "Failed to process " << size << " input bytes.";

      return false;

    }

    data += consumed;

    size -= consumed;



    if (nti->writer->Finished()) {

      // We have written all the bytes desired by this writer.
@@ -1142,9 +1239,13 @@ static int PerformCommandNew(CommandParameters& params) { 
  if (params.canwrite) {

    LOG(INFO) << " writing " << tgt.blocks() << " blocks of new data";



    RangeSinkWriter writer(params.fd, tgt);

    pthread_mutex_lock(&params.nti.mu);

    params.nti.writer = &writer;

    if (params.nti.brotli_compressed) {

      params.nti.writer =

          std::make_unique<BrotliNewDataWriter>(params.fd, tgt, params.nti.brotli_decoder_state);

    } else {

      params.nti.writer = std::make_unique<RangeSinkWriter>(params.fd, tgt);

    }

    pthread_cond_broadcast(&params.nti.cv);



    while (params.nti.writer != nullptr) {
@@ -1384,6 +1485,12 @@ static Value* PerformBlockImageUpdate(const char* name, State* state, 
  if (params.canwrite) {

    params.nti.za = za;

    params.nti.entry = new_entry;

    // The entry is compressed by brotli if has a 'br' extension.

    params.nti.brotli_compressed = android::base::EndsWith(new_data_fn->data, ".br");

    if (params.nti.brotli_compressed) {

      // Initialize brotli decoder state.

      params.nti.brotli_decoder_state = BrotliDecoderCreateInstance(nullptr, nullptr, nullptr);

    }

    params.nti.receiver_available = true;



    pthread_mutex_init(&params.nti.mu, nullptr);
@@ -1526,6 +1633,10 @@ pbiudone: 
  }

  // params.fd will be automatically closed because it's a unique_fd.



  if (params.nti.brotli_decoder_state != nullptr) {

    BrotliDecoderDestroyInstance(params.nti.brotli_decoder_state);

  }



  // Only delete the stash if the update cannot be resumed, or it's a verification run and we

  // created the stash.

  if (params.isunresumable || (!params.canwrite && params.createdstash)) {