Fix puffin failure on duplicating deflates

Test: th
Bug: 322549011

Change-Id: I905ffbf5b3ff10f609be52a6acd8c59f16de5608

4 files changed, 43 insertions, 5 deletions

diff --git a/payload_generator/deflate_utils.cc b/payload_generator/deflate_utils.cc
index cc2e4d64..d1967991 100644
--- a/payload_generator/deflate_utils.cc
+++ b/payload_generator/deflate_utils.cc
@@ -27,6 +27,7 @@
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_generator/delta_diff_generator.h"
 #include "update_engine/payload_generator/extent_ranges.h"
+#include "update_engine/payload_generator/extent_utils.h"
 #include "update_engine/payload_generator/squashfs_filesystem.h"
 #include "update_engine/update_metadata.pb.h"
 
@@ -39,6 +40,12 @@ namespace chromeos_update_engine {
 namespace deflate_utils {
 namespace {
 
+constexpr std::ostream& operator<<(std::ostream& out,
+                                   const puffin::BitExtent& ext) {
+  out << "BitExtent(" << ext.offset << "," << ext.length << ")";
+  return out;
+}
+
 // The minimum size for a squashfs image to be processed.
 const uint64_t kMinimumSquashfsImageSize = 1 * 1024 * 1024;  // bytes
 
@@ -254,7 +261,8 @@ bool CompactDeflates(const vector<Extent>& extents,
 
   // All given |in_deflates| items should've been inside one of the extents in
   // |extents|.
-  TEST_AND_RETURN_FALSE(in_deflates.size() == out_deflates->size());
+  TEST_EQ(in_deflates.size(), out_deflates->size());
+  Dedup(out_deflates);
 
   // Make sure all outgoing deflates are ordered and non-overlapping.
   auto result = std::adjacent_find(out_deflates->begin(),
@@ -262,7 +270,11 @@ bool CompactDeflates(const vector<Extent>& extents,
                                    [](const BitExtent& a, const BitExtent& b) {
                                      return (a.offset + a.length) > b.offset;
                                    });
-  TEST_AND_RETURN_FALSE(result == out_deflates->end());
+  if (result != out_deflates->end()) {
+    LOG(ERROR) << "out_deflate is overlapped " << (*result) << ", "
+               << *(++result);
+    return false;
+  }
   return true;
 }
 
diff --git a/payload_generator/delta_diff_utils.h b/payload_generator/delta_diff_utils.h
index bfbcdf79..c40e267d 100644
--- a/payload_generator/delta_diff_utils.h
+++ b/payload_generator/delta_diff_utils.h
@@ -185,12 +185,14 @@ class BestDiffGenerator {
     // Find all deflate positions inside the given extents and then put all
     // deflates together because we have already read all the extents into
     // one buffer.
+    Dedup(&old_deflates_);
+    Dedup(&new_deflates_);
     vector<puffin::BitExtent> src_deflates;
-    TEST_AND_RETURN(deflate_utils::FindAndCompactDeflates(
+    CHECK(deflate_utils::FindAndCompactDeflates(
         src_extents_, old_deflates_, &src_deflates));
 
     vector<puffin::BitExtent> dst_deflates;
-    TEST_AND_RETURN(deflate_utils::FindAndCompactDeflates(
+    CHECK(deflate_utils::FindAndCompactDeflates(
         dst_extents_, new_deflates_, &dst_deflates));
     puffin::RemoveEqualBitExtents(
         old_data_, new_data_, &src_deflates, &dst_deflates);
diff --git a/payload_generator/delta_diff_utils_unittest.cc b/payload_generator/delta_diff_utils_unittest.cc
index b6983391..53bbeaad 100644
--- a/payload_generator/delta_diff_utils_unittest.cc
+++ b/payload_generator/delta_diff_utils_unittest.cc
@@ -26,8 +26,10 @@
 #include <base/strings/stringprintf.h>
 #include <bsdiff/patch_writer.h>
 #include <gtest/gtest.h>
+#include <puffin/common.h>
 
-#include "payload_generator/filesystem_interface.h"
+#include "update_engine/payload_generator/deflate_utils.h"
+#include "update_engine/payload_generator/filesystem_interface.h"
 #include "update_engine/common/test_utils.h"
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_generator/delta_diff_generator.h"
@@ -845,4 +847,19 @@ TEST_F(DeltaDiffUtilsTest, XorOpsStrided) {
   ASSERT_EQ(aop.xor_ops[3].dst_extent().start_block(), 702UL);
 }
 
+TEST_F(DeltaDiffUtilsTest, FindAndCompactDeflates) {
+  std::vector<puffin::BitExtent> bit_extents{{114122 * 8 * kBlockSize, 1024},
+                                             {114122 * 8 * kBlockSize, 1024}};
+
+  std::vector<Extent> extents = {ExtentForRange(114122, 295),
+                                 ExtentForRange(114418, 16654),
+                                 ExtentForRange(131102, 1),
+                                 ExtentForRange(131104, 307),
+                                 ExtentForRange(131414, 4143),
+                                 ExtentForRange(135559, 8528)};
+  std::vector<puffin::BitExtent> out_deflates;
+  ASSERT_TRUE(deflate_utils::FindAndCompactDeflates(
+      extents, bit_extents, &out_deflates));
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/extent_utils.h b/payload_generator/extent_utils.h
index 52b6d1e0..08636a14 100644
--- a/payload_generator/extent_utils.h
+++ b/payload_generator/extent_utils.h
@@ -174,6 +174,13 @@ constexpr bool ExtentContains(const Extent& big, const Extent& small) {
              big.start_block() + big.num_blocks();
 }
 
+template <typename T>
+constexpr void Dedup(T* container) {
+  std::sort(container->begin(), container->end());
+  container->erase(std::unique(container->begin(), container->end()),
+                   container->end());
+}
+
 }  // namespace chromeos_update_engine
 
 #endif  // UPDATE_ENGINE_PAYLOAD_GENERATOR_EXTENT_UTILS_H_