Make NaluReader skip encrypted portion

Change-Id: Ibb47a1e62cd8ac3057c8f1512a88280991e48b62
This commit is contained in:
Rintaro Kuroiwa 2017-05-19 15:21:02 -07:00
parent a678948db4
commit e2401f02ec
3 changed files with 394 additions and 5 deletions

View File

@ -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_,
&nalu_start_off, &annexb_start_code_size)) {
if (!FindStartCodeInClearRange(
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,
&nalu_size_without_start_code, &next_start_code_size)) {
if (!FindStartCodeInClearRange(
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;

View File

@ -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);
};

View File

@ -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