Mirror-Tools
/
shaka-packager
mirror of https://github.com/shaka-project/shaka-packager.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Handle encryption in TsSegmenter 

- Move GetTrackTypeForEncryption() to muxer_util.h.

Change-Id: I315957cb8983f4e7c4acff6907dfdd6ad6907c82
This commit is contained in:
Rintaro Kuroiwa 2016-04-27 15:04:50 -07:00
parent 0ba35147c8
commit e422b4eb0e
9 changed files with 265 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
packager/media/base/muxer_util.cc
View File

@ -15,6 +15,7 @@
#include "packager/base/strings/string_number_conversions.h"
#include "packager/base/strings/string_split.h"
#include "packager/base/strings/stringprintf.h"
#include "packager/media/base/video_stream_info.h"
namespace edash_packager {
namespace {
@ -153,5 +154,21 @@ std::string GetSegmentName(const std::string& segment_template,
return segment_name;
}
KeySource::TrackType GetTrackTypeForEncryption(const StreamInfo& stream_info,
uint32_t max_sd_pixels) {
if (stream_info.stream_type() == kStreamAudio)
return KeySource::TRACK_TYPE_AUDIO;
if (stream_info.stream_type() != kStreamVideo)
return KeySource::TRACK_TYPE_UNKNOWN;
DCHECK_EQ(kStreamVideo, stream_info.stream_type());
const VideoStreamInfo& video_stream_info =
static_cast<const VideoStreamInfo&>(stream_info);
uint32_t pixels = video_stream_info.width() * video_stream_info.height();
return (pixels > max_sd_pixels) ? KeySource::TRACK_TYPE_HD
: KeySource::TRACK_TYPE_SD;
}
} // namespace media
} // namespace edash_packager

12
packager/media/base/muxer_util.h
View File

@ -13,9 +13,13 @@
#include <string>
#include "packager/media/base/key_source.h"
namespace edash_packager {
namespace media {
class StreamInfo;
/// Validates the segment template against segment URL construction rule
/// specified in ISO/IEC 23009-1:2012 5.3.9.4.4.
/// @param segment_template is the template to be validated.
@ -35,6 +39,14 @@ std::string GetSegmentName(const std::string& segment_template,
uint32_t segment_index,
uint32_t bandwidth);
/// Determine the track type for encryption from input.
/// @param stream_info is the info of the stream.
/// @param max_sd_pixels is the maximum number of pixels to be considered SD.
/// Anything above is HD.
/// @return track type for encryption.
KeySource::TrackType GetTrackTypeForEncryption(const StreamInfo& stream_info,
uint32_t max_sd_pixels);
} // namespace media
} // namespace edash_packager

4
packager/media/formats/mp2t/ts_muxer.cc
View File

@ -22,7 +22,9 @@ Status TsMuxer::Initialize() {
return Status(error::MUXER_FAILURE, "Cannot handle more than one streams.");
segmenter_.reset(new TsSegmenter(options(), muxer_listener()));
Status status = segmenter_->Initialize(*streams()[0]->info());
Status status =
segmenter_->Initialize(*streams()[0]->info(), encryption_key_source(),
max_sd_pixels(), clear_lead_in_seconds());
FireOnMediaStartEvent();
return status;
}

47
packager/media/formats/mp2t/ts_segmenter.cc
View File

