4 files changed, 151 insertions, 127 deletions

diff --git a/tests/component/update_verifier_test.cpp b/tests/component/update_verifier_test.cpp
index 5fc7ef63f..b04e1185e 100644
--- a/tests/component/update_verifier_test.cpp
+++ b/tests/component/update_verifier_test.cpp
@@ -81,3 +81,16 @@ TEST_F(UpdateVerifierTest, verify_image_malformed_care_map) {
   ASSERT_TRUE(android::base::WriteStringToFile(content, temp_file.path));
   ASSERT_FALSE(verify_image(temp_file.path));
 }
+
+TEST_F(UpdateVerifierTest, verify_image_legacy_care_map) {
+  // This test relies on dm-verity support.
+  if (!verity_supported) {
+    GTEST_LOG_(INFO) << "Test skipped on devices without dm-verity support.";
+    return;
+  }
+
+  TemporaryFile temp_file;
+  std::string content = "/dev/block/bootdevice/by-name/system\n2,1,0";
+  ASSERT_TRUE(android::base::WriteStringToFile(content, temp_file.path));
+  ASSERT_TRUE(verify_image(temp_file.path));
+}
diff --git a/tests/component/updater_test.cpp b/tests/component/updater_test.cpp
index 01b86f224..6c341c111 100644
--- a/tests/component/updater_test.cpp
+++ b/tests/component/updater_test.cpp
@@ -592,10 +592,10 @@ TEST_F(UpdaterTest, brotli_new_data) {
   ASSERT_EQ(0, zip_writer.StartEntry("new.dat.br", 0));
 
   auto generator = []() { return rand() % 128; };
-  // Generate 2048 blocks of random data.
+  // Generate 100 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);
+  brotli_new_data.reserve(4096 * 100);
+  generate_n(back_inserter(brotli_new_data), 4096 * 100, generator);
 
   size_t encoded_size = BrotliEncoderMaxCompressedSize(brotli_new_data.size());
   std::vector<uint8_t> encoded_data(encoded_size);
@@ -609,8 +609,19 @@ TEST_F(UpdaterTest, brotli_new_data) {
   ASSERT_EQ(0, zip_writer.StartEntry("patch_data", 0));
   ASSERT_EQ(0, zip_writer.FinishEntry());
 
+  // Write a few small chunks of new data, then a large chunk, and finally a few small chunks.
+  // This helps us to catch potential short writes.
   std::vector<std::string> transfer_list = {
-    "4", "2048", "0", "0", "new 4,0,512,512,1024", "new 2,1024,2048",
+    "4",
+    "100",
+    "0",
+    "0",
+    "new 2,0,1",
+    "new 2,1,2",
+    "new 4,2,50,50,97",
+    "new 2,97,98",
+    "new 2,98,99",
+    "new 2,99,100",
   };
   ASSERT_EQ(0, zip_writer.StartEntry("transfer_list", 0));
   std::string commands = android::base::Join(transfer_list, '\n');
diff --git a/update_verifier/update_verifier.cpp b/update_verifier/update_verifier.cpp
index d3a5185b8..b49011a12 100644
--- a/update_verifier/update_verifier.cpp
+++ b/update_verifier/update_verifier.cpp
@@ -175,40 +175,53 @@ static bool read_blocks(const std::string& partition, const std::string& range_s
   return true;
 }
 
+// Returns true to indicate a passing verification (or the error should be ignored); Otherwise
+// returns false on fatal errors, where we should reject the current boot and trigger a fallback.
+// Note that update_verifier should be backward compatible to not reject care_map.txt from old
+// releases, which could otherwise fail to boot into the new release. For example, we've changed
+// the care_map format between N and O. An O update_verifier would fail to work with N
+// care_map.txt. This could be a result of sideloading an O OTA while the device having a pending N
+// update.
 bool verify_image(const std::string& care_map_name) {
-    android::base::unique_fd care_map_fd(TEMP_FAILURE_RETRY(open(care_map_name.c_str(), O_RDONLY)));
-    // If the device is flashed before the current boot, it may not have care_map.txt
-    // in /data/ota_package. To allow the device to continue booting in this situation,
-    // we should print a warning and skip the block verification.
-    if (care_map_fd.get() == -1) {
-        PLOG(WARNING) << "Failed to open " << care_map_name;
-        return true;
-    }
-    // Care map file has four lines (two lines if vendor partition is not present):
-    // First line has the block partition name (system/vendor).
-    // Second line holds all ranges of blocks to verify.
-    // The next two lines have the same format but for vendor partition.
-    std::string file_content;
-    if (!android::base::ReadFdToString(care_map_fd.get(), &file_content)) {
-        LOG(ERROR) << "Error reading care map contents to string.";
-        return false;
-    }
+  android::base::unique_fd care_map_fd(TEMP_FAILURE_RETRY(open(care_map_name.c_str(), O_RDONLY)));
+  // If the device is flashed before the current boot, it may not have care_map.txt
+  // in /data/ota_package. To allow the device to continue booting in this situation,
+  // we should print a warning and skip the block verification.
+  if (care_map_fd.get() == -1) {
+    PLOG(WARNING) << "Failed to open " << care_map_name;
+    return true;
+  }
+  // Care map file has four lines (two lines if vendor partition is not present):
+  // First line has the block partition name (system/vendor).
+  // Second line holds all ranges of blocks to verify.
+  // The next two lines have the same format but for vendor partition.
+  std::string file_content;
+  if (!android::base::ReadFdToString(care_map_fd.get(), &file_content)) {
+    LOG(ERROR) << "Error reading care map contents to string.";
+    return false;
+  }
 
-    std::vector<std::string> lines;
-    lines = android::base::Split(android::base::Trim(file_content), "\n");
-    if (lines.size() != 2 && lines.size() != 4) {
-        LOG(ERROR) << "Invalid lines in care_map: found " << lines.size()
-                   << " lines, expecting 2 or 4 lines.";
-        return false;
-    }
+  std::vector<std::string> lines;
+  lines = android::base::Split(android::base::Trim(file_content), "\n");
+  if (lines.size() != 2 && lines.size() != 4) {
+    LOG(ERROR) << "Invalid lines in care_map: found " << lines.size()
+               << " lines, expecting 2 or 4 lines.";
+    return false;
+  }
 
-    for (size_t i = 0; i < lines.size(); i += 2) {
-        if (!read_blocks(lines[i], lines[i+1])) {
-            return false;
-        }
+  for (size_t i = 0; i < lines.size(); i += 2) {
+    // We're seeing an N care_map.txt. Skip the verification since it's not compatible with O
+    // update_verifier (the last few metadata blocks can't be read via device mapper).
+    if (android::base::StartsWith(lines[i], "/dev/block/")) {
+      LOG(WARNING) << "Found legacy care_map.txt; skipped.";
+      return true;
+    }
+    if (!read_blocks(lines[i], lines[i+1])) {
+      return false;
     }
+  }
 
-    return true;
+  return true;
 }
 
 static int reboot_device() {
diff --git a/updater/blockimg.cpp b/updater/blockimg.cpp
index 2bec487fe..a0b9ad233 100644
--- a/updater/blockimg.cpp
+++ b/updater/blockimg.cpp
@@ -158,20 +158,22 @@ class RangeSinkWriter {
     CHECK_NE(tgt.size(), static_cast<size_t>(0));
   };
 
-  virtual ~RangeSinkWriter() {};
-
   bool Finished() const {
     return next_range_ == tgt_.size() && current_range_left_ == 0;
   }
 
-  // Return number of bytes consumed; and 0 indicates a writing failure.
-  virtual size_t Write(const uint8_t* data, size_t size) {
+  size_t AvailableSpace() const {
+    return tgt_.blocks() * BLOCKSIZE - bytes_written_;
+  }
+
+  // Return number of bytes written; and 0 indicates a writing failure.
+  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 consumed = 0;
+    size_t written = 0;
     while (size > 0) {
       // Move to the next range as needed.
       if (!SeekToOutputRange()) {
@@ -191,18 +193,18 @@ class RangeSinkWriter {
       size -= write_now;
 
       current_range_left_ -= write_now;
-      consumed += write_now;
+      written += write_now;
     }
 
-    bytes_written_ += consumed;
-    return consumed;
+    bytes_written_ += written;
+    return written;
   }
 
   size_t BytesWritten() const {
     return bytes_written_;
   }
 
- protected:
+ private:
   // Set up the output cursor, move to next range if needed.
   bool SeekToOutputRange() {
     // We haven't finished the current range yet.
@@ -241,75 +243,6 @@ 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
@@ -354,16 +287,73 @@ 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 consumed = nti->writer->Write(data, size);
+    size_t write_now = std::min(size, nti->writer->AvailableSpace());
+    if (nti->writer->Write(data, write_now) != write_now) {
+      LOG(ERROR) << "Failed to write " << write_now << " bytes.";
+      return false;
+    }
+
+    data += write_now;
+    size -= write_now;
+
+    if (nti->writer->Finished()) {
+      // We have written all the bytes desired by this writer.
+
+      pthread_mutex_lock(&nti->mu);
+      nti->writer = nullptr;
+      pthread_cond_broadcast(&nti->cv);
+      pthread_mutex_unlock(&nti->mu);
+    }
+  }
+
+  return true;
+}
+
+static bool receive_brotli_new_data(const uint8_t* data, size_t size, void* cookie) {
+  NewThreadInfo* nti = static_cast<NewThreadInfo*>(cookie);
+
+  while (size > 0 || BrotliDecoderHasMoreOutput(nti->brotli_decoder_state)) {
+    // Wait for nti->writer to be non-null, indicating some of this data is wanted.
+    pthread_mutex_lock(&nti->mu);
+    while (nti->writer == nullptr) {
+      pthread_cond_wait(&nti->cv, &nti->mu);
+    }
+    pthread_mutex_unlock(&nti->mu);
+
+    // At this point nti->writer is set, and we own it. The main thread is waiting for it to
+    // disappear from nti.
+
+    size_t buffer_size = std::min<size_t>(32768, nti->writer->AvailableSpace());
+    if (buffer_size == 0) {
+      LOG(ERROR) << "No space left in output range";
+      return false;
+    }
+    uint8_t buffer[buffer_size];
+    size_t available_in = size;
+    size_t available_out = buffer_size;
+    uint8_t* next_out = buffer;
+
+    // The brotli decoder will update |data|, |available_in|, |next_out| and |available_out|.
+    BrotliDecoderResult result = BrotliDecoderDecompressStream(
+        nti->brotli_decoder_state, &available_in, &data, &available_out, &next_out, nullptr);
 
-    // 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.";
+    if (result == BROTLI_DECODER_RESULT_ERROR) {
+      LOG(ERROR) << "Decompression failed with "
+                 << BrotliDecoderErrorString(BrotliDecoderGetErrorCode(nti->brotli_decoder_state));
       return false;
     }
-    data += consumed;
-    size -= consumed;
+
+    LOG(DEBUG) << "bytes to write: " << buffer_size - available_out << ", bytes consumed "
+               << size - available_in << ", decoder status " << result;
+
+    size_t write_now = buffer_size - available_out;
+    if (nti->writer->Write(buffer, write_now) != write_now) {
+      LOG(ERROR) << "Failed to write " << write_now << " bytes.";
+      return false;
+    }
+
+    // Update the remaining size. The input data ptr is already updated by brotli decoder function.
+    size = available_in;
 
     if (nti->writer->Finished()) {
       // We have written all the bytes desired by this writer.
@@ -380,8 +370,11 @@ static bool receive_new_data(const uint8_t* data, size_t size, void* cookie) {
 
 static void* unzip_new_data(void* cookie) {
   NewThreadInfo* nti = static_cast<NewThreadInfo*>(cookie);
-  ProcessZipEntryContents(nti->za, &nti->entry, receive_new_data, nti);
-
+  if (nti->brotli_compressed) {
+    ProcessZipEntryContents(nti->za, &nti->entry, receive_brotli_new_data, nti);
+  } else {
+    ProcessZipEntryContents(nti->za, &nti->entry, receive_new_data, nti);
+  }
   pthread_mutex_lock(&nti->mu);
   nti->receiver_available = false;
   if (nti->writer != nullptr) {
@@ -1240,12 +1233,7 @@ static int PerformCommandNew(CommandParameters& params) {
     LOG(INFO) << " writing " << tgt.blocks() << " blocks of new data";
 
     pthread_mutex_lock(&params.nti.mu);
-    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);
-    }
+    params.nti.writer = std::make_unique<RangeSinkWriter>(params.fd, tgt);
     pthread_cond_broadcast(&params.nti.cv);
 
     while (params.nti.writer != nullptr) {
@@ -1485,7 +1473,6 @@ 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.