Fix a bug in the client am: dfec7fab93

am: 3243fc4fa6

Change-Id: Icab788734adfd9e4eedd094157067f6df7afb4eb

11 files changed, 184 insertions, 4 deletions

diff --git a/src/bit_io_unittest.cc b/src/bit_io_unittest.cc
index 87fc3f1..de5f0a7 100644
--- a/src/bit_io_unittest.cc
+++ b/src/bit_io_unittest.cc
@@ -62,4 +62,26 @@ TEST(BitIOTest, BitWriterAndBitReaderTest) {
   ASSERT_FALSE(br.CacheBits(1));
 }
 
+TEST(BitIOTest, BitsRemaining) {
+  const size_t kSize = 5;
+  uint8_t buf[kSize];
+
+  BufferBitReader br(buf, kSize);
+  EXPECT_EQ(br.BitsRemaining(), 40);
+  ASSERT_TRUE(br.CacheBits(1));
+  br.DropBits(1);
+  EXPECT_EQ(br.BitsRemaining(), 39);
+
+  ASSERT_TRUE(br.CacheBits(7));
+  br.DropBits(7);
+  EXPECT_EQ(br.BitsRemaining(), 32);
+
+  ASSERT_TRUE(br.CacheBits(31));
+  br.DropBits(31);
+  EXPECT_EQ(br.BitsRemaining(), 1);
+
+  ASSERT_TRUE(br.CacheBits(1));
+  br.DropBits(1);
+  EXPECT_EQ(br.BitsRemaining(), 0);
+}
 }  // namespace puffin
diff --git a/src/bit_reader.cc b/src/bit_reader.cc
index 27f1d8d..0ed68ef 100644
--- a/src/bit_reader.cc
+++ b/src/bit_reader.cc
@@ -68,4 +68,8 @@ uint64_t BufferBitReader::OffsetInBits() const {
   return (index_ * 8) - in_cache_bits_;
 }
 
+uint64_t BufferBitReader::BitsRemaining() const {
+  return ((in_size_ - index_) * 8) + in_cache_bits_;
+}
+
 }  // namespace puffin
diff --git a/src/bit_reader.h b/src/bit_reader.h
index 2d0c805..aa4a97f 100644
--- a/src/bit_reader.h
+++ b/src/bit_reader.h
@@ -68,6 +68,9 @@ class BitReaderInterface {
 
   // Returns the number of bits read (dropped) till now.
   virtual uint64_t OffsetInBits() const = 0;
+
+  // Returns the number of bits remaining to be cached.
+  virtual uint64_t BitsRemaining() const = 0;
 };
 
 // A raw buffer implementation of |BitReaderInterface|.
@@ -97,6 +100,7 @@ class BufferBitReader : public BitReaderInterface {
       std::function<bool(uint8_t* buffer, size_t count)>* read_fn) override;
   size_t Offset() const override;
   uint64_t OffsetInBits() const override;
+  uint64_t BitsRemaining() const override;
 
  private:
   const uint8_t* in_buf_;  // The input buffer.
