2 files changed, 75 insertions, 3 deletions

diff --git a/applypatch/imgpatch.cpp b/applypatch/imgpatch.cpp
index 9794a4878..b06a64f21 100644
--- a/applypatch/imgpatch.cpp
+++ b/applypatch/imgpatch.cpp
@@ -38,6 +38,7 @@
 #include <zlib.h>
 
 #include "edify/expr.h"
+#include "otautil/print_sha1.h"
 
 static inline int64_t Read8(const void *address) {
   return android::base::get_unaligned<int64_t>(address);
@@ -76,8 +77,10 @@ static bool ApplyBSDiffPatchAndStreamOutput(const uint8_t* src_data, size_t src_
   size_t actual_target_length = 0;
   size_t total_written = 0;
   static constexpr size_t buffer_size = 32768;
+  SHA_CTX sha_ctx;
+  SHA1_Init(&sha_ctx);
   auto compression_sink = [&strm, &actual_target_length, &expected_target_length, &total_written,
-                           &ret, &sink](const uint8_t* data, size_t len) -> size_t {
+                           &ret, &sink, &sha_ctx](const uint8_t* data, size_t len) -> size_t {
     // The input patch length for an update never exceeds INT_MAX.
     strm.avail_in = len;
     strm.next_in = data;
@@ -98,6 +101,20 @@ static bool ApplyBSDiffPatchAndStreamOutput(const uint8_t* src_data, size_t src_
 
       size_t have = buffer_size - strm.avail_out;
       total_written += have;
+
+      // TODO(b/67849209) Remove after debugging the unit test flakiness.
+      if (android::base::GetMinimumLogSeverity() <= android::base::LogSeverity::DEBUG &&
+          have != 0) {
+        SHA1_Update(&sha_ctx, data, len - strm.avail_in);
+        SHA_CTX temp_ctx;
+        memcpy(&temp_ctx, &sha_ctx, sizeof(SHA_CTX));
+        uint8_t digest_so_far[SHA_DIGEST_LENGTH];
+        SHA1_Final(digest_so_far, &temp_ctx);
+        LOG(DEBUG) << "Processed " << actual_target_length + len - strm.avail_in
+                   << " bytes input data in the sink function, sha1 so far: "
+                   << short_sha1(digest_so_far);
+      }
+
       if (sink(buffer.data(), have) != have) {
         LOG(ERROR) << "Failed to write " << have << " compressed bytes to output.";
         return 0;
@@ -111,6 +128,11 @@ static bool ApplyBSDiffPatchAndStreamOutput(const uint8_t* src_data, size_t src_
   int bspatch_result = ApplyBSDiffPatch(src_data, src_len, patch, patch_offset, compression_sink);
   deflateEnd(&strm);
 
+  if (android::base::GetMinimumLogSeverity() <= android::base::LogSeverity::DEBUG) {
+    uint8_t digest[SHA_DIGEST_LENGTH];
+    SHA1_Final(digest, &sha_ctx);
+    LOG(DEBUG) << "sha1 of " << actual_target_length << " bytes input data: " << short_sha1(digest);
+  }
   if (bspatch_result != 0) {
     return false;
   }
diff --git a/tests/component/applypatch_test.cpp b/tests/component/applypatch_test.cpp
index 292d76e43..04055b9fe 100644
--- a/tests/component/applypatch_test.cpp
+++ b/tests/component/applypatch_test.cpp
@@ -37,6 +37,7 @@
 #include <android-base/unique_fd.h>
 #include <bsdiff/bsdiff.h>
 #include <openssl/sha.h>
+#include <zlib.h>
 
 #include "applypatch/applypatch.h"
 #include "applypatch/applypatch_modes.h"
@@ -46,6 +47,47 @@
 
 using namespace std::string_literals;
 
+// TODO(b/67849209) Remove after debug the flakiness.
+static void DecompressAndDumpRecoveryImage(const std::string& image_path) {
+  // Expected recovery_image structure
+  // chunk normal:  45066 bytes
+  // chunk deflate: 479442 bytes
+  // chunk normal:  5199 bytes
+  std::string recovery_content;
+  ASSERT_TRUE(android::base::ReadFileToString(image_path, &recovery_content));
+  ASSERT_GT(recovery_content.size(), 45066 + 5199);
+
+  z_stream strm = {};
+  strm.avail_in = recovery_content.size() - 45066 - 5199;
+  strm.next_in =
+      const_cast<uint8_t*>(reinterpret_cast<const uint8_t*>(recovery_content.data())) + 45066;
+
+  ASSERT_EQ(Z_OK, inflateInit2(&strm, -15));
+
+  constexpr unsigned int BUFFER_SIZE = 32768;
+  std::vector<uint8_t> uncompressed_data(BUFFER_SIZE);
+  size_t uncompressed_length = 0;
+  SHA_CTX ctx;
+  SHA1_Init(&ctx);
+  int ret;
+  do {
+    strm.avail_out = BUFFER_SIZE;
+    strm.next_out = uncompressed_data.data();
+
+    ret = inflate(&strm, Z_NO_FLUSH);
+    ASSERT_GE(ret, 0);
+
+    SHA1_Update(&ctx, uncompressed_data.data(), BUFFER_SIZE - strm.avail_out);
+    uncompressed_length += BUFFER_SIZE - strm.avail_out;
+  } while (ret != Z_STREAM_END);
+  inflateEnd(&strm);
+
+  uint8_t digest[SHA_DIGEST_LENGTH];
+  SHA1_Final(digest, &ctx);
+  GTEST_LOG_(INFO) << "uncompressed length " << uncompressed_length
+                   << " sha1: " << short_sha1(digest);
+}
+
 static void sha1sum(const std::string& fname, std::string* sha1, size_t* fsize = nullptr) {
   ASSERT_TRUE(sha1 != nullptr);
 
@@ -317,7 +359,11 @@ TEST_F(ApplyPatchModesTest, PatchModeEmmcTargetWithoutBonusFile) {
                                     recovery_img_sha1.c_str(),
                                     recovery_img_size_arg.c_str(),
                                     patch_arg.c_str() };
-  ASSERT_EQ(0, applypatch_modes(args.size(), args.data()));
+
+  if (applypatch_modes(args.size(), args.data()) != 0) {
+    DecompressAndDumpRecoveryImage(tgt_file.path);
+    FAIL();
+  }
 }
 
 TEST_F(ApplyPatchModesTest, PatchModeEmmcTargetWithMultiplePatches) {
@@ -360,7 +406,11 @@ TEST_F(ApplyPatchModesTest, PatchModeEmmcTargetWithMultiplePatches) {
   for (const auto& arg : args) {
     printf("  %s\n", arg);
   }
-  ASSERT_EQ(0, applypatch_modes(args.size(), args.data()));
+
+  if (applypatch_modes(args.size(), args.data()) != 0) {
+    DecompressAndDumpRecoveryImage(tgt_file.path);
+    FAIL();
+  }
 }
 
 // Ensures that applypatch works with a bsdiff based recovery-from-boot.p.