From e422b4eb0efb85e377c593cc6fc37623aa660ef6 Mon Sep 17 00:00:00 2001 From: Rintaro Kuroiwa Date: Wed, 27 Apr 2016 15:04:50 -0700 Subject: [PATCH] Handle encryption in TsSegmenter - Move GetTrackTypeForEncryption() to muxer_util.h. Change-Id: I315957cb8983f4e7c4acff6907dfdd6ad6907c82 --- packager/media/base/muxer_util.cc | 17 ++ packager/media/base/muxer_util.h | 12 ++ packager/media/formats/mp2t/ts_muxer.cc | 4 +- packager/media/formats/mp2t/ts_segmenter.cc | 47 ++++- packager/media/formats/mp2t/ts_segmenter.h | 19 +- .../formats/mp2t/ts_segmenter_unittest.cc | 173 ++++++++++++++++-- packager/media/formats/mp2t/ts_writer.h | 2 +- packager/media/formats/mp4/segmenter.cc | 14 +- packager/media/formats/webm/segmenter.cc | 22 +-- 9 files changed, 265 insertions(+), 45 deletions(-) diff --git a/packager/media/base/muxer_util.cc b/packager/media/base/muxer_util.cc index f93ff80f99..dcdd57423c 100644 --- a/packager/media/base/muxer_util.cc +++ b/packager/media/base/muxer_util.cc @@ -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(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 diff --git a/packager/media/base/muxer_util.h b/packager/media/base/muxer_util.h index ff437c7f03..bbe5a7d363 100644 --- a/packager/media/base/muxer_util.h +++ b/packager/media/base/muxer_util.h @@ -13,9 +13,13 @@ #include +#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 diff --git a/packager/media/formats/mp2t/ts_muxer.cc b/packager/media/formats/mp2t/ts_muxer.cc index 19c2b34efd..2fa1ad4a1e 100644 --- a/packager/media/formats/mp2t/ts_muxer.cc +++ b/packager/media/formats/mp2t/ts_muxer.cc @@ -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; } diff --git a/packager/media/formats/mp2t/ts_segmenter.cc b/packager/media/formats/mp2t/ts_segmenter.cc index 099202b948..17470bf95a 100644 --- a/packager/media/formats/mp2t/ts_segmenter.cc +++ b/packager/media/formats/mp2t/ts_segmenter.cc @@ -8,9 +8,13 @@ #include +#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 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; } diff --git a/packager/media/formats/mp2t/ts_segmenter.h b/packager/media/formats/mp2t/ts_segmenter.h index 5a13562db3..167eb1d201 100644 --- a/packager/media/formats/mp2t/ts_segmenter.h +++ b/packager/media/formats/mp2t/ts_segmenter.h @@ -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 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); }; diff --git a/packager/media/formats/mp2t/ts_segmenter_unittest.cc b/packager/media/formats/mp2t/ts_segmenter_unittest.cc index 2535007b58..19ba5cf846 100644 --- a/packager/media/formats/mp2t/ts_segmenter_unittest.cc +++ b/packager/media/formats/mp2t/ts_segmenter_unittest.cc @@ -8,6 +8,7 @@ #include #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 sample)); + MOCK_METHOD1(SetEncryptionKeyMock, bool(EncryptionKey* encryption_key)); + bool SetEncryptionKey(scoped_ptr 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 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 sample1 = MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame); scoped_refptr 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 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 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 sample1 = + MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame); + scoped_refptr 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 diff --git a/packager/media/formats/mp2t/ts_writer.h b/packager/media/formats/mp2t/ts_writer.h index ffcbf46ea9..7bd31673fc 100644 --- a/packager/media/formats/mp2t/ts_writer.h +++ b/packager/media/formats/mp2t/ts_writer.h @@ -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. diff --git a/packager/media/formats/mp4/segmenter.cc b/packager/media/formats/mp4/segmenter.cc index b4d48f75e9..25a6222805 100644 --- a/packager/media/formats/mp4/segmenter.cc +++ b/packager/media/formats/mp4/segmenter.cc @@ -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(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, diff --git a/packager/media/formats/webm/segmenter.cc b/packager/media/formats/webm/segmenter.cc index 72c219de3a..2d5ff28eb8 100644 --- a/packager/media/formats/webm/segmenter.cc +++ b/packager/media/formats/webm/segmenter.cc @@ -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(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) {