diff --git a/src/include/puffin/puffer.h b/src/include/puffin/puffer.h
index b9ff6bc..f7b7c98 100644
--- a/src/include/puffin/puffer.h
+++ b/src/include/puffin/puffer.h
@@ -19,6 +19,15 @@ class HuffmanTable;
 
 class PUFFIN_EXPORT Puffer {
  public:
+  // In older versions of puffin, there is a bug in the client which incorrectly
+  // identifies the number of bits to cache when number of bits for the current
+  // distance plus the number of bits for end of block Huffman code is smaller
+  // than the maximum number of bits needed for distance. If this situations
+  // happens at the very end of the block, it incorrectly tries to cache more
+  // bits than we have and crashes as a result. If |exclude_bad_distance_caches|
+  // is true, we identify those problematic deflate buffers and exclude them
+  // from the list of available deflates. The default is false.
+  explicit Puffer(bool exclude_bad_distance_caches);
   Puffer();
   ~Puffer();
 
@@ -39,6 +48,8 @@ class PUFFIN_EXPORT Puffer {
   std::unique_ptr<HuffmanTable> dyn_ht_;
   std::unique_ptr<HuffmanTable> fix_ht_;
 
+  bool exclude_bad_distance_caches_;
+
   DISALLOW_COPY_AND_ASSIGN(Puffer);
 };
 
diff --git a/src/include/puffin/utils.h b/src/include/puffin/utils.h
index f8eddbb..517d575 100644
--- a/src/include/puffin/utils.h
+++ b/src/include/puffin/utils.h
@@ -92,6 +92,12 @@ void RemoveEqualBitExtents(const Buffer& data1,
                            std::vector<BitExtent>* extents1,
                            std::vector<BitExtent>* extents2);
 
+// Using |data| it removes all the deflate extents from |deflates| which have
+// the problem identified in crbug.com/915559. Each element of |deflates| should
+// contain exactly one deflate block otherwire it returns false.
+bool RemoveDeflatesWithBadDistanceCaches(const Buffer& data,
+                                         std::vector<BitExtent>* deflates);
+
 }  // namespace puffin
 
 #endif  // SRC_INCLUDE_PUFFIN_UTILS_H_