@ -8,9 +8,13 @@
#include <memory>
#include "packager/media/base/aes_encryptor.h"
#include "packager/media/base/key_source.h"
#include "packager/media/base/muxer_util.h"
#include "packager/media/base/status.h"
#include "packager/media/base/video_stream_info.h"
#include "packager/media/event/muxer_listener.h"
#include "packager/media/event/progress_listener.h"
namespace edash_packager {
namespace media {
@ -27,7 +31,10 @@ TsSegmenter::TsSegmenter(const MuxerOptions& options, MuxerListener* listener)
pes_packet_generator_(new PesPacketGenerator()) {}
TsSegmenter::~TsSegmenter() {}
Status TsSegmenter::Initialize(const StreamInfo& stream_info) {
Status TsSegmenter::Initialize(const StreamInfo& stream_info,
KeySource* encryption_key_source,
uint32_t max_sd_pixels,
double clear_lead_in_seconds) {
if (muxer_options_.segment_template.empty())
return Status(error::MUXER_FAILURE, "Segment template not specified.");
if (!ts_writer_->Initialize(stream_info, false))
@ -37,6 +44,26 @@ Status TsSegmenter::Initialize(const StreamInfo& stream_info) {
"Failed to initialize PesPacketGenerator.");
}
if (encryption_key_source) {
scoped_ptr<EncryptionKey> encryption_key(new EncryptionKey());
const KeySource::TrackType type =
GetTrackTypeForEncryption(stream_info, max_sd_pixels);
Status status = encryption_key_source->GetKey(type, encryption_key.get());
if (encryption_key->iv.empty()) {
if (!AesCryptor::GenerateRandomIv(FOURCC_cbcs, &encryption_key->iv)) {
return Status(error::INTERNAL_ERROR, "Failed to generate random iv.");
}
}
if (!status.ok())
return status;
encryption_key_ = encryption_key.Pass();
clear_lead_in_seconds_ = clear_lead_in_seconds;
status = NotifyEncrypted();
if (!status.ok())
return status;
}
timescale_scale_ = kTsTimescale / stream_info.time_scale();
return Status::OK;
}
@ -137,10 +164,28 @@ Status TsSegmenter::Flush() {
current_segment_total_sample_duration_ * kTsTimescale, file_size);
}
ts_writer_file_opened_ = false;
total_duration_in_seconds_ += current_segment_total_sample_duration_;
}
current_segment_total_sample_duration_ = 0.0;
current_segment_start_time_ = 0;
current_segment_path_.clear();
return NotifyEncrypted();
}
Status TsSegmenter::NotifyEncrypted() {
if (encryption_key_ && total_duration_in_seconds_ >= clear_lead_in_seconds_) {
if (listener_) {
// For now this only happens once, so send true.
const bool kIsInitialEncryptionInfo = true;
listener_->OnEncryptionInfoReady(
kIsInitialEncryptionInfo, FOURCC_cbcs, encryption_key_->key_id,
encryption_key_->iv, encryption_key_->key_system_info);
}
if (!pes_packet_generator_->SetEncryptionKey(encryption_key_.Pass()))
return Status(error::INTERNAL_ERROR, "Failed to set encryption key.");
ts_writer_->SignalEncypted();
}
return Status::OK;
}

19
packager/media/formats/mp2t/ts_segmenter.h
View File

