From e7c0dbf09db013a56c1ef9bd2042717d4262e942 Mon Sep 17 00:00:00 2001
From: Rintaro Kuroiwa <rkuroiwa@google.com>
Date: Tue, 8 Nov 2016 14:58:13 -0800
Subject: [PATCH] Ensure some method calls NotifyEncryptionStart()

- If OnEncryptionStart() is called OnMediaStart(), someone should be
  calling NotifyEncrytionStart().

Fixes #169

Change-Id: I597f317d2ddeb3091b7e9d238eb9695b4cc9c69d
---
 .../media/event/hls_notify_muxer_listener.cc  | 15 ++++++-
 .../media/event/hls_notify_muxer_listener.h   |  1 +
 .../hls_notify_muxer_listener_unittest.cc     | 42 ++++++++++++++++++-
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/packager/media/event/hls_notify_muxer_listener.cc b/packager/media/event/hls_notify_muxer_listener.cc
index df89279dc6..a7b0d9bf7b 100644
--- a/packager/media/event/hls_notify_muxer_listener.cc
+++ b/packager/media/event/hls_notify_muxer_listener.cc
@@ -30,6 +30,15 @@ HlsNotifyMuxerListener::HlsNotifyMuxerListener(
 
 HlsNotifyMuxerListener::~HlsNotifyMuxerListener() {}
 
+// These methods work together to notify that the media is encrypted.
+// If OnEncryptionInfoReady() is called before the media has been started, then
+// the information is stored and handled when OnEncryptionStart() is called.
+// if OnEncryptionStart() is called before the media has been started then
+// OnMediaStart() is responsible for notifying that the segments are encrypted
+// right away i.e. call OnEncryptionStart().
+// For now (because Live HLS is not implemented yet) this should be called once,
+// before media is started. So the logic after the first if statement should not
+// be taken.
 void HlsNotifyMuxerListener::OnEncryptionInfoReady(
     bool is_initial_encryption_info,
     FourCC protection_scheme,
@@ -51,7 +60,7 @@ void HlsNotifyMuxerListener::OnEncryptionInfoReady(
 
 void HlsNotifyMuxerListener::OnEncryptionStart() {
   if (!media_started_) {
-    DLOG(WARNING) << "Media not started, cannot notify encryption start.";
+    must_notify_encryption_start_ = true;
     return;
   }
   if (next_key_id_.empty()) {
@@ -69,6 +78,7 @@ void HlsNotifyMuxerListener::OnEncryptionStart() {
   next_key_id_.clear();
   next_iv_.clear();
   next_key_system_infos_.clear();
+  must_notify_encryption_start_ = false;
 }
 
 void HlsNotifyMuxerListener::OnMediaStart(const MuxerOptions& muxer_options,
@@ -90,6 +100,9 @@ void HlsNotifyMuxerListener::OnMediaStart(const MuxerOptions& muxer_options,
   }
 
   media_started_ = true;
+  if (must_notify_encryption_start_) {
+    OnEncryptionStart();
+  }
 }
 
 void HlsNotifyMuxerListener::OnSampleDurationReady(uint32_t sample_duration) {}
diff --git a/packager/media/event/hls_notify_muxer_listener.h b/packager/media/event/hls_notify_muxer_listener.h
index b402ea144f..e1cb61e947 100644
--- a/packager/media/event/hls_notify_muxer_listener.h
+++ b/packager/media/event/hls_notify_muxer_listener.h
@@ -73,6 +73,7 @@ class HlsNotifyMuxerListener : public MuxerListener {
   uint32_t stream_id_ = 0;
 
   bool media_started_ = false;
+  bool must_notify_encryption_start_ = false;
   // Cached encryption info before OnMediaStart() is called.
   std::vector<uint8_t> next_key_id_;
   std::vector<uint8_t> next_iv_;
diff --git a/packager/media/event/hls_notify_muxer_listener_unittest.cc b/packager/media/event/hls_notify_muxer_listener_unittest.cc
index be8739db87..88dde1d79a 100644
--- a/packager/media/event/hls_notify_muxer_listener_unittest.cc
+++ b/packager/media/event/hls_notify_muxer_listener_unittest.cc
@@ -119,7 +119,7 @@ TEST_F(HlsNotifyMuxerListenerTest, OnMediaStart) {
                          MuxerListener::kContainerMpeg2ts);
 }
 
-// OnEncryptionStart() should NotifyEncryptionUpdate() after
+// OnEncryptionStart() should call MuxerListener::NotifyEncryptionUpdate() after
 // OnEncryptionInfoReady() and OnMediaStart().
 TEST_F(HlsNotifyMuxerListenerTest, OnEncryptionStart) {
   ProtectionSystemSpecificInfo info;
@@ -158,6 +158,46 @@ TEST_F(HlsNotifyMuxerListenerTest, OnEncryptionStart) {
   listener_.OnEncryptionStart();
 }
 
+// If OnEncryptionStart() is called before media start,
+// HlsNotiifer::NotifyEncryptionUpdate() should be called by the end of
+// OnMediaStart().
+TEST_F(HlsNotifyMuxerListenerTest, OnEncryptionStartBeforeMediaStart) {
+  ProtectionSystemSpecificInfo info;
+  std::vector<uint8_t> system_id(kAnySystemId,
+                                 kAnySystemId + arraysize(kAnySystemId));
+  info.set_system_id(system_id.data(), system_id.size());
+  std::vector<uint8_t> pssh_data(kAnyData, kAnyData + arraysize(kAnyData));
+  info.set_pssh_data(pssh_data);
+
+  std::vector<uint8_t> key_id(16, 0x05);
+  std::vector<ProtectionSystemSpecificInfo> key_system_infos;
+  key_system_infos.push_back(info);
+
+  std::vector<uint8_t> iv(16, 0x54);
+
+  EXPECT_CALL(mock_notifier_, NotifyEncryptionUpdate(_, _, _, _, _)).Times(0);
+  listener_.OnEncryptionInfoReady(kInitialEncryptionInfo, FOURCC_cbcs, key_id,
+                                  iv, key_system_infos);
+  ::testing::Mock::VerifyAndClearExpectations(&mock_notifier_);
+
+  ON_CALL(mock_notifier_, NotifyNewStream(_, _, _, _, _))
+      .WillByDefault(Return(true));
+  VideoStreamInfoParameters video_params = GetDefaultVideoStreamInfoParams();
+  scoped_refptr<StreamInfo> video_stream_info =
+      CreateVideoStreamInfo(video_params);
+  MuxerOptions muxer_options;
+
+  // It doesn't really matter when this is called, could be called right away in
+  // OnEncryptionStart() if that is possible. Just matters that it is called by
+  // the time OnMediaStart() returns.
+  EXPECT_CALL(mock_notifier_,
+              NotifyEncryptionUpdate(_, key_id, system_id, iv, pssh_data))
+      .WillOnce(Return(true));
+  listener_.OnEncryptionStart();
+  listener_.OnMediaStart(muxer_options, *video_stream_info, 90000,
+                         MuxerListener::kContainerMpeg2ts);
+}
+
 // NotifyEncryptionUpdate() should not be called if NotifyNewStream() fails in
 // OnMediaStart().
 TEST_F(HlsNotifyMuxerListenerTest, NoEncryptionUpdateIfNotifyNewStreamFails) {