Move PSSH generation to EncryptionHandler.

The KeySource now only handles fetching the keys and loading any PSSH
info from the license; it will not handle generating new PSSH info
based on the config.

This will allow the PSSH generation to access to the full
EncryptionConfig so we can add additional options to it.

Issue #756

Change-Id: Ia67387aa3d5ec0d723b7f5f21fc517f64c840393

packager/app/packager_util.cc

@@ -85,9 +85,7 @@ std::unique_ptr<KeySource> CreateEncryptionKeySource(
       break;
     }
     case KeyProvider::kRawKey: {
-      encryption_key_source = RawKeySource::Create(
+      encryption_key_source = RawKeySource::Create(encryption_params.raw_key);
-          encryption_params.raw_key, encryption_params.protection_systems,
-          protection_scheme);
       break;
     }
     case KeyProvider::kPlayReady: {
@@ -114,11 +112,10 @@ std::unique_ptr<KeySource> CreateEncryptionKeySource(
               playready.key_server_url, playready.client_cert_file,
               playready.client_cert_private_key_file,
               playready.client_cert_private_key_password,
-              encryption_params.protection_systems, protection_scheme));
+              encryption_params.protection_systems));
         } else {
           playready_key_source.reset(new PlayReadyKeySource(
-              playready.key_server_url, encryption_params.protection_systems,
+              playready.key_server_url, encryption_params.protection_systems));
-              protection_scheme));
         }
         if (!playready.ca_file.empty()) {
           playready_key_source->SetCaFile(playready.ca_file);
@@ -170,10 +167,7 @@ std::unique_ptr<KeySource> CreateDecryptionKeySource(
       break;
     }
     case KeyProvider::kRawKey: {
-      decryption_key_source = RawKeySource::Create(
+      decryption_key_source = RawKeySource::Create(decryption_params.raw_key);
-          decryption_params.raw_key,
-          ProtectionSystem::kCommon /* value does not matter here */,
-          FOURCC_NULL /* value does not matter here */);
       break;
     }
     default:

packager/app/test/packager_test.py

@@ -282,6 +282,8 @@ class PackagerAppTest(unittest.TestCase):
     else:
       self.encryption_iv = '3334353637383930'
     self.widevine_content_id = '3031323334353637'
+    self.pssh = ('0000002070737368000000001077efecc0b24d02ace33c1e52e2fb4b000'
+                 '00000')
     # TS files may have a non-zero start, which could result in the first
     # segment to be less than 1 second. Set clear_lead to be less than 1
     # so only the first segment is left in clear.
@@ -1057,6 +1059,13 @@ class PackagerFunctionalTest(PackagerAppTest):
     self.assertPackageSuccess(streams, flags)
     self._CheckTestResults('encryption-multi-keys-with-stream-label')
 
+  def testExplicitPssh(self):
+    flags = self._GetFlags(encryption=True, output_dash=True) + [
+        '--pssh={0}'.format(self.pssh),
+    ]
+    self.assertPackageSuccess(self._GetStreams(['audio', 'video']), flags)
+    self._CheckTestResults('encryption-using-explicit-pssh')
+
   def testEncryptionOfOnlyVideoStream(self):
     streams = [
         self._GetStream('audio', skip_encryption=True),

packager/app/test/testdata/encryption-using-explicit-pssh/output.mpd

@@ -0,0 +1,31 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--Generated with https://github.com/google/shaka-packager version <tag>-<hash>-<test>-->
+<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="urn:mpeg:dash:schema:mpd:2011 DASH-MPD.xsd" xmlns:cenc="urn:mpeg:cenc:2013" profiles="urn:mpeg:dash:profile:isoff-on-demand:2011" minBufferTime="PT2S" type="static" mediaPresentationDuration="PT2.7360665798187256S">
+  <Period id="0">
+    <AdaptationSet id="0" contentType="video" width="640" height="360" frameRate="30000/1001" subsegmentAlignment="true" par="16:9">
+      <ContentProtection value="cenc" schemeIdUri="urn:mpeg:dash:mp4protection:2011" cenc:default_KID="31323334-3536-3738-3930-313233343536"/>
+      <ContentProtection schemeIdUri="urn:uuid:1077efec-c0b2-4d02-ace3-3c1e52e2fb4b">
+        <cenc:pssh>AAAAIHBzc2gAAAAAEHfv7MCyTQKs4zweUuL7SwAAAAA=</cenc:pssh>
+      </ContentProtection>
+      <Representation id="0" bandwidth="977743" codecs="avc1.64001e" mimeType="video/mp4" sar="1:1">
+        <BaseURL>bear-640x360-video.mp4</BaseURL>
+        <SegmentBase indexRange="1107-1174" timescale="30000">
+          <Initialization range="0-1106"/>
+        </SegmentBase>
+      </Representation>
+    </AdaptationSet>
+    <AdaptationSet id="1" contentType="audio" subsegmentAlignment="true">
+      <ContentProtection value="cenc" schemeIdUri="urn:mpeg:dash:mp4protection:2011" cenc:default_KID="31323334-3536-3738-3930-313233343536"/>
+      <ContentProtection schemeIdUri="urn:uuid:1077efec-c0b2-4d02-ace3-3c1e52e2fb4b">
+        <cenc:pssh>AAAAIHBzc2gAAAAAEHfv7MCyTQKs4zweUuL7SwAAAAA=</cenc:pssh>
+      </ContentProtection>
+      <Representation id="1" bandwidth="133663" codecs="mp4a.40.2" mimeType="audio/mp4" audioSamplingRate="44100">
+        <AudioChannelConfiguration schemeIdUri="urn:mpeg:dash:23003:3:audio_channel_configuration:2011" value="2"/>
+        <BaseURL>bear-640x360-audio.mp4</BaseURL>
+        <SegmentBase indexRange="983-1050" timescale="44100">
+          <Initialization range="0-982"/>
+        </SegmentBase>
+      </Representation>
+    </AdaptationSet>
+  </Period>
+</MPD>

packager/media/base/key_source.cc

@@ -6,80 +6,12 @@
 
 #include "packager/media/base/key_source.h"
 
-#include "packager/base/logging.h"
-#include "packager/media/base/common_pssh_generator.h"
-#include "packager/media/base/playready_pssh_generator.h"
-#include "packager/media/base/protection_system_ids.h"
-#include "packager/media/base/widevine_pssh_generator.h"
-#include "packager/status_macros.h"
-
 namespace shaka {
 namespace media {
 
-KeySource::KeySource(ProtectionSystem protection_systems,
+KeySource::KeySource() = default;
-                     FourCC protection_scheme) {
-  if (has_flag(protection_systems, ProtectionSystem::kCommon)) {
-    pssh_generators_.emplace_back(new CommonPsshGenerator());
-  }
-
-  if (has_flag(protection_systems, ProtectionSystem::kPlayReady)) {
-    pssh_generators_.emplace_back(
-        new PlayReadyPsshGenerator(protection_scheme));
-  }
-
-  if (has_flag(protection_systems, ProtectionSystem::kWidevine)) {
-    pssh_generators_.emplace_back(new WidevinePsshGenerator(protection_scheme));
-  }
-
-  if (has_flag(protection_systems, ProtectionSystem::kFairPlay)) {
-    no_pssh_systems_.emplace_back(std::begin(kFairPlaySystemId),
-                                  std::end(kFairPlaySystemId));
-  }
-  // We only support Marlin Adaptive Streaming Specification – Simple Profile
-  // with Implicit Content ID Mapping, which does not need a PSSH. Marlin
-  // specific PSSH with Explicit Content ID Mapping is not generated.
-  if (has_flag(protection_systems, ProtectionSystem::kMarlin)) {
-    no_pssh_systems_.emplace_back(std::begin(kMarlinSystemId),
-                                  std::end(kMarlinSystemId));
-  }
-}
 
 KeySource::~KeySource() = default;
 
-Status KeySource::UpdateProtectionSystemInfo(
-    EncryptionKeyMap* encryption_key_map) {
-  for (const auto& pssh_generator : pssh_generators_) {
-    const bool support_multiple_keys = pssh_generator->SupportMultipleKeys();
-    if (support_multiple_keys) {
-      ProtectionSystemSpecificInfo info;
-      std::vector<std::vector<uint8_t>> key_ids;
-      for (const EncryptionKeyMap::value_type& pair : *encryption_key_map) {
-        key_ids.push_back(pair.second->key_id);
-      }
-      RETURN_IF_ERROR(pssh_generator->GeneratePsshFromKeyIds(key_ids, &info));
-      for (const EncryptionKeyMap::value_type& pair : *encryption_key_map) {
-        pair.second->key_system_info.push_back(info);
-      }
-    } else {
-      for (const EncryptionKeyMap::value_type& pair : *encryption_key_map) {
-        ProtectionSystemSpecificInfo info;
-        RETURN_IF_ERROR(pssh_generator->GeneratePsshFromKeyIdAndKey(
-            pair.second->key_id, pair.second->key, &info));
-        pair.second->key_system_info.push_back(info);
-      }
-    }
-  }
-
-  for (const auto& no_pssh_system : no_pssh_systems_) {
-    ProtectionSystemSpecificInfo info;
-    info.system_id = no_pssh_system;
-    for (const EncryptionKeyMap::value_type& pair : *encryption_key_map) {
-      pair.second->key_system_info.push_back(info);
-    }
-  }
-
-  return Status::OK;
-}
-
 }  // namespace media
 }  // namespace shaka

packager/media/base/key_source.h

@@ -37,7 +37,10 @@ enum class EmeInitDataType {
 
 struct EncryptionKey {
   std::vector<ProtectionSystemSpecificInfo> key_system_info;
+  /// The ID of this key.
   std::vector<uint8_t> key_id;
+  /// The IDs of the other keys to include in PSSH info.
+  std::vector<std::vector<uint8_t>> key_ids;
   std::vector<uint8_t> key;
   std::vector<uint8_t> iv;
 };
@@ -47,7 +50,7 @@ typedef std::map<std::string, std::unique_ptr<EncryptionKey>> EncryptionKeyMap;
 /// KeySource is responsible for encryption key acquisition.
 class KeySource {
  public:
-  KeySource(ProtectionSystem protection_systems, FourCC protection_scheme);
+  KeySource();
 
   virtual ~KeySource();
 
@@ -88,15 +91,7 @@ class KeySource {
                                     const std::string& stream_label,
                                     EncryptionKey* key) = 0;
 
- protected:
-  /// Update the protection sysmtem specific info for the encryption keys.
-  /// @param encryption_key_map is a map of encryption keys for all tracks.
-  Status UpdateProtectionSystemInfo(EncryptionKeyMap* encryption_key_map);
-
  private:
-  std::vector<std::unique_ptr<PsshGenerator>> pssh_generators_;
-  std::vector<std::vector<uint8_t>> no_pssh_systems_;
-
   DISALLOW_COPY_AND_ASSIGN(KeySource);
 };
 

packager/media/base/media_handler_test_base.h

@@ -43,6 +43,15 @@ bool TryMatch(const T& value,
   return true;
 }
 
+MATCHER_P(IsPsshInfoWithSystemId,
+          system_id,
+          std::string(negation ? "doesn't " : "") + " have system ID " +
+              testing::PrintToString(system_id)) {
+  *result_listener << "which is (" << testing::PrintToString(arg.system_id)
+                   << ")";
+  return arg.system_id == system_id;
+}
+
 MATCHER_P4(IsStreamInfo, stream_index, time_scale, encrypted, language, "") {
   if (!TryMatchStreamDataType(arg->stream_data_type,
                               StreamDataType::kStreamInfo, result_listener)) {

packager/media/base/playready_key_source.cc

@@ -63,12 +63,9 @@ bool Base64StringToBytes(const std::string& base64_string,
 }
 
 PlayReadyKeySource::PlayReadyKeySource(const std::string& server_url,
-                                       ProtectionSystem protection_systems,
+                                       ProtectionSystem protection_systems)
-                                       FourCC protection_scheme)
     // PlayReady PSSH is retrived from PlayReady server response.
-    : KeySource(protection_systems & ~ProtectionSystem::kPlayReady,
+    : generate_playready_protection_system_(
-                protection_scheme),
-      generate_playready_protection_system_(
           // Generate PlayReady protection system if there are no other
           // protection system specified.
           protection_systems == ProtectionSystem::kNone ||
@@ -81,12 +78,9 @@ PlayReadyKeySource::PlayReadyKeySource(
     const std::string& client_cert_file,
     const std::string& client_cert_private_key_file,
     const std::string& client_cert_private_key_password,
-    ProtectionSystem protection_systems,
+    ProtectionSystem protection_systems)
-    FourCC protection_scheme)
     // PlayReady PSSH is retrived from PlayReady server response.
-    : KeySource(protection_systems & ~ProtectionSystem::kPlayReady,
+    : encryption_key_(new EncryptionKey),
-                protection_scheme),
-      encryption_key_(new EncryptionKey),
       server_url_(server_url),
       client_cert_file_(client_cert_file),
       client_cert_private_key_file_(client_cert_private_key_file),
@@ -141,6 +135,7 @@ Status SetKeyInformationFromServerResponse(
     LOG(ERROR) << "Cannot parse key, " << key_data_b64;
     return Status(error::SERVER_ERROR, "Cannot parse key.");
   }
+  encryption_key->key_ids.emplace_back(encryption_key->key_id);
 
   if (generate_playready_protection_system) {
     std::string pssh_data_b64;
@@ -189,11 +184,7 @@ Status PlayReadyKeySource::FetchKeysWithProgramIdentifier(
       encryption_key.get()));
 
   // PlayReady does not specify different streams.
-  const char kEmptyDrmLabel[] = "";
+  encryption_key_ = std::move(encryption_key);
-  EncryptionKeyMap encryption_key_map;
-  encryption_key_map[kEmptyDrmLabel] = std::move(encryption_key);
-  RETURN_IF_ERROR(UpdateProtectionSystemInfo(&encryption_key_map));
-  encryption_key_ = std::move(encryption_key_map[kEmptyDrmLabel]);
   return Status::OK;
 }
 

packager/media/base/playready_key_source.h

@@ -23,12 +23,8 @@ class PlayReadyKeySource : public KeySource {
   /// @param server_url PlayReady packaging server url.
   /// @param protection_systems is the enum indicating which PSSH should
   ///        be included.
-  /// @param protection_scheme is the Protection Scheme to be used for
-  ///        encryption. It needs to be signalled in Widevine PSSH. This
-  ///        argument can be ignored if Widevine PSSH is not generated.
   PlayReadyKeySource(const std::string& server_url,
-                     ProtectionSystem protection_systems,
+                     ProtectionSystem protection_systems);
-                     FourCC protection_scheme);
   /// Creates a new PlayReadyKeySource from the given packaging information.
   /// @param server_url PlayReady packaging server url.
   /// @param client_cert_file absolute path to a client certificate.
@@ -37,15 +33,11 @@ class PlayReadyKeySource : public KeySource {
   /// @param client_cert_private_key_password password for the private key.
   /// @param protection_systems is the enum indicating which PSSH should
   ///        be included.
-  /// @param protection_scheme is the Protection Scheme to be used for
-  ///        encryption. It needs to be signalled in Widevine PSSH. This
-  ///        argument can be ignored if Widevine PSSH is not generated.
   PlayReadyKeySource(const std::string& server_url,
                      const std::string& client_cert_file,
                      const std::string& client_cert_private_key_file,
                      const std::string& client_cert_private_key_password,
-                     ProtectionSystem protection_systems,
+                     ProtectionSystem protection_systems);
-                     FourCC protection_scheme);
   ~PlayReadyKeySource() override;
 
   /// @name KeySource implementation overrides.

packager/media/base/playready_pssh_generator.cc

@@ -86,6 +86,7 @@ Status GeneratePlayReadyPsshData(const std::vector<uint8_t>& key_id,
   std::string playready_header;
 
   switch (protection_scheme) {
+    case kAppleSampleAesProtectionScheme:
     case FOURCC_cbc1:
     case FOURCC_cbcs:
       playready_header = kPlayHeaderObject_4_3;

packager/media/base/raw_key_source.cc

@@ -76,34 +76,16 @@ Status RawKeySource::GetCryptoPeriodKey(uint32_t crypto_period_index,
   std::rotate(key->key.begin(),
               key->key.begin() + (crypto_period_index % key->key.size()),
               key->key.end());
-
+  key->key_ids.clear();
-  // Clear |key->key_system_info| to prepare for update. The original
+  key->key_ids.emplace_back(key->key_id);
-  // |key_system_info| is saved as it may still be useful later.
-  std::vector<ProtectionSystemSpecificInfo> original_key_system_info;
-  key->key_system_info.swap(original_key_system_info);
-
-  EncryptionKeyMap encryption_key_map;
-  encryption_key_map[stream_label].reset(new EncryptionKey(*key));
-  RETURN_IF_ERROR(UpdateProtectionSystemInfo(&encryption_key_map));
-  key->key_system_info = encryption_key_map[stream_label]->key_system_info;
-
-  // It is possible that the generated |key_system_info| is empty. This happens
-  // when RawKeyParams.pssh is specified. Restore the original key system info
-  // in this case.
-  if (key->key_system_info.empty())
-    key->key_system_info.swap(original_key_system_info);
 
   return Status::OK;
 }
 
 std::unique_ptr<RawKeySource> RawKeySource::Create(
-    const RawKeyParams& raw_key,
+    const RawKeyParams& raw_key) {
-    ProtectionSystem protection_systems,
-    FourCC protection_scheme) {
   std::vector<ProtectionSystemSpecificInfo> key_system_info;
-  bool pssh_provided = false;
   if (!raw_key.pssh.empty()) {
-    pssh_provided = true;
     if (!ProtectionSystemSpecificInfo::ParseBoxes(
             raw_key.pssh.data(), raw_key.pssh.size(), &key_system_info)) {
       LOG(ERROR) << "--pssh argument should be full PSSH boxes.";
@@ -111,6 +93,10 @@ std::unique_ptr<RawKeySource> RawKeySource::Create(
     }
   }
 
+  std::vector<std::vector<uint8_t>> key_ids;
+  for (const auto& entry : raw_key.key_map)
+    key_ids.emplace_back(entry.second.key_id);
+
   EncryptionKeyMap encryption_key_map;
   for (const auto& entry : raw_key.key_map) {
     const std::string& drm_label = entry.first;
@@ -130,32 +116,21 @@ std::unique_ptr<RawKeySource> RawKeySource::Create(
 
     std::unique_ptr<EncryptionKey> encryption_key(new EncryptionKey);
     encryption_key->key_id = key_pair.key_id;
+    encryption_key->key_ids = key_ids;
     encryption_key->key = key_pair.key;
     encryption_key->iv = raw_key.iv;
     encryption_key->key_system_info = key_system_info;
     encryption_key_map[drm_label] = std::move(encryption_key);
   }
 
-  // Generate common protection system if no other protection system is
+  return std::unique_ptr<RawKeySource>(
-  // specified.
+      new RawKeySource(std::move(encryption_key_map)));
-  if (!pssh_provided && protection_systems == ProtectionSystem::kNone) {
-    protection_systems = ProtectionSystem::kCommon;
 }
 
-  return std::unique_ptr<RawKeySource>(new RawKeySource(
+RawKeySource::RawKeySource() {}
-      std::move(encryption_key_map), protection_systems, protection_scheme));
-}
 
-RawKeySource::RawKeySource()
+RawKeySource::RawKeySource(EncryptionKeyMap&& encryption_key_map)
-    : KeySource(ProtectionSystem::kNone, FOURCC_NULL) {}
+    : encryption_key_map_(std::move(encryption_key_map)) {}
-
-RawKeySource::RawKeySource(EncryptionKeyMap&& encryption_key_map,
-                           ProtectionSystem protection_systems,
-                           FourCC protection_scheme)
-    : KeySource(protection_systems, protection_scheme),
-      encryption_key_map_(std::move(encryption_key_map)) {
-  UpdateProtectionSystemInfo(&encryption_key_map_);
-}
 
 }  // namespace media
 }  // namespace shaka

packager/media/base/raw_key_source.h

@@ -38,24 +38,14 @@ class RawKeySource : public KeySource {
   /// Creates a new RawKeySource from the given data.  Returns null
   /// if the parameter is malformed.
   /// @param raw_key contains parameters to setup the key source.
-  /// @param protection_systems is the enum indicating which PSSH should
+  static std::unique_ptr<RawKeySource> Create(const RawKeyParams& raw_key);
-  ///        be included.
-  /// @param protection_scheme is the Protection Scheme to be used for
-  ///        encryption. It needs to be signalled in Widevine PSSH. This
-  ///        argument can be ignored if Widevine PSSH is not generated.
-  static std::unique_ptr<RawKeySource> Create(
-      const RawKeyParams& raw_key,
-      ProtectionSystem protection_systems,
-      FourCC protection_scheme);
 
  protected:
   // Allow default constructor for mock key sources.
   RawKeySource();
 
  private:
-  RawKeySource(EncryptionKeyMap&& encryption_key_map,
+  RawKeySource(EncryptionKeyMap&& encryption_key_map);
-               ProtectionSystem protection_systems,
-               FourCC protection_scheme);
   RawKeySource(const RawKeySource&) = delete;
   RawKeySource& operator=(const RawKeySource&) = delete;
 

packager/media/base/raw_key_source_unittest.cc

@@ -41,13 +41,6 @@ const char kPsshBox2Hex[] =
     "bbbbbbbbbaaaaaaaaaaaaddddddddddd"
     "0000000e"
     "fffffffff000000000000ddddddd";
-const char kDefaultPsshBoxHex[] =
-    "000000447073736801000000"
-    "1077efecc0b24d02ace33c1e52e2fb4b"
-    "00000002"
-    "1111121315180d1522375990e9000000"
-    "0101020305080d1522375990e9000000"
-    "00000000";
 const char kDrmLabel[] = "SomeDrmLabel";
 const char kAnotherDrmLabel[] = "AnotherDrmLabel";
 const char kEmptyDrmLabel[] = "";
@@ -68,8 +61,8 @@ TEST(RawKeySourceTest, Success) {
   raw_key_params.iv = HexStringToVector(kIvHex);
   raw_key_params.pssh =
       HexStringToVector(std::string(kPsshBox1Hex) + kPsshBox2Hex);
-  std::unique_ptr<RawKeySource> key_source = RawKeySource::Create(
+  std::unique_ptr<RawKeySource> key_source =
-      raw_key_params, ProtectionSystem::kNone, FOURCC_NULL);
+      RawKeySource::Create(raw_key_params);
   ASSERT_NE(nullptr, key_source);
 
   EncryptionKey key_from_drm_label;
@@ -98,27 +91,6 @@ TEST(RawKeySourceTest, Success) {
   EXPECT_HEX_EQ(kPsshBox2Hex, key_from_drm_label.key_system_info[1].psshs);
 }
 
-TEST(RawKeySourceTest, EmptyPssh) {
-  RawKeyParams raw_key_params;
-  raw_key_params.key_map[kDrmLabel].key_id = HexStringToVector(kKeyIdHex);
-  raw_key_params.key_map[kDrmLabel].key = HexStringToVector(kKeyHex);
-  raw_key_params.key_map[kAnotherDrmLabel].key_id =
-      HexStringToVector(kKeyId2Hex);
-  raw_key_params.key_map[kAnotherDrmLabel].key = HexStringToVector(kKey2Hex);
-  raw_key_params.iv = HexStringToVector(kIvHex);
-  std::unique_ptr<RawKeySource> key_source = RawKeySource::Create(
-      raw_key_params, ProtectionSystem::kNone, FOURCC_NULL);
-  ASSERT_NE(nullptr, key_source);
-
-  EncryptionKey key;
-  ASSERT_OK(key_source->GetKey(kDrmLabel, &key));
-  EXPECT_HEX_EQ(kKeyIdHex, key.key_id);
-  EXPECT_HEX_EQ(kKeyHex, key.key);
-  EXPECT_HEX_EQ(kIvHex, key.iv);
-  ASSERT_EQ(1u, key.key_system_info.size());
-  EXPECT_HEX_EQ(kDefaultPsshBoxHex, key.key_system_info[0].psshs);
-}
-
 TEST(RawKeySourceTest, Failure) {
   // Invalid key id size.
   RawKeyParams raw_key_params;
@@ -127,15 +99,14 @@ TEST(RawKeySourceTest, Failure) {
   raw_key_params.key_map[kEmptyDrmLabel].key = HexStringToVector(kKeyHex);
   raw_key_params.pssh = HexStringToVector(kPsshBox1Hex);
   raw_key_params.iv = HexStringToVector(kIvHex);
-  std::unique_ptr<RawKeySource> key_source = RawKeySource::Create(
+  std::unique_ptr<RawKeySource> key_source =
-      raw_key_params, ProtectionSystem::kNone, FOURCC_NULL);
+      RawKeySource::Create(raw_key_params);
   EXPECT_EQ(nullptr, key_source);
 
   // Invalid pssh box.
   raw_key_params.key_map[kEmptyDrmLabel].key_id = HexStringToVector(kKeyIdHex);
   raw_key_params.pssh = HexStringToVector("000102030405");
-  key_source = RawKeySource::Create(raw_key_params, ProtectionSystem::kNone,
+  key_source = RawKeySource::Create(raw_key_params);
-                                    FOURCC_NULL);
   EXPECT_EQ(nullptr, key_source);
 }
 

packager/media/base/widevine_key_source.cc

@@ -87,9 +87,7 @@ WidevineKeySource::WidevineKeySource(const std::string& server_url,
                                      ProtectionSystem protection_systems,
                                      FourCC protection_scheme)
     // Widevine PSSH is fetched from Widevine license server.
-    : KeySource(protection_systems & ~ProtectionSystem::kWidevine,
+    : generate_widevine_protection_system_(
-                protection_scheme),
-      generate_widevine_protection_system_(
           // Generate Widevine protection system if there are no other
           // protection system specified.
           protection_systems == ProtectionSystem::kNone ||
@@ -433,6 +431,12 @@ bool WidevineKeySource::ExtractEncryptionKey(
 
   uint32_t current_crypto_period_index = first_crypto_period_index_;
 
+  std::vector<std::vector<uint8_t>> key_ids;
+  for (const auto& track : response_proto.tracks()) {
+    if (!widevine_classic)
+      key_ids.emplace_back(track.key_id().begin(), track.key_id().end());
+  }
+
   EncryptionKeyMap encryption_key_map;
   for (const auto& track : response_proto.tracks()) {
     VLOG(2) << "track " << track.ShortDebugString();
@@ -463,6 +467,7 @@ bool WidevineKeySource::ExtractEncryptionKey(
       encryption_key->key_id.assign(track.key_id().begin(),
                                     track.key_id().end());
       encryption_key->iv.assign(track.iv().begin(), track.iv().end());
+      encryption_key->key_ids = key_ids;
 
       if (generate_widevine_protection_system_) {
         if (track.pssh_size() != 1) {
@@ -477,12 +482,6 @@ bool WidevineKeySource::ExtractEncryptionKey(
     encryption_key_map[stream_label] = std::move(encryption_key);
   }
 
-  if (!widevine_classic) {
-    if (!UpdateProtectionSystemInfo(&encryption_key_map).ok()) {
-      return false;
-    }
-  }
-
   DCHECK(!encryption_key_map.empty());
   if (!enable_key_rotation) {
     // Merge with previously requested keys.
@@ -490,6 +489,7 @@ bool WidevineKeySource::ExtractEncryptionKey(
       encryption_key_map_[pair.first] = std::move(pair.second);
     return true;
   }
+
   return PushToKeyPool(&encryption_key_map);
 }
 

packager/media/base/widevine_key_source_unittest.cc

@@ -270,7 +270,7 @@ class WidevineKeySourceTest : public Test {
       EXPECT_EQ(GetMockKey(stream_label), ToString(encryption_key.key));
       if (!classic) {
         size_t num_key_system_info =
-            (add_widevine_pssh_ && add_common_pssh_) ? 2 : 1;
+            add_widevine_pssh_ || !add_common_pssh_ ? 1 : 0;
         ASSERT_EQ(num_key_system_info, encryption_key.key_system_info.size());
         EXPECT_EQ(GetMockKeyId(stream_label), ToString(encryption_key.key_id));
         if (has_iv)
@@ -294,12 +294,6 @@ class WidevineKeySourceTest : public Test {
 
           ++key_system_info_iter;
         }
-
-        if (add_common_pssh_) {
-          const std::vector<uint8_t> common_system_id(
-              std::begin(kCommonSystemId), std::end(kCommonSystemId));
-          ASSERT_EQ(common_system_id, key_system_info_iter->system_id);
-        }
       }
     }
   }

packager/media/crypto/encryption_handler.cc

@@ -13,10 +13,14 @@
 
 #include "packager/media/base/aes_encryptor.h"
 #include "packager/media/base/audio_stream_info.h"
+#include "packager/media/base/common_pssh_generator.h"
 #include "packager/media/base/key_source.h"
 #include "packager/media/base/macros.h"
 #include "packager/media/base/media_sample.h"
+#include "packager/media/base/playready_pssh_generator.h"
+#include "packager/media/base/protection_system_ids.h"
 #include "packager/media/base/video_stream_info.h"
+#include "packager/media/base/widevine_pssh_generator.h"
 #include "packager/media/crypto/aes_encryptor_factory.h"
 #include "packager/media/crypto/subsample_generator.h"
 #include "packager/status_macros.h"
@@ -65,6 +69,94 @@ bool IsPatternEncryptionScheme(FourCC protection_scheme) {
          protection_scheme == FOURCC_cbcs || protection_scheme == FOURCC_cens;
 }
 
+void FillPsshGenerators(
+    const EncryptionParams& encryption_params,
+    std::vector<std::unique_ptr<PsshGenerator>>* pssh_generators,
+    std::vector<std::vector<uint8_t>>* no_pssh_systems) {
+  if (has_flag(encryption_params.protection_systems,
+               ProtectionSystem::kCommon)) {
+    pssh_generators->emplace_back(new CommonPsshGenerator());
+  }
+
+  if (has_flag(encryption_params.protection_systems,
+               ProtectionSystem::kPlayReady)) {
+    pssh_generators->emplace_back(new PlayReadyPsshGenerator(
+        static_cast<FourCC>(encryption_params.protection_scheme)));
+  }
+
+  if (has_flag(encryption_params.protection_systems,
+               ProtectionSystem::kWidevine)) {
+    pssh_generators->emplace_back(new WidevinePsshGenerator(
+        static_cast<FourCC>(encryption_params.protection_scheme)));
+  }
+
+  if (has_flag(encryption_params.protection_systems,
+               ProtectionSystem::kFairPlay)) {
+    no_pssh_systems->emplace_back(std::begin(kFairPlaySystemId),
+                                  std::end(kFairPlaySystemId));
+  }
+  // We only support Marlin Adaptive Streaming Specification – Simple Profile
+  // with Implicit Content ID Mapping, which does not need a PSSH. Marlin
+  // specific PSSH with Explicit Content ID Mapping is not generated.
+  if (has_flag(encryption_params.protection_systems,
+               ProtectionSystem::kMarlin)) {
+    no_pssh_systems->emplace_back(std::begin(kMarlinSystemId),
+                                  std::end(kMarlinSystemId));
+  }
+
+  if (pssh_generators->empty() && no_pssh_systems->empty() &&
+      (encryption_params.key_provider != KeyProvider::kRawKey ||
+       encryption_params.raw_key.pssh.empty())) {
+    pssh_generators->emplace_back(new CommonPsshGenerator());
+  }
+}
+
+void AddProtectionSystemIfNotExist(
+    const ProtectionSystemSpecificInfo& pssh_info,
+    EncryptionConfig* encryption_config) {
+  for (const auto& info : encryption_config->key_system_info) {
+    if (info.system_id == pssh_info.system_id)
+      return;
+  }
+  encryption_config->key_system_info.push_back(pssh_info);
+}
+
+Status FillProtectionSystemInfo(const EncryptionParams& encryption_params,
+                                const EncryptionKey& encryption_key,
+                                EncryptionConfig* encryption_config) {
+  // If generating dummy keys for key rotation, don't generate PSSH info.
+  if (encryption_key.key_ids.empty())
+    return Status::OK;
+
+  std::vector<std::unique_ptr<PsshGenerator>> pssh_generators;
+  std::vector<std::vector<uint8_t>> no_pssh_systems;
+  FillPsshGenerators(encryption_params, &pssh_generators, &no_pssh_systems);
+
+  encryption_config->key_system_info = encryption_key.key_system_info;
+  for (const auto& pssh_generator : pssh_generators) {
+    const bool support_multiple_keys = pssh_generator->SupportMultipleKeys();
+    if (support_multiple_keys) {
+      ProtectionSystemSpecificInfo info;
+      RETURN_IF_ERROR(pssh_generator->GeneratePsshFromKeyIds(
+          encryption_key.key_ids, &info));
+      AddProtectionSystemIfNotExist(info, encryption_config);
+    } else {
+      ProtectionSystemSpecificInfo info;
+      RETURN_IF_ERROR(pssh_generator->GeneratePsshFromKeyIdAndKey(
+          encryption_key.key_id, encryption_key.key, &info));
+      AddProtectionSystemIfNotExist(info, encryption_config);
+    }
+  }
+
+  for (const auto& no_pssh_system : no_pssh_systems) {
+    ProtectionSystemSpecificInfo info;
+    info.system_id = no_pssh_system;
+    AddProtectionSystemIfNotExist(info, encryption_config);
+  }
+
+  return Status::OK;
+}
+
 }  // namespace
 
 EncryptionHandler::EncryptionHandler(const EncryptionParams& encryption_params,
@@ -286,8 +378,9 @@ bool EncryptionHandler::CreateEncryptor(const EncryptionKey& encryption_key) {
   }
 
   encryption_config_->key_id = encryption_key.key_id;
-  encryption_config_->key_system_info = encryption_key.key_system_info;
+  const auto status = FillProtectionSystemInfo(
-  return true;
+      encryption_params_, encryption_key, encryption_config_.get());
+  return status.ok();
 }
 
 void EncryptionHandler::EncryptBytes(const uint8_t* source,

packager/media/crypto/encryption_handler_unittest.cc

@@ -12,6 +12,7 @@
 #include "packager/media/base/aes_cryptor.h"
 #include "packager/media/base/media_handler_test_base.h"
 #include "packager/media/base/mock_aes_cryptor.h"
+#include "packager/media/base/protection_system_ids.h"
 #include "packager/media/base/raw_key_source.h"
 #include "packager/media/crypto/aes_encryptor_factory.h"
 #include "packager/media/crypto/subsample_generator.h"
@@ -31,6 +32,7 @@ using ::testing::Mock;
 using ::testing::Return;
 using ::testing::SetArgPointee;
 using ::testing::StrictMock;
+using ::testing::UnorderedElementsAre;
 using ::testing::Values;
 using ::testing::ValuesIn;
 using ::testing::WithParamInterface;
@@ -127,6 +129,7 @@ class EncryptionHandlerTest : public MediaHandlerGraphTestBase {
   EncryptionKey GetMockEncryptionKey() {
     EncryptionKey encryption_key;
     encryption_key.key_id.assign(kKeyId, kKeyId + sizeof(kKeyId));
+    encryption_key.key_ids.emplace_back(encryption_key.key_id);
     encryption_key.key.assign(kKey, kKey + sizeof(kKey));
     encryption_key.iv.assign(kIv, kIv + sizeof(kIv));
     return encryption_key;
@@ -602,5 +605,68 @@ TEST_F(EncryptionHandlerTrackTypeTest, VideoTrackType) {
   EXPECT_EQ(captured_stream_attributes.oneof.video.height, kHeight);
 }
 
+class EncryptionHandlerPsshTest : public EncryptionHandlerTest {};
+
+TEST_F(EncryptionHandlerPsshTest, GeneratesPssh) {
+  EncryptionParams encryption_params;
+  encryption_params.protection_scheme = FOURCC_cenc;
+  encryption_params.protection_systems =
+      ProtectionSystem::kWidevine | ProtectionSystem::kPlayReady;
+  SetUpEncryptionHandler(encryption_params);
+
+  const EncryptionKey mock_encryption_key = GetMockEncryptionKey();
+  EXPECT_CALL(mock_key_source_, GetKey(_, _))
+      .WillOnce(
+          DoAll(SetArgPointee<1>(mock_encryption_key), Return(Status::OK)));
+
+  ASSERT_OK(Process(StreamData::FromStreamInfo(
+      kStreamIndex, GetVideoStreamInfo(kTimeScale, kCodecH264))));
+
+  EXPECT_THAT(GetOutputStreamDataVector(),
+              ElementsAre(IsStreamInfo(_, kTimeScale, kEncrypted, _)));
+  const StreamInfo* stream_info =
+      GetOutputStreamDataVector().back()->stream_info.get();
+
+  std::vector<uint8_t> widevine_system_id(
+      kWidevineSystemId, kWidevineSystemId + arraysize(kWidevineSystemId));
+  std::vector<uint8_t> playready_system_id(
+      kPlayReadySystemId, kPlayReadySystemId + arraysize(kPlayReadySystemId));
+  ASSERT_THAT(
+      stream_info->encryption_config().key_system_info,
+      UnorderedElementsAre(IsPsshInfoWithSystemId(widevine_system_id),
+                           IsPsshInfoWithSystemId(playready_system_id)));
+}
+
+TEST_F(EncryptionHandlerPsshTest, UsesKeyInfoFirst) {
+  EncryptionParams encryption_params;
+  encryption_params.protection_scheme = FOURCC_cenc;
+  encryption_params.protection_systems = ProtectionSystem::kWidevine;
+  SetUpEncryptionHandler(encryption_params);
+
+  std::vector<uint8_t> widevine_system_id(
+      kWidevineSystemId, kWidevineSystemId + arraysize(kWidevineSystemId));
+  EncryptionKey mock_encryption_key = GetMockEncryptionKey();
+  ProtectionSystemSpecificInfo protection_info;
+  protection_info.system_id = widevine_system_id;
+  mock_encryption_key.key_system_info.emplace_back(protection_info);
+  EXPECT_CALL(mock_key_source_, GetKey(_, _))
+      .WillOnce(
+          DoAll(SetArgPointee<1>(mock_encryption_key), Return(Status::OK)));
+
+  ASSERT_OK(Process(StreamData::FromStreamInfo(
+      kStreamIndex, GetVideoStreamInfo(kTimeScale, kCodecH264))));
+
+  EXPECT_THAT(GetOutputStreamDataVector(),
+              ElementsAre(IsStreamInfo(_, kTimeScale, kEncrypted, _)));
+  const StreamInfo* stream_info =
+      GetOutputStreamDataVector().back()->stream_info.get();
+
+  ASSERT_THAT(stream_info->encryption_config().key_system_info,
+              ElementsAre(IsPsshInfoWithSystemId(widevine_system_id)));
+  // Should use above info, not generate new info.
+  ASSERT_TRUE(
+      stream_info->encryption_config().key_system_info[0].psshs.empty());
+}
+
 }  // namespace media
 }  // namespace shaka