From d1caa29c8da8ed28a5eb39901e35013540d8d639 Mon Sep 17 00:00:00 2001
From: KongQun Yang <kqyang@google.com>
Date: Tue, 26 Jun 2018 11:48:47 -0700
Subject: [PATCH] Disregard trailing null bytes when locating RBSP stop bit

This is a more faithful implementation of more_rbsp_data().

There could be trailing null bytes in NAL units. This isn't valid per
H264 specification, but the referenced bug includes a sample where the
PPS in the avcC record includes a trailing null byte.

Workaround the problem so packager does not fail.

A similar problem is workarounded in Chrome:
https://codereview.chromium.org/1107593004

Closes #418

Change-Id: I28cb8a9371945dc094f766c3e559d7a66859b451
---
 packager/media/codecs/h26x_bit_reader.cc      | 35 ++++++++++++++-----
 .../media/codecs/h26x_bit_reader_unittest.cc  | 16 +++++++++
 2 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/packager/media/codecs/h26x_bit_reader.cc b/packager/media/codecs/h26x_bit_reader.cc
index e1ef625009..ae3a8bcf1d 100644
--- a/packager/media/codecs/h26x_bit_reader.cc
+++ b/packager/media/codecs/h26x_bit_reader.cc
@@ -7,6 +7,14 @@
 
 namespace shaka {
 namespace media {
+namespace {
+
+// Check if any bits in the least significant |valid_bits| are set to 1.
+bool CheckAnyBitsSet(int byte, int valid_bits) {
+  return (byte & ((1 << valid_bits) - 1)) != 0;
+}
+
+}  // namespace
 
 H26xBitReader::H26xBitReader()
     : data_(NULL),
@@ -145,20 +153,29 @@ off_t H26xBitReader::NumBitsLeft() {
 }
 
 bool H26xBitReader::HasMoreRBSPData() {
-  // Make sure we have more bits, if we are at 0 bits in current byte
-  // and updating current byte fails, we don't have more data anyway.
+  // Make sure we have more bits, if we are at 0 bits in current byte and
+  // updating current byte fails, we don't have more data anyway.
   if (num_remaining_bits_in_curr_byte_ == 0 && !UpdateCurrByte())
     return false;
 
-  // On last byte?
-  if (bytes_left_)
+  // If there is no more RBSP data, then the remaining bits is the stop bit
+  // followed by zero paddings. So if there are 1s in the remaining bits
+  // excluding the current bit, then the current bit is not a stop bit,
+  // regardless of whether it is 1 or not. Therefore there is more data.
+  if (CheckAnyBitsSet(curr_byte_, num_remaining_bits_in_curr_byte_ - 1))
     return true;
 
-  // Last byte, look for stop bit;
-  // We have more RBSP data if the last non-zero bit we find is not the
-  // first available bit.
-  return (curr_byte_ &
-          ((1 << (num_remaining_bits_in_curr_byte_ - 1)) - 1)) != 0;
+  // While the spec disallows it (7.4.1: "The last byte of the NAL unit shall
+  // not be equal to 0x00"), some streams have trailing null bytes anyway. We
+  // don't handle emulation prevention sequences because HasMoreRBSPData() is
+  // not used when parsing slices (where cabac_zero_word elements are legal).
+  for (off_t i = 0; i < bytes_left_; i++) {
+    if (data_[i] != 0)
+      return true;
+  }
+
+  bytes_left_ = 0;
+  return false;
 }
 
 size_t H26xBitReader::NumEmulationPreventionBytesRead() {
diff --git a/packager/media/codecs/h26x_bit_reader_unittest.cc b/packager/media/codecs/h26x_bit_reader_unittest.cc
index 0dff514765..66f37830f7 100644
--- a/packager/media/codecs/h26x_bit_reader_unittest.cc
+++ b/packager/media/codecs/h26x_bit_reader_unittest.cc
@@ -42,6 +42,22 @@ TEST(H26xBitReaderTest, ReadStreamWithoutEscapeAndTrailingZeroBytes) {
   EXPECT_FALSE(reader.HasMoreRBSPData());
 }
 
+TEST(H26xBitReaderTest, ReadPpsWithTrailingZeroByte) {
+  H26xBitReader reader;
+
+  // Data copied from https://github.com/google/shaka-packager/issues/418.
+  const unsigned char pps_rbsp[] = {0xee, 0x3c, 0x80, 0x00};
+  EXPECT_TRUE(reader.Initialize(pps_rbsp, sizeof(pps_rbsp)));
+
+  // Skips all the fields in PPS (kind of simulates ParsePps).
+  EXPECT_TRUE(reader.SkipBits(16));
+
+  EXPECT_EQ(reader.NumBitsLeft(), 16);
+  // The remaining data is '80 00'. The trailing null byte is ignored. There
+  // are no bits before the stop bit, so there is no more RBSP data.
+  EXPECT_FALSE(reader.HasMoreRBSPData());
+}
+
 TEST(H26xBitReaderTest, SingleByteStream) {
   H26xBitReader reader;
   const unsigned char rbsp[] = {0x18};