diff --git a/src/puffer.cc b/src/puffer.cc
index a514968..f3a54f8 100644
--- a/src/puffer.cc
+++ b/src/puffer.cc
@@ -23,7 +23,12 @@ using std::vector;
 
 namespace puffin {
 
-Puffer::Puffer() : dyn_ht_(new HuffmanTable()), fix_ht_(new HuffmanTable()) {}
+Puffer::Puffer(bool exclude_bad_distance_caches)
+    : dyn_ht_(new HuffmanTable()),
+      fix_ht_(new HuffmanTable()),
+      exclude_bad_distance_caches_(exclude_bad_distance_caches) {}
+
+Puffer::Puffer() : Puffer(false) {}
 
 Puffer::~Puffer() {}
 
@@ -129,6 +134,10 @@ bool Puffer::PuffDeflate(BitReaderInterface* br,
         return false;
     }
 
+    // If true and the list of output |deflates| is non-null, the current
+    // deflate location will be added to that list.
+    bool include_deflate = true;
+
     while (true) {  // Breaks when the end of block is reached.
       auto max_bits = cur_ht->LitLenMaxBits();
       if (!br->CacheBits(max_bits)) {
@@ -152,7 +161,7 @@ bool Puffer::PuffDeflate(BitReaderInterface* br,
       } else if (256 == lit_len_alphabet) {
         pd.type = PuffData::Type::kEndOfBlock;
         TEST_AND_RETURN_FALSE(pw->Insert(pd));
-        if (deflates != nullptr) {
+        if (deflates != nullptr && include_deflate) {
           deflates->emplace_back(start_bit_offset,
                                  br->OffsetInBits() - start_bit_offset);
         }
@@ -170,8 +179,21 @@ bool Puffer::PuffDeflate(BitReaderInterface* br,
         }
         auto length = kLengthBases[len_code_start] + extra_bits_value;
 
-        TEST_AND_RETURN_FALSE(br->CacheBits(cur_ht->DistanceMaxBits()));
-        auto bits = br->ReadBits(cur_ht->DistanceMaxBits());
+        auto bits_to_cache = cur_ht->DistanceMaxBits();
+        if (!br->CacheBits(bits_to_cache)) {
+          // This is a corner case that is present in the older versions of the
+          // puffin. So we need to catch it and correctly discard this kind of
+          // deflate when we encounter it. See crbug.com/915559 for more info.
+          bits_to_cache = br->BitsRemaining();
+          TEST_AND_RETURN_FALSE(br->CacheBits(bits_to_cache));
+          if (exclude_bad_distance_caches_) {
+            include_deflate = false;
+          }
+          LOG(WARNING) << "A rare condition that older puffin clients fail to"
+                       << " recognize happened. Nothing to worry about."
+                       << " See crbug.com/915559";
+        }
+        auto bits = br->ReadBits(bits_to_cache);
         uint16_t distance_alphabet;
         size_t nbits;
         TEST_AND_RETURN_FALSE(
diff --git a/src/puffin_unittest.cc b/src/puffin_unittest.cc
index 2ccf23f..322e97b 100644
--- a/src/puffin_unittest.cc
+++ b/src/puffin_unittest.cc
@@ -574,4 +574,33 @@ TEST_F(PuffinTest, BitExtentPuffAndHuffTest) {
                              kGapPuffs, kGapPuffExtents);
 }
 
+TEST_F(PuffinTest, ExcludeBadDistanceCaches) {
+  BufferBitReader br(kProblematicCache.data(), kProblematicCache.size());
+  BufferPuffWriter pw(nullptr, 0);
+
+  // The first two bits of this data should be ignored.
+  br.CacheBits(2);
+  br.DropBits(2);
+
+  vector<BitExtent> deflates, empty;
+  Puffer puffer(true);
+  EXPECT_TRUE(puffer.PuffDeflate(&br, &pw, &deflates));
+  EXPECT_EQ(deflates, empty);
+}
+
+TEST_F(PuffinTest, NoExcludeBadDistanceCaches) {
+  BufferBitReader br(kProblematicCache.data(), kProblematicCache.size());
+  BufferPuffWriter pw(nullptr, 0);
+
+  // The first two bits of this data should be ignored.
+  br.CacheBits(2);
+  br.DropBits(2);
+
+  vector<BitExtent> deflates;
+  Puffer puffer;  // The default value for excluding bad distance cache should
+                  // be false.
+  EXPECT_TRUE(puffer.PuffDeflate(&br, &pw, &deflates));
+  EXPECT_EQ(deflates, kProblematicCacheDeflateExtents);
+}
+
 }  // namespace puffin
diff --git a/src/unittest_common.cc b/src/unittest_common.cc
index 7b16011..6bc3b1d 100644
--- a/src/unittest_common.cc
+++ b/src/unittest_common.cc
@@ -76,4 +76,16 @@ const std::vector<ByteExtent> kPuffExtentsSample2 = {
   {0, 11}, {14, 11}, {25, 7}};
 // clang-format on
 
+// This data is taken from the failed instances described in crbug.com/915559.
+const Buffer kProblematicCache = {
+    0x51, 0x74, 0x97, 0x71, 0x51, 0x6e, 0x6d, 0x1b, 0x87, 0x4f, 0x5b,
+    0xb1, 0xbb, 0xb6, 0xdd, 0xdd, 0xdd, 0x89, 0x89, 0xa2, 0x88, 0x9d,
+    0x18, 0x4c, 0x1a, 0x8c, 0x8a, 0x1d, 0xa8, 0xd8, 0x89, 0xdd, 0xdd,
+    0x81, 0x89, 0x62, 0x77, 0xb7, 0x32, 0x81, 0x31, 0x98, 0x88, 0x5d,
+    0x83, 0xbd, 0xff, 0xf3, 0xe1, 0xf8, 0x9d, 0xd7, 0xba, 0xd6, 0x9a,
+    0x7b, 0x86, 0x99, 0x3b, 0xf7, 0xbb, 0xdf, 0xfd, 0x90, 0xf0, 0x45,
+    0x0b, 0xb4, 0x44, 0x2b, 0xb4, 0x46, 0x1b, 0xb4, 0xc5, 0xff};
+const std::vector<BitExtent> kProblematicCacheDeflateExtents = {{2, 606}};
+const std::vector<BitExtent> kProblematicCachePuffExtents = {{1, 185}};
+
 }  // namespace puffin
diff --git a/src/unittest_common.h b/src/unittest_common.h
index c73d9b3..eac61e6 100644
--- a/src/unittest_common.h
+++ b/src/unittest_common.h
@@ -46,6 +46,10 @@ extern const std::vector<ByteExtent> kDeflateExtentsSample2;
 extern const std::vector<BitExtent> kSubblockDeflateExtentsSample2;
 extern const std::vector<ByteExtent> kPuffExtentsSample2;
 
+extern const Buffer kProblematicCache;
+extern const std::vector<BitExtent> kProblematicCacheDeflateExtents;
+extern const std::vector<BitExtent> kProblematicCachePuffExtents;
+
 }  // namespace puffin
 
 #endif  // SRC_UNITTEST_COMMON_H_
diff --git a/src/utils.cc b/src/utils.cc
index de3cc6c..e77cf59 100644
--- a/src/utils.cc
+++ b/src/utils.cc
@@ -411,4 +411,34 @@ void RemoveEqualBitExtents(const Buffer& data1,
                      }),
       extents1->end());
 }