@ -18,6 +18,7 @@
namespace edash_packager {
namespace media {
class KeySource;
class MuxerListener;
namespace mp2t {
@ -27,6 +28,7 @@ namespace mp2t {
// segmenters.
class TsSegmenter {
public:
// TODO(rkuroiwa): Add progress listener?
/// @param options is the options for this muxer. This must stay valid
/// throughout the life time of the instance.
/// @param listener is the MuxerListener that should be used to notify events.
@ -35,9 +37,13 @@ class TsSegmenter {
~TsSegmenter();
/// Initialize the object.
/// Key rotation is not supported.
/// @param stream_info is the stream info for the segmenter.
/// @return OK on success.
Status Initialize(const StreamInfo& stream_info);
Status Initialize(const StreamInfo& stream_info,
KeySource* encryption_key_source,
uint32_t max_sd_pixels,
double clear_lead_in_seconds);
/// Finalize the segmenter.
/// @return OK on success.
@ -69,6 +75,9 @@ class TsSegmenter {
// before calling this, this will open one and write them to file.
Status Flush();
// If conditions are met, notify objects that the data is encrypted.
Status NotifyEncrypted();
const MuxerOptions& muxer_options_;
MuxerListener* const listener_;
@ -97,6 +106,14 @@ class TsSegmenter {
// the segment has been finalized.
std::string current_segment_path_;
scoped_ptr<EncryptionKey> encryption_key_;
double clear_lead_in_seconds_ = 0;
// The total duration of the segments that it has segmented. This only
// includes segments that have been finailzed. IOW, this does not count the
// current segments duration.
double total_duration_in_seconds_ = 0.0;
DISALLOW_COPY_AND_ASSIGN(TsSegmenter);
};

173
packager/media/formats/mp2t/ts_segmenter_unittest.cc
View File

@ -8,6 +8,7 @@
#include <gtest/gtest.h>
#include "packager/media/base/audio_stream_info.h"
#include "packager/media/base/fixed_key_source.h"
#include "packager/media/base/test/status_test_util.h"
#include "packager/media/base/video_stream_info.h"
#include "packager/media/event/mock_muxer_listener.h"
@ -45,10 +46,19 @@ const uint16_t kTrickPlayRate = 1;
const uint8_t kNaluLengthSize = 1;
const bool kIsEncrypted = false;
const uint8_t kAnyData[] = {
0x01, 0x0F, 0x3C,
};
class MockPesPacketGenerator : public PesPacketGenerator {
public:
MOCK_METHOD1(Initialize, bool(const StreamInfo& info));
MOCK_METHOD1(PushSample, bool(scoped_refptr<MediaSample> sample));
MOCK_METHOD1(SetEncryptionKeyMock, bool(EncryptionKey* encryption_key));
bool SetEncryptionKey(scoped_ptr<EncryptionKey> encryption_key) override {
return SetEncryptionKeyMock(encryption_key.get());
}
MOCK_METHOD0(NumberOfReadyPesPackets, size_t());
// Hack found at the URL below for mocking methods that return scoped_ptr.
@ -66,6 +76,7 @@ class MockTsWriter : public TsWriter {
MOCK_METHOD2(Initialize,
bool(const StreamInfo& stream_info, bool will_be_encrypted));
MOCK_METHOD1(NewSegment, bool(const std::string& file_name));
MOCK_METHOD0(SignalEncypted, void());
MOCK_METHOD0(FinalizeSegment, bool());
// Similar to the hack above but takes a scoped_ptr.
@ -76,6 +87,12 @@ class MockTsWriter : public TsWriter {
}
};
// TODO(rkuroiwa): Add mock_key_source.{h,cc} in media/base.
class MockKeySource : public FixedKeySource {
public:
MOCK_METHOD2(GetKey, Status(TrackType track_type, EncryptionKey* key));
};
} // namespace
class TsSegmenterTest : public ::testing::Test {
@ -106,7 +123,7 @@ TEST_F(TsSegmenterTest, Initialize) {
segmenter.InjectPesPacketGeneratorForTesting(
mock_pes_packet_generator_.Pass());
EXPECT_OK(segmenter.Initialize(*stream_info));
EXPECT_OK(segmenter.Initialize(*stream_info, nullptr, 0, 0));
}
TEST_F(TsSegmenterTest, AddSample) {
@ -123,9 +140,6 @@ TEST_F(TsSegmenterTest, AddSample) {
EXPECT_CALL(*mock_pes_packet_generator_, Initialize(_))
.WillOnce(Return(true));
const uint8_t kAnyData[] = {
0x01, 0x0F, 0x3C,
};
scoped_refptr<MediaSample> sample =
MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame);
@ -156,7 +170,7 @@ TEST_F(TsSegmenterTest, AddSample) {
segmenter.InjectPesPacketGeneratorForTesting(
mock_pes_packet_generator_.Pass());
EXPECT_OK(segmenter.Initialize(*stream_info));
EXPECT_OK(segmenter.Initialize(*stream_info, nullptr, 0, 0));
EXPECT_OK(segmenter.AddSample(sample));
}
@ -185,9 +199,6 @@ TEST_F(TsSegmenterTest, PassedSegmentDuration) {
EXPECT_CALL(*mock_pes_packet_generator_, Initialize(_))
.WillOnce(Return(true));
const uint8_t kAnyData[] = {
0x01, 0x0F, 0x3C,
};
scoped_refptr<MediaSample> sample1 =
MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame);
scoped_refptr<MediaSample> sample2 =
@ -263,7 +274,7 @@ TEST_F(TsSegmenterTest, PassedSegmentDuration) {
segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
segmenter.InjectPesPacketGeneratorForTesting(
mock_pes_packet_generator_.Pass());
EXPECT_OK(segmenter.Initialize(*stream_info));
EXPECT_OK(segmenter.Initialize(*stream_info, nullptr, 0, 0));
EXPECT_OK(segmenter.AddSample(sample1));
EXPECT_OK(segmenter.AddSample(sample2));
}
@ -290,7 +301,7 @@ TEST_F(TsSegmenterTest, InitializeThenFinalize) {
segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
segmenter.InjectPesPacketGeneratorForTesting(
mock_pes_packet_generator_.Pass());
EXPECT_OK(segmenter.Initialize(*stream_info));
EXPECT_OK(segmenter.Initialize(*stream_info, nullptr, 0, 0));
EXPECT_OK(segmenter.Finalize());
}
@ -321,7 +332,7 @@ TEST_F(TsSegmenterTest, Finalize) {
segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
segmenter.InjectPesPacketGeneratorForTesting(
mock_pes_packet_generator_.Pass());
EXPECT_OK(segmenter.Initialize(*stream_info));
EXPECT_OK(segmenter.Initialize(*stream_info, nullptr, 0, 0));
segmenter.SetTsWriterFileOpenedForTesting(true);
EXPECT_OK(segmenter.Finalize());
}
@ -427,12 +438,150 @@ TEST_F(TsSegmenterTest, SegmentOnlyBeforeKeyFrame) {
segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
segmenter.InjectPesPacketGeneratorForTesting(
mock_pes_packet_generator_.Pass());
EXPECT_OK(segmenter.Initialize(*stream_info));
EXPECT_OK(segmenter.Initialize(*stream_info, nullptr, 0, 0));
EXPECT_OK(segmenter.AddSample(key_frame_sample1));
EXPECT_OK(segmenter.AddSample(non_key_frame_sample));
EXPECT_OK(segmenter.AddSample(key_frame_sample2));
}
TEST_F(TsSegmenterTest, WithEncryptionNoClearLead) {
scoped_refptr<VideoStreamInfo> stream_info(new VideoStreamInfo(
kTrackId, kTimeScale, kDuration, kH264VideoCodec, kCodecString, kLanguage,
kWidth, kHeight, kPixelWidth, kPixelHeight, kTrickPlayRate,
kNaluLengthSize, kExtraData, arraysize(kExtraData), kIsEncrypted));
MuxerOptions options;
options.segment_duration = 10.0;
options.segment_template = "file$Number$.ts";
MockMuxerListener mock_listener;
EXPECT_CALL(mock_listener, OnEncryptionInfoReady(_, _, _, _, _));
TsSegmenter segmenter(options, &mock_listener);
EXPECT_CALL(*mock_ts_writer_, Initialize(_, _)).WillOnce(Return(true));
EXPECT_CALL(*mock_ts_writer_, SignalEncypted());
EXPECT_CALL(*mock_pes_packet_generator_, Initialize(_))
.WillOnce(Return(true));
EXPECT_CALL(*mock_pes_packet_generator_, SetEncryptionKeyMock(_))
.WillOnce(Return(true));
segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
segmenter.InjectPesPacketGeneratorForTesting(
mock_pes_packet_generator_.Pass());
MockKeySource mock_key_source;
EXPECT_CALL(mock_key_source, GetKey(KeySource::TRACK_TYPE_HD, _))
.WillOnce(Return(Status::OK));
const uint32_t k480pPixels = 640 * 480;
// Set this to 0 so that Finalize will call
// PesPacketGenerator::SetEncryptionKey().
// Even tho no samples have been added.
const double kClearLeadSeconds = 0;
EXPECT_OK(segmenter.Initialize(*stream_info, &mock_key_source, k480pPixels,
kClearLeadSeconds));
}
// Verify that encryption notification is sent to objects after clear lead.
TEST_F(TsSegmenterTest, WithEncryptionWithClearLead) {
scoped_refptr<VideoStreamInfo> stream_info(new VideoStreamInfo(
kTrackId, kTimeScale, kDuration, kH264VideoCodec, kCodecString,
kLanguage, kWidth, kHeight, kPixelWidth, kPixelHeight, kTrickPlayRate,
kNaluLengthSize, kExtraData, arraysize(kExtraData), kIsEncrypted));
MuxerOptions options;
options.segment_duration = 1.0;
const double kClearLeadSeconds = 1.0;
options.segment_template = "file$Number$.ts";
MockMuxerListener mock_listener;
TsSegmenter segmenter(options, &mock_listener);
ON_CALL(*mock_ts_writer_, Initialize(_, _)).WillByDefault(Return(true));
ON_CALL(*mock_ts_writer_, NewSegment(_)).WillByDefault(Return(true));
ON_CALL(*mock_ts_writer_, FinalizeSegment()).WillByDefault(Return(true));
ON_CALL(*mock_ts_writer_, AddPesPacketMock(_)).WillByDefault(Return(true));
ON_CALL(*mock_pes_packet_generator_, Initialize(_))
.WillByDefault(Return(true));
ON_CALL(*mock_pes_packet_generator_, Flush()).WillByDefault(Return(true));
const uint8_t kAnyData[] = {
0x01, 0x0F, 0x3C,
};
scoped_refptr<MediaSample> sample1 =
MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame);
scoped_refptr<MediaSample> sample2 =
MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame);
// Something longer than 1.0 (segment duration and clear lead).
sample1->set_duration(kTimeScale * 2);
// The length of the second sample doesn't really matter.
sample2->set_duration(kTimeScale * 3);
EXPECT_CALL(*mock_pes_packet_generator_, PushSample(_))
.Times(2)
.WillRepeatedly(Return(true));
Sequence ready_pes_sequence;
// First AddSample().
EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
.InSequence(ready_pes_sequence)
.WillOnce(Return(1u));
EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
.InSequence(ready_pes_sequence)
.WillOnce(Return(0u));
// When Flush() is called, inside second AddSample().
EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
.InSequence(ready_pes_sequence)
.WillOnce(Return(0u));
// Still inside AddSample() but after Flush().
EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
.InSequence(ready_pes_sequence)
.WillOnce(Return(1u));
EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
.InSequence(ready_pes_sequence)
.WillOnce(Return(0u));
EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_))
.Times(2)
.WillRepeatedly(Return(true));
// The pointers are released inside the segmenter.
Sequence pes_packet_sequence;
EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
.InSequence(pes_packet_sequence)
.WillOnce(Return(new PesPacket()));
EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
.InSequence(pes_packet_sequence)
.WillOnce(Return(new PesPacket()));
MockPesPacketGenerator* mock_pes_packet_generator_raw =
mock_pes_packet_generator_.get();
MockTsWriter* mock_ts_writer_raw = mock_ts_writer_.get();
segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
segmenter.InjectPesPacketGeneratorForTesting(
mock_pes_packet_generator_.Pass());
MockKeySource mock_key_source;
// This should be called AFTER the first AddSample().
EXPECT_CALL(mock_key_source, GetKey(KeySource::TRACK_TYPE_HD, _))
.WillOnce(Return(Status::OK));
EXPECT_OK(segmenter.Initialize(*stream_info, &mock_key_source, 0,
kClearLeadSeconds));
EXPECT_OK(segmenter.AddSample(sample1));
// These should be called AFTER the first AddSample(), before the second
// segment.
EXPECT_CALL(mock_listener, OnEncryptionInfoReady(_, _, _, _, _));
EXPECT_CALL(*mock_pes_packet_generator_raw, SetEncryptionKeyMock(_))
.WillOnce(Return(true));
EXPECT_CALL(*mock_ts_writer_raw, SignalEncypted());
EXPECT_OK(segmenter.AddSample(sample2));
}
} // namespace mp2t
} // namespace media
} // namespace edash_packager

2
packager/media/formats/mp2t/ts_writer.h
View File

@ -47,7 +47,7 @@ class TsWriter {
/// Signals the writer that the rest of the segments are encrypted.
/// |will_be_encrypted| passed to Initialize() should be true.
void SignalEncypted();
virtual void SignalEncypted();
/// Flush all the pending PesPackets that have not been written to file and
/// close the file.

14
packager/media/formats/mp4/segmenter.cc
View File

@ -16,6 +16,7 @@
#include "packager/media/base/media_sample.h"
#include "packager/media/base/media_stream.h"
#include "packager/media/base/muxer_options.h"
#include "packager/media/base/muxer_util.h"
#include "packager/media/base/video_stream_info.h"
#include "packager/media/event/muxer_listener.h"
#include "packager/media/event/progress_listener.h"
@ -122,19 +123,6 @@ void GenerateEncryptedSampleEntry(const EncryptionKey& encryption_key,
}
}
KeySource::TrackType GetTrackTypeForEncryption(const StreamInfo& stream_info,
uint32_t max_sd_pixels) {
if (stream_info.stream_type() == kStreamAudio)
return KeySource::TRACK_TYPE_AUDIO;
DCHECK_EQ(kStreamVideo, stream_info.stream_type());
const VideoStreamInfo& video_stream_info =
static_cast<const VideoStreamInfo&>(stream_info);
uint32_t pixels = video_stream_info.width() * video_stream_info.height();
return (pixels > max_sd_pixels) ? KeySource::TRACK_TYPE_HD
: KeySource::TRACK_TYPE_SD;
}
} // namespace
Segmenter::Segmenter(const MuxerOptions& options,

22
packager/media/formats/webm/segmenter.cc
View File

@ -11,6 +11,7 @@
#include "packager/media/base/media_sample.h"
#include "packager/media/base/media_stream.h"
#include "packager/media/base/muxer_options.h"
#include "packager/media/base/muxer_util.h"
#include "packager/media/base/stream_info.h"
#include "packager/media/base/video_stream_info.h"
#include "packager/media/event/muxer_listener.h"
@ -334,22 +335,11 @@ Status Segmenter::CreateAudioTrack(AudioStreamInfo* info) {
Status Segmenter::InitializeEncryptor(KeySource* key_source,
uint32_t max_sd_pixels) {
encryptor_.reset(new Encryptor());
switch (info_->stream_type()) {
case kStreamVideo: {
VideoStreamInfo* video_info = static_cast<VideoStreamInfo*>(info_);
uint32_t pixels = video_info->width() * video_info->height();
KeySource::TrackType type = (pixels > max_sd_pixels)
? KeySource::TRACK_TYPE_HD
: KeySource::TRACK_TYPE_SD;
return encryptor_->Initialize(muxer_listener_, type, key_source);
}
case kStreamAudio:
return encryptor_->Initialize(
muxer_listener_, KeySource::TrackType::TRACK_TYPE_AUDIO, key_source);
default:
// Other streams are not encrypted.
return Status::OK;
}
const KeySource::TrackType track_type =
GetTrackTypeForEncryption(*info_, max_sd_pixels);
if (track_type == KeySource::TrackType::TRACK_TYPE_UNKNOWN)
return Status::OK;
return encryptor_->Initialize(muxer_listener_, track_type, key_source);
}
Status Segmenter::WriteFrame(bool write_duration) {