Add encrypted range handling to H265Parser

BUG=chromium:1141237,b:153111783
TEST=media_unittests, media_h265_parser_fuzzer

Change-Id: I71ff94421a5dea803b7cd26bb4da36c47ccb7f8d
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2510136Reviewed-by: Sergey Volk <servolk@chromium.org>
Reviewed-by: Dale Curtis <dalecurtis@chromium.org>
Commit-Queue: Jeffrey Kardatzke <jkardatzke@google.com>
Cr-Commit-Position: refs/heads/master@{#822837}

media/video/h265_parser.cc

@@ -44,6 +44,38 @@ constexpr int kTableSarHeight[] = {0,  1,  11, 11, 11, 33, 11, 11, 11,
 static_assert(base::size(kTableSarWidth) == base::size(kTableSarHeight),
               "sar tables must have the same size");
 
+// Converts [|start|, |end|) range with |encrypted_ranges| into a vector of
+// SubsampleEntry. |encrypted_ranges| must be with in the range defined by
+// |start| and |end|.
+// It is OK to pass in empty |encrypted_ranges|; this will return a vector
+// with single SubsampleEntry with clear_bytes set to the size of the buffer.
+std::vector<SubsampleEntry> EncryptedRangesToSubsampleEntry(
+    const uint8_t* start,
+    const uint8_t* end,
+    const Ranges<const uint8_t*>& encrypted_ranges) {
+  std::vector<SubsampleEntry> subsamples(encrypted_ranges.size());
+  const uint8_t* cur = start;
+  for (size_t i = 0; i < encrypted_ranges.size(); ++i) {
+    const uint8_t* encrypted_start = encrypted_ranges.start(i);
+    DCHECK_GE(encrypted_start, cur)
+        << "Encrypted range started before the current buffer pointer.";
+    subsamples[i].clear_bytes = encrypted_start - cur;
+
+    const uint8_t* encrypted_end = encrypted_ranges.end(i);
+    subsamples[i].cypher_bytes = encrypted_end - encrypted_start;
+
+    cur = encrypted_end;
+    DCHECK_LE(cur, end) << "Encrypted range is outside the buffer range.";
+  }
+
+  // If there is more data in the buffer but not covered by encrypted_ranges,
+  // then it must be in the clear.
+  if (cur < end)
+    subsamples.emplace_back(end - cur, 0);
+
+  return subsamples;
+}
+
 void FillInDefaultScalingListData(H265ScalingListData* scaling_list_data,
                                   int size_id,
                                   int matrix_id) {
@@ -278,6 +310,7 @@ void H265Parser::Reset() {
   stream_ = NULL;
   bytes_left_ = 0;
   encrypted_ranges_.clear();
+  previous_nalu_range_.clear();
 }
 
 void H265Parser::SetStream(const uint8_t* stream, off_t stream_size) {
@@ -294,6 +327,7 @@ void H265Parser::SetEncryptedStream(
 
   stream_ = stream;
   bytes_left_ = stream_size;
+  previous_nalu_range_.clear();
 
   encrypted_ranges_.clear();
   const uint8_t* start = stream;
@@ -435,6 +469,8 @@ H265Parser::Result H265Parser::AdvanceToNextNALU(H265NALU* nalu) {
            << " at: " << reinterpret_cast<const void*>(nalu->data)
            << " size: " << nalu->size;
 
+  previous_nalu_range_.clear();
+  previous_nalu_range_.Add(nalu->data, nalu->data + nalu->size);
   return kOk;
 }
 
@@ -1210,6 +1246,16 @@ VideoCodecProfile H265Parser::ProfileIDCToVideoCodecProfile(int profile_idc) {
   }
 }
 
+std::vector<SubsampleEntry> H265Parser::GetCurrentSubsamples() {
+  DCHECK_EQ(previous_nalu_range_.size(), 1u)
+      << "This should only be called after a "
+         "successful call to AdvanceToNextNalu()";
+
+  auto intersection = encrypted_ranges_.IntersectionWith(previous_nalu_range_);
+  return EncryptedRangesToSubsampleEntry(
+      previous_nalu_range_.start(0), previous_nalu_range_.end(0), intersection);
+}
+
 H265Parser::Result H265Parser::ParseProfileTierLevel(
     bool profile_present,
     int max_num_sub_layers_minus1,

media/video/h265_parser.h

@@ -482,6 +482,12 @@ class MEDIA_EXPORT H265Parser {
 
   static VideoCodecProfile ProfileIDCToVideoCodecProfile(int profile_idc);
 
+  // The return value of this method changes for every successful call to
+  // AdvanceToNextNALU().
+  // This returns the subsample information for the last NALU that was output
+  // from AdvanceToNextNALU().
+  std::vector<SubsampleEntry> GetCurrentSubsamples();
+
  private:
   // Move the stream pointer to the beginning of the next NALU,
   // i.e. pointing at the next start code.
@@ -533,6 +539,10 @@ class MEDIA_EXPORT H265Parser {
   // Ranges of encrypted bytes in the buffer passed to SetEncryptedStream().
   Ranges<const uint8_t*> encrypted_ranges_;
 
+  // This contains the range of the previous NALU found in
+  // AdvanceToNextNalu(). Holds exactly one range.
+  Ranges<const uint8_t*> previous_nalu_range_;
+
   DISALLOW_COPY_AND_ASSIGN(H265Parser);
 };
 

media/video/h265_parser_unittest.cc

@@ -8,6 +8,7 @@
 #include "base/files/file_path.h"
 #include "base/files/memory_mapped_file.h"
 #include "base/logging.h"
+#include "media/base/subsample_entry.h"
 #include "media/base/test_data_util.h"
 #include "media/video/h265_parser.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -276,4 +277,113 @@ TEST_F(H265ParserTest, SliceHeaderParsing) {
   EXPECT_TRUE(shdr.slice_loop_filter_across_slices_enabled_flag);
 }
 
+// Verify that GetCurrentSubsamples works.
+TEST_F(H265ParserTest, GetCurrentSubsamplesNormal) {
+  constexpr uint8_t kStream[] = {
+      // First NALU.
+      // Clear bytes = 5.
+      0x00, 0x00, 0x01,  // start code.
+      0x28, 0x00,        // Nalu type = 20, IDR slice.
+      // Below is bogus data.
+      // Encrypted bytes = 15.
+      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x00, 0x01, 0x02, 0x03,
+      0x04, 0x05, 0x06,
+      // Clear bytes = 5.
+      0x07, 0x00, 0x01, 0x02, 0x03,
+      // Encrypted until next NALU. Encrypted bytes = 20.
+      0x04, 0x05, 0x06, 0x07, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+      // Note that this is still in the encrypted region but looks like a start
+      // code.
+      0x00, 0x00, 0x01, 0x03, 0x04, 0x05, 0x06, 0x07,
+      // Second NALU. Completely clear.
+      // Clear bytes = 11.
+      0x00, 0x00, 0x01,  // start code.
+      0x42, 0x00,        // nalu type = 33, SPS.
+      // Bogus data.
+      0xff, 0xfe, 0xfd, 0xee, 0x12, 0x33,
+  };
+  std::vector<SubsampleEntry> subsamples;
+  subsamples.emplace_back(5u, 15u);
+  subsamples.emplace_back(5u, 20u);
+  subsamples.emplace_back(11u, 0u);
+  H265Parser parser;
+  parser.SetEncryptedStream(kStream, base::size(kStream), subsamples);
+
+  H265NALU nalu;
+  EXPECT_EQ(H265Parser::kOk, parser.AdvanceToNextNALU(&nalu));
+  EXPECT_EQ(H265NALU::IDR_N_LP, nalu.nal_unit_type);
+  auto nalu_subsamples = parser.GetCurrentSubsamples();
+  EXPECT_EQ(2u, nalu_subsamples.size());
+
+  // Note that nalu->data starts from the NALU header, i.e. does not include
+  // the start code.
+  EXPECT_EQ(2u, nalu_subsamples[0].clear_bytes);
+  EXPECT_EQ(15u, nalu_subsamples[0].cypher_bytes);
+  EXPECT_EQ(5u, nalu_subsamples[1].clear_bytes);
+  EXPECT_EQ(20u, nalu_subsamples[1].cypher_bytes);
+
+  // Make sure that it reached the next NALU.
+  EXPECT_EQ(H265Parser::kOk, parser.AdvanceToNextNALU(&nalu));
+  EXPECT_EQ(H265NALU::SPS_NUT, nalu.nal_unit_type);
+  nalu_subsamples = parser.GetCurrentSubsamples();
+  EXPECT_EQ(1u, nalu_subsamples.size());
+
+  EXPECT_EQ(8u, nalu_subsamples[0].clear_bytes);
+  EXPECT_EQ(0u, nalu_subsamples[0].cypher_bytes);
+}
+
+// Verify that subsamples starting at non-NALU boundary also works.
+TEST_F(H265ParserTest, GetCurrentSubsamplesSubsampleNotStartingAtNaluBoundary) {
+  constexpr uint8_t kStream[] = {
+      // First NALU.
+      // Clear bytes = 5.
+      0x00, 0x00, 0x01,  // start code.
+      0x28, 0x00,        // Nalu type = 20, IDR slice.
+      // Below is bogus data.
+      // Encrypted bytes = 24.
+      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x00, 0x01, 0x02, 0x03,
+      0x04, 0x05, 0x06, 0x07, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+      // Clear bytes = 19. The rest is in the clear. Note that this is not at
+      // a NALU boundary and a NALU starts below.
+      0xaa, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
+      // Second NALU. Completely clear.
+      0x00, 0x00, 0x01,  // start code.
+      0x42, 0x00,        // nalu type = 33, SPS.
+      // Bogus data.
+      0xff, 0xfe, 0xfd, 0xee, 0x12, 0x33,
+  };
+
+  std::vector<SubsampleEntry> subsamples;
+  subsamples.emplace_back(5u, 24u);
+  subsamples.emplace_back(19u, 0u);
+  H265Parser parser;
+  parser.SetEncryptedStream(kStream, base::size(kStream), subsamples);
+
+  H265NALU nalu;
+  EXPECT_EQ(H265Parser::kOk, parser.AdvanceToNextNALU(&nalu));
+  EXPECT_EQ(H265NALU::IDR_N_LP, nalu.nal_unit_type);
+  auto nalu_subsamples = parser.GetCurrentSubsamples();
+  EXPECT_EQ(2u, nalu_subsamples.size());
+
+  // Note that nalu->data starts from the NALU header, i.e. does not include
+  // the start code.
+  EXPECT_EQ(2u, nalu_subsamples[0].clear_bytes);
+  EXPECT_EQ(24u, nalu_subsamples[0].cypher_bytes);
+
+  // The nalu ends with 8 more clear bytes. The last 10 bytes should be
+  // associated with the next nalu.
+  EXPECT_EQ(8u, nalu_subsamples[1].clear_bytes);
+  EXPECT_EQ(0u, nalu_subsamples[1].cypher_bytes);
+
+  EXPECT_EQ(H265Parser::kOk, parser.AdvanceToNextNALU(&nalu));
+  EXPECT_EQ(H265NALU::SPS_NUT, nalu.nal_unit_type);
+  nalu_subsamples = parser.GetCurrentSubsamples();
+  EXPECT_EQ(1u, nalu_subsamples.size());
+
+  // Although the input had 10 more bytes, since nalu->data starts from the nalu
+  // header, there's only 7 more bytes left.
+  EXPECT_EQ(8u, nalu_subsamples[0].clear_bytes);
+  EXPECT_EQ(0u, nalu_subsamples[0].cypher_bytes);
+}
+
 }  // namespace media