+
+bool RemoveDeflatesWithBadDistanceCaches(const Buffer& data,
+                                         vector<BitExtent>* deflates) {
+  Puffer puffer(true /* exclude_bad_distance_caches */);
+  for (auto def = deflates->begin(); def != deflates->end();) {
+    uint64_t offset = def->offset / 8;
+    uint64_t length = (def->offset + def->length + 7) / 8 - offset;
+    BufferBitReader br(&data[offset], length);
+    BufferPuffWriter pw(nullptr, 0);
+
+    // Drop the first few bits in the buffer so we start exactly where the
+    // deflate starts.
+    uint64_t bits_to_drop = def->offset % 8;
+    TEST_AND_RETURN_FALSE(br.CacheBits(bits_to_drop));
+    br.DropBits(bits_to_drop);
+
+    vector<BitExtent> defs_out;
+    TEST_AND_RETURN_FALSE(puffer.PuffDeflate(&br, &pw, &defs_out));
+
+    TEST_AND_RETURN_FALSE(defs_out.size() <= 1);
+    if (defs_out.size() == 0) {
+      // This is a deflate we were looking for, remove it.
+      def = deflates->erase(def);
+    } else {
+      ++def;
+    }
+  }
+  return true;
+}
+
 }  // namespace puffin
diff --git a/src/utils_unittest.cc b/src/utils_unittest.cc
index 43f18d0..bbeecbf 100644
--- a/src/utils_unittest.cc
+++ b/src/utils_unittest.cc
@@ -242,4 +242,40 @@ TEST(UtilsTest, RemoveEqualBitExtents) {
   EXPECT_EQ(expected_ext2, ext2);
 }
 
+TEST(UtilsTest, RemoveDeflatesWithBadDistanceCaches) {
+  vector<BitExtent> deflates(kProblematicCacheDeflateExtents), empty;
+  EXPECT_TRUE(
+      RemoveDeflatesWithBadDistanceCaches(kProblematicCache, &deflates));
+  EXPECT_EQ(deflates, empty);
+
+  // Just a sanity check to make sure this function is not removing anything
+  // else.
+  deflates = kSubblockDeflateExtentsSample1;
+  EXPECT_TRUE(RemoveDeflatesWithBadDistanceCaches(kDeflatesSample1, &deflates));
+  EXPECT_EQ(deflates, kSubblockDeflateExtentsSample1);
+
+  // Now combine three deflates and make sure it is doing the right job.
+  Buffer data;
+  data.insert(data.end(), kDeflatesSample1.begin(), kDeflatesSample1.end());
+  data.insert(data.end(), kProblematicCache.begin(), kProblematicCache.end());
+  data.insert(data.end(), kDeflatesSample1.begin(), kDeflatesSample1.end());
+
+  deflates = kSubblockDeflateExtentsSample1;
+  size_t offset = kDeflatesSample1.size() * 8;
+  for (const auto& deflate : kProblematicCacheDeflateExtents) {
+    deflates.emplace_back(deflate.offset + offset, deflate.length);
+  }
+  offset += kProblematicCache.size() * 8;
+  for (const auto& deflate : kSubblockDeflateExtentsSample1) {
+    deflates.emplace_back(deflate.offset + offset, deflate.length);
+  }
+
+  auto expected_deflates(deflates);
+  expected_deflates.erase(expected_deflates.begin() +
+                          kSubblockDeflateExtentsSample1.size());
+
+  EXPECT_TRUE(RemoveDeflatesWithBadDistanceCaches(data, &deflates));
+  EXPECT_EQ(deflates, expected_deflates);
+}
+
 }  // namespace puffin