Make NaluReader skip encrypted portion

Change-Id: Ibb47a1e62cd8ac3057c8f1512a88280991e48b62

packager/media/codecs/nalu_reader.cc

@@ -19,6 +19,55 @@ namespace {
 inline bool IsStartCode(const uint8_t* data) {
   return data[0] == 0x00 && data[1] == 0x00 && data[2] == 0x01;
 }
+
+// Edits |subsamples| given the number of consumed bytes.
+void UpdateSubsamples(uint64_t consumed_bytes,
+                      std::vector<SubsampleEntry>* subsamples) {
+  if (consumed_bytes == 0 || subsamples->empty()) {
+    return;
+  }
+  size_t num_entries_to_delete = 0;
+  for (SubsampleEntry& subsample : *subsamples) {
+    if (subsample.clear_bytes > consumed_bytes) {
+      subsample.clear_bytes -= consumed_bytes;
+      consumed_bytes = 0;
+      break;
+    }
+    consumed_bytes -= subsample.clear_bytes;
+    subsample.clear_bytes = 0;
+
+    if (subsample.cipher_bytes > consumed_bytes) {
+      subsample.cipher_bytes -= consumed_bytes;
+      consumed_bytes = 0;
+      break;
+    }
+    consumed_bytes -= subsample.cipher_bytes;
+    subsample.cipher_bytes = 0;
+    ++num_entries_to_delete;
+  }
+
+  subsamples->erase(subsamples->begin(),
+                    subsamples->begin() + num_entries_to_delete);
+}
+
+bool IsNaluLengthEncrypted(
+    uint8_t nalu_length_size,
+    const std::vector<SubsampleEntry>& subsamples) {
+  if (subsamples.empty())
+    return false;
+
+  for (const SubsampleEntry& subsample : subsamples) {
+    if (subsample.clear_bytes >= nalu_length_size) {
+      return false;
+    }
+    nalu_length_size -= subsample.clear_bytes;
+    if (subsample.cipher_bytes > 0) {
+      return true;
+    }
+  }
+  // Ran out of subsamples. Assume the rest is in the clear.
+  return false;
+}
 }  // namespace
 
 Nalu::Nalu() = default;
@@ -164,14 +213,27 @@ NaluReader::NaluReader(Nalu::CodecType type,
                        uint8_t nal_length_size,
                        const uint8_t* stream,
                        uint64_t stream_size)
+    : NaluReader(type,
+                 nal_length_size,
+                 stream,
+                 stream_size,
+                 std::vector<SubsampleEntry>()) {}
+
+NaluReader::NaluReader(Nalu::CodecType type,
+                       uint8_t nal_length_size,
+                       const uint8_t* stream,
+                       uint64_t stream_size,
+                       const std::vector<SubsampleEntry>& subsamples)
     : stream_(stream),
       stream_size_(stream_size),
       nalu_type_(type),
       nalu_length_size_(nal_length_size),
       format_(nal_length_size == 0 ? kAnnexbByteStreamFormat
-                                   : kNalUnitStreamFormat) {
+                                   : kNalUnitStreamFormat),
+      subsamples_(subsamples) {
   DCHECK(stream);
 }
+
 NaluReader::~NaluReader() {}
 
 NaluReader::Result NaluReader::Advance(Nalu* nalu) {
@@ -195,6 +257,10 @@ NaluReader::Result NaluReader::Advance(Nalu* nalu) {
     nalu_length = nalu_length_with_header - nalu_length_size_or_start_code_size;
   } else {
     BufferReader reader(stream_, stream_size_);
+    if (IsNaluLengthEncrypted(nalu_length_size_, subsamples_)) {
+      LOG(ERROR) << "NALU length is encrypted.";
+      return NaluReader::kInvalidStream;
+    }
     if (!reader.ReadNBytesInto8(&nalu_length, nalu_length_size_))
       return NaluReader::kInvalidStream;
     nalu_length_size_or_start_code_size = nalu_length_size_;
@@ -218,6 +284,8 @@ NaluReader::Result NaluReader::Advance(Nalu* nalu) {
   // is called, we will effectively be skipping it.
   stream_ += nalu_length_size_or_start_code_size + nalu_length;
   stream_size_ -= nalu_length_size_or_start_code_size + nalu_length;
+  UpdateSubsamples(nalu_length_size_or_start_code_size + nalu_length,
+                   &subsamples_);
 
   DVLOG(4) << "NALU type: " << static_cast<int>(nalu->type())
            << " at: " << reinterpret_cast<const void*>(nalu->data())
@@ -272,13 +340,72 @@ bool NaluReader::FindStartCode(const uint8_t* data,
   return false;
 }
 
+// static
+bool NaluReader::FindStartCodeInClearRange(
+    const uint8_t* data,
+    uint64_t data_size,
+    uint64_t* offset,
+    uint8_t* start_code_size,
+    const std::vector<SubsampleEntry>& subsamples) {
+  if (subsamples.empty()) {
+    return FindStartCode(data, data_size, offset, start_code_size);
+  }
+
+  uint64_t current_offset = 0;
+  for (const SubsampleEntry& subsample : subsamples) {
+    uint16_t clear_bytes = subsample.clear_bytes;
+    if (current_offset + clear_bytes > data_size) {
+      LOG(WARNING) << "The sum of subsample sizes is greater than data_size.";
+      clear_bytes = data_size - current_offset;
+    }
+
+    // Note that calling FindStartCode() here should get the correct
+    // start_code_size, even tho data + current_offset may be in the middle of
+    // the buffer because data + current_offset - 1 is either it shouldn't be
+    // accessed because it's data - 1 or it is encrypted.
+    const bool found_start_code = FindStartCode(
+        data + current_offset, clear_bytes, offset, start_code_size);
+    if (found_start_code) {
+      *offset += current_offset;
+      return true;
+    }
+    const uint64_t subsample_size =
+        subsample.clear_bytes + subsample.cipher_bytes;
+    current_offset += subsample_size;
+    if (current_offset > data_size) {
+      // Assign data_size here so that the returned offset points to the end of
+      // the data.
+      current_offset = data_size;
+      LOG(WARNING) << "The sum of subsamples is greater than data_size.";
+      break;
+    }
+  }
+
+  // If there is more that's not specified by the subsample entries, assume it
+  // is in the clear.
+  if (current_offset < data_size) {
+    const bool found_start_code =
+        FindStartCode(data + current_offset, data_size - current_offset, offset,
+                      start_code_size);
+    *offset += current_offset;
+    return found_start_code;
+  }
+
+  // End of data: offset is pointing to the first byte that was not considered
+  // as a possible start of a start code.
+  *offset = current_offset;
+  *start_code_size = 0;
+  return false;
+}
+
 bool NaluReader::LocateNaluByStartCode(uint64_t* nalu_size,
                                        uint8_t* start_code_size) {
   // Find the start code of next NALU.
   uint64_t nalu_start_off = 0;
   uint8_t annexb_start_code_size = 0;
-  if (!FindStartCode(stream_, stream_size_,
+  if (!FindStartCodeInClearRange(
-                     &nalu_start_off, &annexb_start_code_size)) {
+          stream_, stream_size_,
+          &nalu_start_off, &annexb_start_code_size, subsamples_)) {
     DVLOG(4) << "Could not find start code, end of stream?";
     return false;
   }
@@ -286,8 +413,18 @@ bool NaluReader::LocateNaluByStartCode(uint64_t* nalu_size,
   // Move the stream to the beginning of the NALU (pointing at the start code).
   stream_ += nalu_start_off;
   stream_size_ -= nalu_start_off;
+  // Shift the subsamples so that next call to FindStartCode() takes the updated
+  // subsample info.
+  UpdateSubsamples(nalu_start_off, &subsamples_);
 
   const uint8_t* nalu_data = stream_ + annexb_start_code_size;
+  // This is a temporary subsample entries for finding next nalu. subsamples_
+  // should not be updated below.
+  std::vector<SubsampleEntry> subsamples_for_finding_next_nalu;
+  if (!subsamples_.empty()) {
+    subsamples_for_finding_next_nalu = subsamples_;
+    UpdateSubsamples(annexb_start_code_size, &subsamples_for_finding_next_nalu);
+  }
   uint64_t max_nalu_data_size = stream_size_ - annexb_start_code_size;
   if (max_nalu_data_size <= 0) {
     DVLOG(3) << "End of stream";
@@ -303,14 +440,18 @@ bool NaluReader::LocateNaluByStartCode(uint64_t* nalu_size,
   uint64_t nalu_size_without_start_code = 0;
   uint8_t next_start_code_size = 0;
   while (true) {
-    if (!FindStartCode(nalu_data, max_nalu_data_size,
+    if (!FindStartCodeInClearRange(
-                       &nalu_size_without_start_code, &next_start_code_size)) {
+            nalu_data, max_nalu_data_size,
+            &nalu_size_without_start_code, &next_start_code_size,
+            subsamples_for_finding_next_nalu)) {
       nalu_data += max_nalu_data_size;
       break;
     }
 
     nalu_data += nalu_size_without_start_code + next_start_code_size;
     max_nalu_data_size -= nalu_size_without_start_code + next_start_code_size;
+    UpdateSubsamples(nalu_size_without_start_code + next_start_code_size,
+                     &subsamples_for_finding_next_nalu);
     // If it is not a valid NAL unit, we will continue searching. This is to
     // handle the case where emulation prevention are not applied.
     Nalu nalu;

packager/media/codecs/nalu_reader.h

@@ -12,6 +12,7 @@
 
 #include "packager/base/compiler_specific.h"
 #include "packager/base/macros.h"
+#include "packager/media/base/decrypt_config.h"
 
 namespace shaka {
 namespace media {
@@ -160,6 +161,22 @@ class NaluReader {
              uint8_t nal_length_size,
              const uint8_t* stream,
              uint64_t stream_size);
+
+  /// @param type is the codec type of the NALU unit.
+  /// @param nalu_length_size should be set to 0 for AnnexB byte streams;
+  ///        otherwise, it indicates the size of NAL unit length for the NAL
+  ///        unit stream.
+  /// @param stream is the input stream.
+  /// @param stream_size is the size of @a stream.
+  /// @param subsamples specifies the clear and encrypted sections of the
+  ///        @a stream starting from the beginning of the @a stream. If
+  ///        @a subsamples doesn't cover the entire stream, then the rest is
+  ///        assumed to be in the clear.
+  NaluReader(Nalu::CodecType type,
+             uint8_t nal_length_size,
+             const uint8_t* stream,
+             uint64_t stream_size,
+             const std::vector<SubsampleEntry>& subsamples);
   ~NaluReader();
 
   // Find offset from start of data to next NALU start code
@@ -176,6 +193,21 @@ class NaluReader {
                             uint64_t* offset,
                             uint8_t* start_code_size);
 
+  /// Same as FindStartCode() but also specify the subsamples. This searches for
+  /// start codes in the clear section and will not scan for start codes in the
+  /// encrypted section. Even if there is a real NALU start code in the
+  /// encrypted section, this will skip them.
+  /// @param subsamples starting from the start of @a data. If @a subsamples
+  ///        does not cover the whole @a data, the rest is assumed to be in the
+  ///        clear.
+  /// @return true if it finds a NALU. false otherwise.
+  static bool FindStartCodeInClearRange(
+      const uint8_t* data,
+      uint64_t data_size,
+      uint64_t* offset,
+      uint8_t* start_code_size,
+      const std::vector<SubsampleEntry>& subsamples);
+
   /// Reads a NALU from the stream into |*nalu|, if one exists, and then
   /// advances to the next NALU.
   /// @param nalu contains the NALU read if it exists.
@@ -213,6 +245,9 @@ class NaluReader {
   // The format of the stream.
   Format format_;
 
+  // subsamples left in stream_.
+  std::vector<SubsampleEntry> subsamples_;
+
   DISALLOW_COPY_AND_ASSIGN(NaluReader);
 };
 

packager/media/codecs/nalu_reader_unittest.cc

@@ -187,5 +187,218 @@ TEST(NaluReaderTest, ErrorForZeroSize) {
   EXPECT_FALSE(nalu.Initialize(Nalu::kH265, kNaluData, 0));
 }
 
+TEST(NaluReaderTest, SubsamplesAnnexB) {
+  const uint8_t kNaluData[] = {
+      // This array contains 1 nalu starting with a NALU start code.
+      // what looks like NALU start codes below are "encrypted" portion.
+      0x00, 0x00, 0x01, 0x14,
+      // This is in the encrypted portion and none of the following sequence
+      // should be recognized as a NALU start code.
+      0x00, 0x00, 0x01, 0x65, 0x00, 0x00, 0x00, 0x01, 0x67,
+  };
+  std::vector<SubsampleEntry> subsamples;
+  subsamples.emplace_back(SubsampleEntry(4, 9));
+  NaluReader reader(Nalu::kH264, kIsAnnexbByteStream, kNaluData,
+                    arraysize(kNaluData), subsamples);
+
+  Nalu nalu;
+  ASSERT_EQ(NaluReader::kOk, reader.Advance(&nalu));
+  EXPECT_EQ(kNaluData + 3, nalu.data());
+  EXPECT_EQ(9u, nalu.payload_size());
+  EXPECT_EQ(1u, nalu.header_size());
+  EXPECT_EQ(0, nalu.ref_idc());
+  EXPECT_EQ(0x14, nalu.type());
+}
+
+TEST(NaluReaderTest, MultiSubsamplesAnnexB) {
+  const uint8_t kNaluData[] = {
+      // Clear
+      0x00,
+      // Encrypted. Should not recognize this as a NALU start code.
+      0x00, 0x01, 0x14,
+      // Clear. Valid NALU start code + NALU header.
+      0x00, 0x00, 0x01, 0x65,
+      // Encrypted.
+      0x00, 0x00, 0x00, 0x01, 0x67,
+  };
+  std::vector<SubsampleEntry> subsamples;
+  subsamples.emplace_back(SubsampleEntry(1, 3));
+  subsamples.emplace_back(SubsampleEntry(4, 5));
+  NaluReader reader(Nalu::kH264, kIsAnnexbByteStream, kNaluData,
+                    arraysize(kNaluData), subsamples);
+
+  Nalu nalu;
+  ASSERT_EQ(NaluReader::kOk, reader.Advance(&nalu));
+  EXPECT_EQ(kNaluData + 7, nalu.data());
+  EXPECT_EQ(5u, nalu.payload_size());
+  EXPECT_EQ(1u, nalu.header_size());
+  EXPECT_EQ(3, nalu.ref_idc());
+  EXPECT_EQ(5, nalu.type());
+}
+
+// Verify that data outside subsamples is treated as clear data.
+TEST(NaluReaderTest, BufferBiggerThanSubsamplesAnnexB) {
+  const uint8_t kNaluData[] = {
+      // This array contains 1 nalu starting with a NALU start code.
+      // what looks like NALU start codes below are "encrypted" portion.
+      0x00, 0x00, 0x01, 0x14,
+      // This is in the encrypted portion and none of the following sequence
+      // should be recognized as a NALU start code.
+      0x00, 0x00, 0x01, 0x65, 0x00, 0x00, 0x00, 0x01, 0x67,
+      // Start of second NALU not specified by subsamples.
+      0x00, 0x00, 0x00, 0x01, 0x67, 0xbb, 0xcc, 0xdd,
+  };
+  std::vector<SubsampleEntry> subsamples;
+  subsamples.emplace_back(SubsampleEntry(4, 9));
+  NaluReader reader(Nalu::kH264, kIsAnnexbByteStream, kNaluData,
+                    arraysize(kNaluData), subsamples);
+
+  Nalu nalu;
+  ASSERT_EQ(NaluReader::kOk, reader.Advance(&nalu));
+  EXPECT_EQ(kNaluData + 3, nalu.data());
+  EXPECT_EQ(9u, nalu.payload_size());
+  EXPECT_EQ(1u, nalu.header_size());
+  EXPECT_EQ(0, nalu.ref_idc());
+  EXPECT_EQ(0x14, nalu.type());
+
+  ASSERT_EQ(NaluReader::kOk, reader.Advance(&nalu));
+  EXPECT_EQ(kNaluData + 17, nalu.data());
+  EXPECT_EQ(3u, nalu.payload_size());
+  EXPECT_EQ(1u, nalu.header_size());
+  EXPECT_EQ(3, nalu.ref_idc());
+  EXPECT_EQ(7, nalu.type());
+}
+
+// Finds a NALU start code + header in the clear section but is an invalid NALU.
+TEST(NaluReaderTest, SubsamplesWithInvalidNalu) {
+  const uint8_t kNaluData[] = {
+      // Start with a valid NALU.
+      // Clear.
+      0x00, 0x00, 0x01, 0x14,
+      // Encrypted.
+      0x00, 0x00,
+      // Clear. Has NALU start code but invalid NALU.
+      0x00, 0x00, 0x01, 0x80,
+      // Encrypted.
+      0x00, 0x04, 0x03,
+      // Clear.
+      0x00, 0xFE,
+      // Encrypted.
+      0xFF, 0xFE, 0xFD, 0xFC, 0xFB, 0xFA, 0x00, 0x01,
+      // Clear. Valid NALU. The first NALU should end here.
+      // If subsamples is not updated correctly the parser won't recognize that
+      // this is a NALU start code.
+      0x00, 0x00, 0x01, 0x65,
+      // Encrypted.
+      0xEE, 0xCE, 0x12, 0x44,
+  };
+  std::vector<SubsampleEntry> subsamples;
+  subsamples.emplace_back(SubsampleEntry(4, 2));
+  subsamples.emplace_back(SubsampleEntry(4, 3));
+  subsamples.emplace_back(SubsampleEntry(2, 8));
+  subsamples.emplace_back(SubsampleEntry(4, 4));
+
+  NaluReader reader(Nalu::kH264, kIsAnnexbByteStream, kNaluData,
+                    arraysize(kNaluData), subsamples);
+
+  Nalu nalu;
+  ASSERT_EQ(NaluReader::kOk, reader.Advance(&nalu));
+  EXPECT_EQ(19u, nalu.payload_size());
+  EXPECT_EQ(1u, nalu.header_size());
+  EXPECT_EQ(0, nalu.ref_idc());
+  EXPECT_EQ(0x14, nalu.type());
+}
+
+// No NALU start code in the subsample range. A NALU start code in the buffer
+// not specified by subsamples.
+TEST(NaluReaderTest, FindStartCodeInClearRangeNoNalu) {
+  const uint8_t kNaluData[] = {
+      // Any sequence not NALU start code in the subsample region.
+      0xFF, 0xFE, 0xFD, 0xFC,
+      // End of subsample specified region. No NALU start code.
+      0x00, 0x04, 0x03, 0x14, 0x34, 0x56, 0x78,
+  };
+  std::vector<SubsampleEntry> subsamples;
+  subsamples.emplace_back(SubsampleEntry(2, 2));
+
+  uint64_t offset = 0;
+  uint8_t start_code_size = 0;
+  EXPECT_FALSE(NaluReader::FindStartCodeInClearRange(
+      kNaluData, arraysize(kNaluData), &offset, &start_code_size, subsamples));
+  EXPECT_GT(offset, 4u)
+      << "Expect at least the subsample region should be consumed.";
+}
+
+// If subsamples goes beyond the data size and cannot find a NALU start code,
+// |offset| should not be set to the end of the subsamples. Instead it should be
+// less than or equal to the size of the data as documented in the header.
+TEST(NaluReaderTest, FindStartCodeInClearRangeSubsamplesBiggerThanBuffer) {
+  const uint8_t kNaluData[] = {
+    // The data in here doesn't really matter.
+    0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+  };
+  std::vector<SubsampleEntry> subsamples;
+  subsamples.emplace_back(SubsampleEntry(1, 14));
+
+  uint64_t offset;
+  uint8_t start_code_size;
+  EXPECT_FALSE(NaluReader::FindStartCodeInClearRange(
+      kNaluData, arraysize(kNaluData), &offset, &start_code_size, subsamples));
+  EXPECT_LE(offset, arraysize(kNaluData));
+}
+
+// Verify that it doesn't affect the Nalu stream mode too much.
+TEST(NaluReaderTest, SubsamplesNaluStream) {
+  const uint8_t kNaluData[] = {
+      // This array contains 1 nalu starting with a 1 byte NALU length size.
+      0x0A, 0x14,
+      // This is in the encrypted portion and none of the following sequence
+      // should be recognized as a NALU start code.
+      0x00, 0x00, 0x01, 0x65, 0x00, 0x00, 0x00, 0x01, 0x67,
+  };
+  std::vector<SubsampleEntry> subsamples;
+  subsamples.emplace_back(SubsampleEntry(2, 9));
+  NaluReader reader(Nalu::kH264, 1, kNaluData,
+                    arraysize(kNaluData), subsamples);
+
+  Nalu nalu;
+  ASSERT_EQ(NaluReader::kOk, reader.Advance(&nalu));
+  EXPECT_EQ(kNaluData + 1, nalu.data());
+  EXPECT_EQ(9u, nalu.payload_size());
+  EXPECT_EQ(1u, nalu.header_size());
+  EXPECT_EQ(0, nalu.ref_idc());
+  EXPECT_EQ(0x14, nalu.type());
+}
+
+// Verify that if NALU length is encrypted, NALUs cannot be parsed.
+TEST(NaluReaderTest, EncryptedNaluLengthNaluStream) {
+  const uint8_t kNaluData[] = {
+      // This array contains 1 nalu starting with a 1 byte NALU length size.
+      0x00, 0x0A, 0x14,
+      // This is in the encrypted portion and none of the following sequence
+      // should be recognized as a NALU start code.
+      0x00, 0x00, 0x01, 0x65, 0x00, 0x00, 0x00, 0x01, 0x67,
+      // Second NALU is supposed to start here but the second byte of the length
+      // is encrypted.
+      0x00, 0xFF, 0xFF,
+  };
+
+  std::vector<SubsampleEntry> subsamples;
+  subsamples.emplace_back(SubsampleEntry(3, 9));
+  subsamples.emplace_back(SubsampleEntry(1, 2));
+  NaluReader reader(Nalu::kH264, 2, kNaluData,
+                    arraysize(kNaluData), subsamples);
+
+  Nalu nalu;
+  ASSERT_EQ(NaluReader::kOk, reader.Advance(&nalu));
+  EXPECT_EQ(kNaluData + 2, nalu.data());
+  EXPECT_EQ(9u, nalu.payload_size());
+  EXPECT_EQ(1u, nalu.header_size());
+  EXPECT_EQ(0, nalu.ref_idc());
+  EXPECT_EQ(0x14, nalu.type());
+
+  ASSERT_EQ(NaluReader::kInvalidStream, reader.Advance(&nalu));
+}
+
 }  // namespace media
 }  // namespace shaka