From 1bf1ec24452f887bebedfa8fe8353f534e4b311b Mon Sep 17 00:00:00 2001
From: KongQun Yang <kqyang@google.com>
Date: Wed, 20 Sep 2017 15:49:00 -0700
Subject: [PATCH] Add multi-key support in fixed/raw keys

The keys can be specified in --keys option, with the form of
  label=<label>:key_id=<32-digit hex string>:key=<32-digit hex string>,label=...

There can be multiple "label=..." block. The DRM label can be one
of "AUDIO, SD, HD, UHD1, UHD2" or a custom label.

Mark --key and --key_id as deprecated. We can achieve the same result
with --keys=label=:key_id=<key_id>:key=<key>.

Also add a new optional field in stream descriptor: drm_label, which
is used to overwrite the internally generated DRM label for the stream.
If not provided, the DRM label is generated automatically based on
audio/video information, e.g. resolution.

Code changes:
- Merged RawKeyEncryptionParams and RawKeyDecryptionParams into
  RawKeyParams.
- Make FixedKeySource accepts RawKeyParams as input.

Change-Id: Ic8c2f071cc71188e13f14bc6396fc2b3ffa5cac6
---
 .../options/raw_key_encryption_options.rst    |  19 ++-
 docs/source/options/stream_descriptors.rst    |   8 ++
 packager/app/fixed_key_encryption_flags.cc    |  31 ++++-
 packager/app/fixed_key_encryption_flags.h     |   1 +
 packager/app/packager_main.cc                 |  71 ++++++++--
 packager/app/packager_util.cc                 |  14 +-
 packager/app/stream_descriptor.cc             |   7 +
 packager/app/test/packager_test.py            |  62 ++++++++-
 packager/app/validate_flag.cc                 |   4 +
 packager/app/validate_flag.h                  |   4 +
 packager/media/base/fixed_key_source.cc       | 116 +++++++++-------
 packager/media/base/fixed_key_source.h        |  27 ++--
 .../media/base/fixed_key_source_unittest.cc   | 125 ++++++++++++------
 .../base/protection_system_specific_info.h    |   2 -
 packager/media/public/crypto_params.h         |  33 ++---
 packager/packager.cc                          |   8 +-
 packager/packager.h                           |   5 +
 17 files changed, 365 insertions(+), 172 deletions(-)

diff --git a/docs/source/options/raw_key_encryption_options.rst b/docs/source/options/raw_key_encryption_options.rst
index 2e34d1d6c8..31820f2a0f 100644
--- a/docs/source/options/raw_key_encryption_options.rst
+++ b/docs/source/options/raw_key_encryption_options.rst
@@ -9,14 +9,27 @@ Raw key encryption options
 
     Enable decryption with fixed key.
 
+--keys <key_info_string[,key_info_string][,key_info_string]...>
+
+    **key_info_string** is of the form::
+
+        label={label}:key_id={key_id}:key={key}
+
+    *label* can be an arbitrary string or a predefined DRM label like AUDIO,
+    SD, HD, etc. Label with an empty string indicates the default key and
+    key_id. The *drm_label* in :doc:`/options/stream_descriptors`,
+    which can be implicit, determines which key info is applied to the stream
+    by matching the *drm_label* with the *label* in key info.
+
+    *key_id* and *key* should be 32-digit hex strings.
+
 --key_id <32-digit hex string>
 
-    The key id in hex string format.
-    HEX.
+    The key id in hex string format. Deprecated. Use --keys instead.
 
 --key <32-digit hex string>
 
-    The key in hex string format.
+    The key in hex string format. Deprecated. Use --keys instead.
 
 --iv <16-digit or 32-digit hex string>
 
diff --git a/docs/source/options/stream_descriptors.rst b/docs/source/options/stream_descriptors.rst
index ad0fcc6305..ac97e252ff 100644
--- a/docs/source/options/stream_descriptors.rst
+++ b/docs/source/options/stream_descriptors.rst
@@ -62,6 +62,14 @@ These are the available fields:
     Optional. Defaults to 0 if not specified. If it is set to 1, no encryption
     of the stream will be made.
 
+:drm_label:
+
+    Optional value for custom DRM label, which defines the encryption key
+    applied to the stream. Typically values include AUDIO, SD, HD, UHD1, UHD2.
+    For raw key, it should be a label defined in --keys. If not provided, the
+    DRM label is derived from stream type (video, audio), resolutions, etc.
+    Note that it is case sensitive.
+
 :trick_play_factor (tpf):
 
     Optional value which specifies the trick play, a.k.a. trick mode, stream
diff --git a/packager/app/fixed_key_encryption_flags.cc b/packager/app/fixed_key_encryption_flags.cc
index 6c67421353..63259219e9 100644
--- a/packager/app/fixed_key_encryption_flags.cc
+++ b/packager/app/fixed_key_encryption_flags.cc
@@ -18,6 +18,11 @@ DEFINE_bool(enable_fixed_key_decryption,
             "Enable decryption with fixed key.");
 DEFINE_hex_bytes(key_id, "", "Key id in hex string format.");
 DEFINE_hex_bytes(key, "", "Key in hex string format.");
+DEFINE_string(keys,
+              "",
+              "A list of key information in the form of label=<drm "
+              "label>:key_id=<32-digit key id in hex>:key=<32-digit key in "
+              "hex>,label=...");
 DEFINE_hex_bytes(
     iv,
     "",
@@ -40,13 +45,25 @@ bool ValidateFixedCryptoFlags() {
   const char fixed_crypto_label[] = "--enable_fixed_key_encryption/decryption";
   // --key_id and --key are associated with --enable_fixed_key_encryption and
   // --enable_fixed_key_decryption.
-  if (!ValidateFlag("key_id", FLAGS_key_id_bytes, fixed_crypto, false,
-                    fixed_crypto_label)) {
-    success = false;
-  }
-  if (!ValidateFlag("key", FLAGS_key_bytes, fixed_crypto, false,
-                    fixed_crypto_label)) {
-    success = false;
+  if (FLAGS_keys.empty()) {
+    if (!ValidateFlag("key_id", FLAGS_key_id_bytes, fixed_crypto, false,
+                      fixed_crypto_label)) {
+      success = false;
+    }
+    if (!ValidateFlag("key", FLAGS_key_bytes, fixed_crypto, false,
+                      fixed_crypto_label)) {
+      success = false;
+    }
+    if (success && (!FLAGS_key_id_bytes.empty() || !FLAGS_key_bytes.empty())) {
+      PrintWarning(
+          "--key_id and --key are going to be deprecated. Please switch to "
+          "--keys as soon as possible.");
+    }
+  } else {
+    if (!FLAGS_key_id_bytes.empty() || !FLAGS_key_bytes.empty()) {
+      PrintError("--key_id or --key cannot be used together with --keys.");
+      success = false;
+    }
   }
   if (!ValidateFlag("iv", FLAGS_iv_bytes, FLAGS_enable_fixed_key_encryption,
                     true, "--enable_fixed_key_encryption")) {
diff --git a/packager/app/fixed_key_encryption_flags.h b/packager/app/fixed_key_encryption_flags.h
index b3ef29fbdd..8a93f36e1d 100644
--- a/packager/app/fixed_key_encryption_flags.h
+++ b/packager/app/fixed_key_encryption_flags.h
@@ -18,6 +18,7 @@ DECLARE_bool(enable_fixed_key_encryption);
 DECLARE_bool(enable_fixed_key_decryption);
 DECLARE_hex_bytes(key_id);
 DECLARE_hex_bytes(key);
+DECLARE_string(keys);
 DECLARE_hex_bytes(iv);
 DECLARE_hex_bytes(pssh);
 
diff --git a/packager/app/packager_main.cc b/packager/app/packager_main.cc
index c525a6a605..040389c7f0 100644
--- a/packager/app/packager_main.cc
+++ b/packager/app/packager_main.cc
@@ -82,6 +82,12 @@ const char kUsage[] =
     "    derived from the file extension of the output file.\n"
     "  - skip_encryption=0|1: Optional. Defaults to 0 if not specified. If\n"
     "    it is set to 1, no encryption of the stream will be made.\n"
+    "  - drm_label: Optional value for custom DRM label, which defines the\n"
+    "    encryption key applied to the stream. Typical values include AUDIO,\n"
+    "    SD, HD, UHD1, UHD2. For raw key, it should be a label defined in\n"
+    "    --keys. If not provided, the DRM label is derived from stream type\n"
+    "    (video, audio), resolution, etc.\n"
+    "    Note that it is case sensitive.\n"
     "  - trick_play_factor (tpf): Optional value which specifies the trick\n"
     "    play, a.k.a. trick mode, stream sampling rate among key frames.\n"
     "    If specified, the output is a trick play stream.\n"
@@ -93,6 +99,11 @@ const char kUsage[] =
     "    Usually ends with '.m3u8'. If unspecified, defaults to something of\n"
     "    the form 'stream_0.m3u8', 'stream_1.m3u8', 'stream_2.m3u8', etc.\n";
 
+// Labels for parameters in RawKey key info.
+const char kDrmLabelLabel[] = "label";
+const char kKeyIdLabel[] = "key_id";
+const char kKeyLabel[] = "key";
+
 enum ExitStatus {
   kSuccess = 0,
   kArgumentValidationFailed,
@@ -154,6 +165,50 @@ bool GetProtectionScheme(uint32_t* protection_scheme) {
   return false;
 }
 
+bool ParseKeys(const std::string& keys, RawKeyParams* raw_key) {
+  for (const std::string& key_data : base::SplitString(
+           keys, ",", base::TRIM_WHITESPACE, base::SPLIT_WANT_NONEMPTY)) {
+    base::StringPairs string_pairs;
+    base::SplitStringIntoKeyValuePairs(key_data, '=', ':', &string_pairs);
+
+    std::map<std::string, std::string> value_map;
+    for (const auto& string_pair : string_pairs)
+      value_map[string_pair.first] = string_pair.second;
+    const std::string drm_label = value_map[kDrmLabelLabel];
+    if (raw_key->key_map.find(drm_label) != raw_key->key_map.end()) {
+      LOG(ERROR) << "Seeing duplicated DRM label '" << drm_label << "'.";
+      return false;
+    }
+    auto& key_info = raw_key->key_map[drm_label];
+    if (value_map[kKeyIdLabel].empty() ||
+        !base::HexStringToBytes(value_map[kKeyIdLabel], &key_info.key_id)) {
+      return false;
+    }
+    if (value_map[kKeyLabel].empty() ||
+        !base::HexStringToBytes(value_map[kKeyLabel], &key_info.key)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool GetRawKeyParams(RawKeyParams* raw_key) {
+  raw_key->iv = FLAGS_iv_bytes;
+  raw_key->pssh = FLAGS_pssh_bytes;
+  if (!FLAGS_keys.empty()) {
+    if (!ParseKeys(FLAGS_keys, raw_key)) {
+      LOG(ERROR) << "Failed to parse --keys " << FLAGS_keys;
+      return false;
+    }
+  } else {
+    // An empty StreamLabel specifies the default key info.
+    RawKeyParams::KeyInfo& key_info = raw_key->key_map[""];
+    key_info.key_id = FLAGS_key_id_bytes;
+    key_info.key = FLAGS_key_bytes;
+  }
+  return true;
+}
+
 base::Optional<PackagingParams> GetPackagingParams() {
   PackagingParams packaging_params;
 
@@ -223,13 +278,8 @@ base::Optional<PackagingParams> GetPackagingParams() {
       break;
     }
     case KeyProvider::kRawKey: {
-      RawKeyEncryptionParams& raw_key = encryption_params.raw_key;
-      raw_key.iv = FLAGS_iv_bytes;
-      raw_key.pssh = FLAGS_pssh_bytes;
-      // An empty StreamLabel specifies the default KeyPair.
-      RawKeyEncryptionParams::KeyPair& key_pair = raw_key.key_map[""];
-      key_pair.key_id = FLAGS_key_id_bytes;
-      key_pair.key = FLAGS_key_bytes;
+      if (!GetRawKeyParams(&encryption_params.raw_key))
+        return base::nullopt;
       break;
     }
     case KeyProvider::kNone:
@@ -260,11 +310,8 @@ base::Optional<PackagingParams> GetPackagingParams() {
       break;
     }
     case KeyProvider::kRawKey: {
-      RawKeyDecryptionParams& raw_key = decryption_params.raw_key;
-      // An empty StreamLabel specifies the default KeyPair.
-      RawKeyDecryptionParams::KeyPair& key_pair = raw_key.key_map[""];
-      key_pair.key_id = FLAGS_key_id_bytes;
-      key_pair.key = FLAGS_key_bytes;
+      if (!GetRawKeyParams(&decryption_params.raw_key))
+        return base::nullopt;
       break;
     }
     case KeyProvider::kPlayready:
diff --git a/packager/app/packager_util.cc b/packager/app/packager_util.cc
index a17c298048..32f2fbf3cb 100644
--- a/packager/app/packager_util.cc
+++ b/packager/app/packager_util.cc
@@ -85,12 +85,7 @@ std::unique_ptr<KeySource> CreateEncryptionKeySource(
       break;
     }
     case KeyProvider::kRawKey: {
-      const RawKeyEncryptionParams& raw_key = encryption_params.raw_key;
-      const std::string kDefaultTrackType;
-      // TODO(kqyang): Refactor FixedKeySource.
-      encryption_key_source = FixedKeySource::Create(
-          raw_key.key_map.find("")->second.key_id,
-          raw_key.key_map.find("")->second.key, raw_key.pssh, raw_key.iv);
+      encryption_key_source = FixedKeySource::Create(encryption_params.raw_key);
       break;
     }
     case KeyProvider::kPlayready: {
@@ -171,12 +166,7 @@ std::unique_ptr<KeySource> CreateDecryptionKeySource(
       break;
     }
     case KeyProvider::kRawKey: {
-      const RawKeyDecryptionParams& raw_key = decryption_params.raw_key;
-      const std::vector<uint8_t> kNoPssh;
-      const std::vector<uint8_t> kNoIv;
-      decryption_key_source = FixedKeySource::Create(
-          raw_key.key_map.find("")->second.key_id,
-          raw_key.key_map.find("")->second.key, kNoPssh, kNoIv);
+      decryption_key_source = FixedKeySource::Create(decryption_params.raw_key);
       break;
     }
     case KeyProvider::kNone:
diff --git a/packager/app/stream_descriptor.cc b/packager/app/stream_descriptor.cc
index 724dd09c37..2991f2d54f 100644
--- a/packager/app/stream_descriptor.cc
+++ b/packager/app/stream_descriptor.cc
@@ -28,6 +28,7 @@ enum FieldType {
   kHlsPlaylistNameField,
   kTrickPlayFactorField,
   kSkipEncryptionField,
+  kDrmStreamLabelField,
 };
 
 struct FieldNameToTypeMapping {
@@ -58,6 +59,8 @@ const FieldNameToTypeMapping kFieldNameTypeMappings[] = {
     {"trick_play_factor", kTrickPlayFactorField},
     {"tpf", kTrickPlayFactorField},
     {"skip_encryption", kSkipEncryptionField},
+    {"drm_stream_label", kDrmStreamLabelField},
+    {"drm_label", kDrmStreamLabelField},
 };
 
 FieldType GetFieldType(const std::string& field_name) {
@@ -154,6 +157,10 @@ base::Optional<StreamDescriptor> ParseStreamDescriptor(
         descriptor.skip_encryption = skip_encryption_value > 0;
         break;
       }
+      case kDrmStreamLabelField: {
+        descriptor.drm_label = iter->second;
+        break;
+      }
       default:
         LOG(ERROR) << "Unknown field in stream descriptor (\"" << iter->first
                    << "\").";
diff --git a/packager/app/test/packager_test.py b/packager/app/test/packager_test.py
index 25fee9265b..0a6d4f1a6d 100755
--- a/packager/app/test/packager_test.py
+++ b/packager/app/test/packager_test.py
@@ -151,10 +151,12 @@ class PackagerAppTest(unittest.TestCase):
           '--content_id=' + self.widevine_content_id,
       ]
     elif encryption:
-      flags += ['--enable_fixed_key_encryption',
-                '--key_id=' + self.encryption_key_id,
-                '--key=' + self.encryption_key,
-                '--clear_lead={0}'.format(clear_lead)]
+      flags += [
+          '--enable_fixed_key_encryption',
+          '--keys=label=:key_id={0}:key={1}'.format(self.encryption_key_id,
+                                                    self.encryption_key),
+          '--clear_lead={0}'.format(clear_lead)
+      ]
 
       if not random_iv:
         flags.append('--iv=' + self.encryption_iv)
@@ -164,9 +166,11 @@ class PackagerAppTest(unittest.TestCase):
       flags += ['--vp9_subsample_encryption=false']
 
     if decryption:
-      flags += ['--enable_fixed_key_decryption',
-                '--key_id=' + self.encryption_key_id,
-                '--key=' + self.encryption_key]
+      flags += [
+          '--enable_fixed_key_decryption',
+          '--keys=label=:key_id={0}:key={1}'.format(self.encryption_key_id,
+                                                    self.encryption_key)
+      ]
 
     if key_rotation:
       flags.append('--crypto_period_duration=1')
@@ -567,6 +571,50 @@ class PackagerFunctionalTest(PackagerAppTest):
     self._VerifyDecryption(self.output[0], 'bear-640x360-a-golden.mp4')
     self._VerifyDecryption(self.output[1], 'bear-640x360-v-golden.mp4')
 
+  def testPackageWithEncryptionMultiKeys(self):
+    audio_key_id = '10111213141516171819202122232425'
+    audio_key = '11121314151617181920212223242526'
+    video_key_id = '20212223242526272829303132333435'
+    video_key = '21222324252627282930313233343536'
+    flags = self._GetFlags() + [
+        '--enable_fixed_key_encryption',
+        '--keys=label=AUDIO:key_id={0}:key={1},label=SD:key_id={2}:key={3}'.
+        format(audio_key_id, audio_key, video_key_id, video_key),
+        '--clear_lead={0}'.format(1), '--iv={0}'.format(self.encryption_iv)
+    ]
+    self.assertPackageSuccess(self._GetStreams(['audio', 'video']), flags)
+
+    self.encryption_key_id = audio_key_id
+    self.encryption_key = audio_key
+    self._VerifyDecryption(self.output[0], 'bear-640x360-a-golden.mp4')
+    self.encryption_key_id = video_key_id
+    self.encryption_key = video_key
+    self._VerifyDecryption(self.output[1], 'bear-640x360-v-golden.mp4')
+
+  def testPackageWithEncryptionMultiKeysWithStreamLabel(self):
+    audio_key_id = '20212223242526272829303132333435'
+    audio_key = '21222324252627282930313233343536'
+    video_key_id = '10111213141516171819202122232425'
+    video_key = '11121314151617181920212223242526'
+    flags = self._GetFlags() + [
+        '--enable_fixed_key_encryption',
+        '--keys=label=MyAudio:key_id={0}:key={1},label=:key_id={2}:key={3}'.
+        format(audio_key_id, audio_key, video_key_id, video_key),
+        '--clear_lead={0}'.format(1), '--iv={0}'.format(self.encryption_iv)
+    ]
+    # DRM label 'MyVideo' is not defined, will fall back to the key for the
+    # empty default label.
+    self.assertPackageSuccess(
+        self._GetStreams(['audio,drm_label=MyAudio',
+                          'video,drm_label=MyVideo']), flags)
+
+    self.encryption_key_id = audio_key_id
+    self.encryption_key = audio_key
+    self._VerifyDecryption(self.output[0], 'bear-640x360-a-golden.mp4')
+    self.encryption_key_id = video_key_id
+    self.encryption_key = video_key
+    self._VerifyDecryption(self.output[1], 'bear-640x360-v-golden.mp4')
+
   def testPackageWithEncryptionOfOnlyVideoStream(self):
     self.assertPackageSuccess(
         self._GetStreams(['audio,skip_encryption=1', 'video']),
diff --git a/packager/app/validate_flag.cc b/packager/app/validate_flag.cc
index 583e9aa59b..00565b6b3c 100644
--- a/packager/app/validate_flag.cc
+++ b/packager/app/validate_flag.cc
@@ -16,4 +16,8 @@ void PrintError(const std::string& error_message) {
   fprintf(stderr, "ERROR: %s\n", error_message.c_str());
 }
 
+void PrintWarning(const std::string& warning_message) {
+  fprintf(stderr, "WARNING: %s\n", warning_message.c_str());
+}
+
 }  // namespace shaka
diff --git a/packager/app/validate_flag.h b/packager/app/validate_flag.h
index b7e5285713..d81bcf414c 100644
--- a/packager/app/validate_flag.h
+++ b/packager/app/validate_flag.h
@@ -19,6 +19,10 @@ namespace shaka {
 /// @param error_message specifies the error message.
 void PrintError(const std::string& error_message);
 
+/// Format and print warning message.
+/// @param warning_message specifies the warning message.
+void PrintWarning(const std::string& warning_message);
+
 /// Validate a flag against the given condition.
 /// @param flag_name is the name of the flag.
 /// @param flag_value is the value of the flag.
diff --git a/packager/media/base/fixed_key_source.cc b/packager/media/base/fixed_key_source.cc
index b1f30d7871..8317909c91 100644
--- a/packager/media/base/fixed_key_source.cc
+++ b/packager/media/base/fixed_key_source.cc
@@ -10,6 +10,10 @@
 #include "packager/base/logging.h"
 #include "packager/base/strings/string_number_conversions.h"
 
+namespace {
+const char kEmptyDrmLabel[] = "";
+}  // namespace
+
 namespace shaka {
 namespace media {
 
@@ -24,30 +28,40 @@ Status FixedKeySource::FetchKeys(EmeInitDataType init_data_type,
 Status FixedKeySource::GetKey(const std::string& stream_label,
                               EncryptionKey* key) {
   DCHECK(key);
-  DCHECK(encryption_key_);
-  *key = *encryption_key_;
+  // Try to find the key with label |stream_label|. If it is not available,
+  // fall back to the default empty label if it is available.
+  auto iter = encryption_key_map_.find(stream_label);
+  if (iter == encryption_key_map_.end()) {
+    iter = encryption_key_map_.find(kEmptyDrmLabel);
+    if (iter == encryption_key_map_.end()) {
+      return Status(error::NOT_FOUND,
+                    "Key for '" + stream_label + "' was not found.");
+    }
+  }
+  *key = *iter->second;
   return Status::OK;
 }
 
 Status FixedKeySource::GetKey(const std::vector<uint8_t>& key_id,
                               EncryptionKey* key) {
   DCHECK(key);
-  DCHECK(encryption_key_);
-  if (key_id != encryption_key_->key_id) {
-    return Status(error::NOT_FOUND,
-                  std::string("Key for key ID ") +
-                      base::HexEncode(&key_id[0], key_id.size()) +
-                      " was not found.");
+  for (const auto& pair : encryption_key_map_) {
+    if (pair.second->key_id == key_id) {
+      *key = *pair.second;
+      return Status::OK;
+    }
   }
-  *key = *encryption_key_;
-  return Status::OK;
+  return Status(error::INTERNAL_ERROR,
+                "Key for key_id=" + base::HexEncode(&key_id[0], key_id.size()) +
+                    " was not found.");
 }
 
 Status FixedKeySource::GetCryptoPeriodKey(uint32_t crypto_period_index,
                                           const std::string& stream_label,
                                           EncryptionKey* key) {
-  // Create a copy of the key.
-  *key = *encryption_key_;
+  Status status = GetKey(stream_label, key);
+  if (!status.ok())
+    return status;
 
   // A naive key rotation algorithm is implemented here by left rotating the
   // key, key_id and pssh. Note that this implementation is only intended for
@@ -88,51 +102,59 @@ Status FixedKeySource::GetCryptoPeriodKey(uint32_t crypto_period_index,
 }
 
 std::unique_ptr<FixedKeySource> FixedKeySource::Create(
-    const std::vector<uint8_t>& key_id,
-    const std::vector<uint8_t>& key,
-    const std::vector<uint8_t>& pssh_boxes,
-    const std::vector<uint8_t>& iv) {
-  std::unique_ptr<EncryptionKey> encryption_key(new EncryptionKey());
-
-  if (key_id.size() != 16) {
-    LOG(ERROR) << "Invalid key ID size '" << key_id.size()
-               << "', must be 16 bytes.";
-    return std::unique_ptr<FixedKeySource>();
-  }
-  if (key.size() != 16) {
-    // CENC only supports AES-128, i.e. 16 bytes.
-    LOG(ERROR) << "Invalid key size '" << key.size() << "', must be 16 bytes.";
-    return std::unique_ptr<FixedKeySource>();
+    const RawKeyParams& raw_key) {
+  std::vector<ProtectionSystemSpecificInfo> key_system_info;
+  if (!raw_key.pssh.empty()) {
+    if (!ProtectionSystemSpecificInfo::ParseBoxes(
+            raw_key.pssh.data(), raw_key.pssh.size(), &key_system_info)) {
+      LOG(ERROR) << "--pssh argument should be full PSSH boxes.";
+      return std::unique_ptr<FixedKeySource>();
+    }
+  } else {
+    // If there aren't any PSSH boxes given, create one with the common system
+    // ID.
+    key_system_info.resize(1);
+    for (const auto& entry : raw_key.key_map) {
+      const RawKeyParams::KeyInfo& key_pair = entry.second;
+      key_system_info.back().add_key_id(key_pair.key_id);
+    }
+    key_system_info.back().set_system_id(kCommonSystemId,
+                                         arraysize(kCommonSystemId));
+    key_system_info.back().set_pssh_box_version(1);
   }
 
-  encryption_key->key_id = key_id;
-  encryption_key->key = key;
-  encryption_key->iv = iv;
+  EncryptionKeyMap encryption_key_map;
+  for (const auto& entry : raw_key.key_map) {
+    const std::string& drm_label = entry.first;
+    const RawKeyParams::KeyInfo& key_pair = entry.second;
 
-  if (!ProtectionSystemSpecificInfo::ParseBoxes(
-          pssh_boxes.data(), pssh_boxes.size(),
-          &encryption_key->key_system_info)) {
-    LOG(ERROR) << "--pssh argument should be full PSSH boxes.";
-    return std::unique_ptr<FixedKeySource>();
-  }
+    if (key_pair.key_id.size() != 16) {
+      LOG(ERROR) << "Invalid key ID size '" << key_pair.key_id.size()
+                 << "', must be 16 bytes.";
+      return std::unique_ptr<FixedKeySource>();
+    }
+    if (key_pair.key.size() != 16) {
+      // CENC only supports AES-128, i.e. 16 bytes.
+      LOG(ERROR) << "Invalid key size '" << key_pair.key.size()
+                 << "', must be 16 bytes.";
+      return std::unique_ptr<FixedKeySource>();
+    }
 
-  // If there aren't any PSSH boxes given, create one with the common system ID.
-  if (encryption_key->key_system_info.size() == 0) {
-    ProtectionSystemSpecificInfo info;
-    info.add_key_id(encryption_key->key_id);
-    info.set_system_id(kCommonSystemId, arraysize(kCommonSystemId));
-    info.set_pssh_box_version(1);
-
-    encryption_key->key_system_info.push_back(info);
+    std::unique_ptr<EncryptionKey> encryption_key(new EncryptionKey);
+    encryption_key->key_id = key_pair.key_id;
+    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);
   }
 
   return std::unique_ptr<FixedKeySource>(
-      new FixedKeySource(std::move(encryption_key)));
+      new FixedKeySource(std::move(encryption_key_map)));
 }
 
 FixedKeySource::FixedKeySource() {}
-FixedKeySource::FixedKeySource(std::unique_ptr<EncryptionKey> key)
-    : encryption_key_(std::move(key)) {}
+FixedKeySource::FixedKeySource(EncryptionKeyMap&& encryption_key_map)
+    : encryption_key_map_(std::move(encryption_key_map)) {}
 
 }  // namespace media
 }  // namespace shaka
diff --git a/packager/media/base/fixed_key_source.h b/packager/media/base/fixed_key_source.h
index f7b1c28148..7d263b3058 100644
--- a/packager/media/base/fixed_key_source.h
+++ b/packager/media/base/fixed_key_source.h
@@ -12,6 +12,7 @@
 #include <vector>
 
 #include "packager/media/base/key_source.h"
+#include "packager/media/public/crypto_params.h"
 
 namespace shaka {
 namespace media {
@@ -47,30 +48,22 @@ class FixedKeySource : public KeySource {
   /// @}
 
   /// Creates a new FixedKeySource from the given data.  Returns null
-  /// if the strings are invalid.
-  /// @param key_id is the key identifier. Must be 16 bytes.
-  /// @param key is the encryption / decryption key. Must be 16 bytes.
-  /// @param pssh_boxes is the concatenated pssh boxes.
-  /// @param iv is the initialization vector. If not specified, a randomly
-  ///        generated IV with the default length will be used.
-  /// Note: GetKey on the created key source will always return the same key for
-  /// all track types.
-  static std::unique_ptr<FixedKeySource> Create(
-      const std::vector<uint8_t>& key_id,
-      const std::vector<uint8_t>& key,
-      const std::vector<uint8_t>& pssh_boxes,
-      const std::vector<uint8_t>& iv);
+  /// if the parameter is malformed.
+  /// @param raw_key contains parameters to setup the key source.
+  static std::unique_ptr<FixedKeySource> Create(const RawKeyParams& raw_key);
 
  protected:
   // Allow default constructor for mock key sources.
   FixedKeySource();
 
  private:
-  explicit FixedKeySource(std::unique_ptr<EncryptionKey> key);
+  typedef std::map<std::string, std::unique_ptr<EncryptionKey>>
+      EncryptionKeyMap;
+  explicit FixedKeySource(EncryptionKeyMap&& encryption_key_map);
+  FixedKeySource(const FixedKeySource&) = delete;
+  FixedKeySource& operator=(const FixedKeySource&) = delete;
 
-  std::unique_ptr<EncryptionKey> encryption_key_;
-
-  DISALLOW_COPY_AND_ASSIGN(FixedKeySource);
+  EncryptionKeyMap encryption_key_map_;
 };
 
 }  // namespace media
diff --git a/packager/media/base/fixed_key_source_unittest.cc b/packager/media/base/fixed_key_source_unittest.cc
index eb48065053..a0ab4ad37d 100644
--- a/packager/media/base/fixed_key_source_unittest.cc
+++ b/packager/media/base/fixed_key_source_unittest.cc
@@ -22,82 +22,123 @@ namespace shaka {
 namespace media {
 
 namespace {
-const char kKeyIdHex[] =   "0101020305080d1522375990e9000000";
-const char kKeyHex[] =     "00100100200300500801302103405500";
+const char kKeyIdHex[] = "0101020305080d1522375990e9000000";
+const char kKeyHex[] = "00100100200300500801302103405500";
+const char kKeyId2Hex[] = "1111121315180d1522375990e9000000";
+const char kKey2Hex[] = "10201110300300500801302103405500";
 const char kIvHex[] = "000102030405060708090a0b0c0d0e0f";
+// PSSH boxes generated manually according to PSSH box syntax specified in
+// ISO/IEC 23001-7:2016 8.1.2.
 const char kPsshBox1Hex[] =
-  "000000427073736801000000"
-  "020305070b0d1113171d1f25292b2f35"
-  "00000001"
-  "0101020305080d1522375990e9000000"
-  "0000000e"
-  "6544617368207061636b61676572";
+    "000000427073736801000000"
+    "020305070b0d1113171d1f25292b2f35"
+    "00000001"
+    "0101020305080d1522375990e9000000"
+    "0000000e"
+    "6544617368207061636b61676572";
 const char kPsshBox2Hex[] =
-  "0000002e7073736800000000"
-  "bbbbbbbbbaaaaaaaaaaaaddddddddddd"
-  "0000000e"
-  "fffffffff000000000000ddddddd";
+    "0000002e7073736800000000"
+    "bbbbbbbbbaaaaaaaaaaaaddddddddddd"
+    "0000000e"
+    "fffffffff000000000000ddddddd";
 const char kDefaultPsshBoxHex[] =
-  "000000347073736801000000"
-  "1077efecc0b24d02ace33c1e52e2fb4b"
-  "00000001"
-  "0101020305080d1522375990e9000000"
-  "00000000";
+    "000000447073736801000000"
+    "1077efecc0b24d02ace33c1e52e2fb4b"
+    "00000002"
+    "1111121315180d1522375990e9000000"
+    "0101020305080d1522375990e9000000"
+    "00000000";
+const char kDrmLabel[] = "SomeDrmLabel";
+const char kAnotherDrmLabel[] = "AnotherDrmLabel";
+const char kEmptyDrmLabel[] = "";
 
 std::vector<uint8_t> HexStringToVector(const std::string& str) {
   std::vector<uint8_t> vec;
   CHECK(base::HexStringToBytes(str, &vec));
   return vec;
 }
-}
+}  // namespace
 
 TEST(FixedKeySourceTest, Success) {
-  std::string pssh_boxes = std::string(kPsshBox1Hex) + kPsshBox2Hex;
-  std::unique_ptr<FixedKeySource> key_source = FixedKeySource::Create(
-      HexStringToVector(kKeyIdHex), HexStringToVector(kKeyHex),
-      HexStringToVector(pssh_boxes), HexStringToVector(kIvHex));
+  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[kEmptyDrmLabel].key_id = HexStringToVector(kKeyId2Hex);
+  raw_key_params.key_map[kEmptyDrmLabel].key = HexStringToVector(kKey2Hex);
+  raw_key_params.iv = HexStringToVector(kIvHex);
+  raw_key_params.pssh =
+      HexStringToVector(std::string(kPsshBox1Hex) + kPsshBox2Hex);
+  std::unique_ptr<FixedKeySource> key_source =
+      FixedKeySource::Create(raw_key_params);
   ASSERT_NE(nullptr, key_source);
 
-  EncryptionKey key;
-  ASSERT_OK(key_source->GetKey("SomeStreamLabel", &key));
+  EncryptionKey key_from_drm_label;
+  ASSERT_OK(key_source->GetKey(kDrmLabel, &key_from_drm_label));
+  EXPECT_HEX_EQ(kKeyIdHex, key_from_drm_label.key_id);
+  EXPECT_HEX_EQ(kKeyHex, key_from_drm_label.key);
+  EXPECT_HEX_EQ(kIvHex, key_from_drm_label.iv);
+  ASSERT_EQ(2u, key_from_drm_label.key_system_info.size());
+  EXPECT_HEX_EQ(kPsshBox1Hex,
+                key_from_drm_label.key_system_info[0].CreateBox());
+  EXPECT_HEX_EQ(kPsshBox2Hex,
+                key_from_drm_label.key_system_info[1].CreateBox());
 
-  EXPECT_HEX_EQ(kKeyIdHex, key.key_id);
-  EXPECT_HEX_EQ(kKeyHex, key.key);
-  EXPECT_HEX_EQ(kIvHex, key.iv);
+  // Using Key ID.
+  EncryptionKey key_from_key_id;
+  ASSERT_OK(key_source->GetKey(HexStringToVector(kKeyIdHex), &key_from_key_id));
+  EXPECT_EQ(key_from_key_id.key_id, key_from_drm_label.key_id);
 
-  ASSERT_EQ(2u, key.key_system_info.size());
-  EXPECT_HEX_EQ(kPsshBox1Hex, key.key_system_info[0].CreateBox());
-  EXPECT_HEX_EQ(kPsshBox2Hex, key.key_system_info[1].CreateBox());
+  // |kAnotherDrmLabel| is not present in the key source, but |kEmptyDrmLabel|
+  // is in the key source, the associated key for |kEmptyDrmLabel| will be
+  // returned.
+  ASSERT_OK(key_source->GetKey(kAnotherDrmLabel, &key_from_drm_label));
+  EXPECT_HEX_EQ(kKeyId2Hex, key_from_drm_label.key_id);
+  EXPECT_HEX_EQ(kKey2Hex, key_from_drm_label.key);
+  EXPECT_HEX_EQ(kIvHex, key_from_drm_label.iv);
+  ASSERT_EQ(2u, key_from_drm_label.key_system_info.size());
+  EXPECT_HEX_EQ(kPsshBox1Hex,
+                key_from_drm_label.key_system_info[0].CreateBox());
+  EXPECT_HEX_EQ(kPsshBox2Hex,
+                key_from_drm_label.key_system_info[1].CreateBox());
 }
 
 TEST(FixedKeySourceTest, EmptyPssh) {
-  std::unique_ptr<FixedKeySource> key_source = FixedKeySource::Create(
-      HexStringToVector(kKeyIdHex), HexStringToVector(kKeyHex),
-      std::vector<uint8_t>(), HexStringToVector(kIvHex));
+  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<FixedKeySource> key_source =
+      FixedKeySource::Create(raw_key_params);
   ASSERT_NE(nullptr, key_source);
 
   EncryptionKey key;
-  ASSERT_OK(key_source->GetKey("SomeStreamLabel", &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].CreateBox());
 }
 
 TEST(FixedKeySourceTest, Failure) {
   // Invalid key id size.
-  std::unique_ptr<FixedKeySource> key_source = FixedKeySource::Create(
-      HexStringToVector("000102030405"), HexStringToVector(kKeyHex),
-      HexStringToVector(kPsshBox1Hex), HexStringToVector(kIvHex));
+  RawKeyParams raw_key_params;
+  raw_key_params.key_map[kEmptyDrmLabel].key_id =
+      HexStringToVector("000102030405");
+  raw_key_params.key_map[kEmptyDrmLabel].key = HexStringToVector(kKeyHex);
+  raw_key_params.pssh = HexStringToVector(kPsshBox1Hex);
+  raw_key_params.iv = HexStringToVector(kIvHex);
+  std::unique_ptr<FixedKeySource> key_source =
+      FixedKeySource::Create(raw_key_params);
   EXPECT_EQ(nullptr, key_source);
 
   // Invalid pssh box.
-  key_source = FixedKeySource::Create(
-      HexStringToVector(kKeyIdHex), HexStringToVector(kKeyHex),
-      HexStringToVector("000102030405"), HexStringToVector(kIvHex));
+  raw_key_params.key_map[kEmptyDrmLabel].key_id = HexStringToVector(kKeyIdHex);
+  raw_key_params.pssh = HexStringToVector("000102030405");
+  key_source = FixedKeySource::Create(raw_key_params);
   EXPECT_EQ(nullptr, key_source);
 }
 
diff --git a/packager/media/base/protection_system_specific_info.h b/packager/media/base/protection_system_specific_info.h
index 170cccff6a..54c1d04b3e 100644
--- a/packager/media/base/protection_system_specific_info.h
+++ b/packager/media/base/protection_system_specific_info.h
@@ -50,7 +50,6 @@ class ProtectionSystemSpecificInfo {
     system_id_.assign(system_id, system_id + system_id_size);
   }
   void add_key_id(const std::vector<uint8_t>& key_id) {
-    DCHECK_EQ(16u, key_id.size());
     key_ids_.push_back(key_id);
   }
   void clear_key_ids() { key_ids_.clear(); }
@@ -72,4 +71,3 @@ class ProtectionSystemSpecificInfo {
 }  // namespace shaka
 
 #endif  // MEDIA_BASE_PSSH_H_
-
diff --git a/packager/media/public/crypto_params.h b/packager/media/public/crypto_params.h
index 21c4d446fd..61118283a5 100644
--- a/packager/media/public/crypto_params.h
+++ b/packager/media/public/crypto_params.h
@@ -72,7 +72,7 @@ struct WidevineEncryptionParams {
 ///       required. The presence of other parameters may be necessary depends
 ///       on server configuration.
 ///   (2) Provide the raw key directly. Both `key_id` and `key` are required.
-///       We are planning to merge this mode with `RawKeyEncryptionParams`.
+///       We are planning to merge this mode with `RawKeyParams`.
 struct PlayreadyEncryptionParams {
   /// Playready license / key server URL.
   std::string key_server_url;
@@ -95,24 +95,26 @@ struct PlayreadyEncryptionParams {
   std::vector<uint8_t> key;
 };
 
-/// Raw key encryption parameters, i.e. with key parameters provided.
-struct RawKeyEncryptionParams {
+/// Raw key encryption/decryption parameters, i.e. with key parameters provided.
+struct RawKeyParams {
   /// An optional initialization vector. If not provided, a random `iv` will be
   /// generated. Note that this parameter should only be used during testing.
+  /// Not needed for decryption.
   std::vector<uint8_t> iv;
   /// Inject a custom `pssh` or multiple concatenated `psshs`. If not provided,
   /// a common system pssh will be generated.
+  /// Not needed for decryption.
   std::vector<uint8_t> pssh;
 
   using StreamLabel = std::string;
-  struct KeyPair {
+  struct KeyInfo {
     std::vector<uint8_t> key_id;
     std::vector<uint8_t> key;
   };
-  /// Defines the KeyPair for the streams. An empty `StreamLabel` indicates the
-  /// default `KeyPair`, which applies to all the `StreamLabels` not present in
+  /// Defines the KeyInfo for the streams. An empty `StreamLabel` indicates the
+  /// default `KeyInfo`, which applies to all the `StreamLabels` not present in
   /// `key_map`.
-  std::map<StreamLabel, KeyPair> key_map;
+  std::map<StreamLabel, KeyInfo> key_map;
 };
 
 /// Encryption parameters.
@@ -124,7 +126,7 @@ struct EncryptionParams {
   // Only one of the three fields is valid.
   WidevineEncryptionParams widevine;
   PlayreadyEncryptionParams playready;
-  RawKeyEncryptionParams raw_key;
+  RawKeyParams raw_key;
 
   /// Clear lead duration in seconds.
   double clear_lead_in_seconds = 0;
@@ -182,19 +184,6 @@ struct WidevineDecryptionParams {
   WidevineSigner signer;
 };
 
-/// Raw key decryption parameters, i.e. with key parameters provided.
-struct RawKeyDecryptionParams {
-  using StreamLabel = std::string;
-  struct KeyPair {
-    std::vector<uint8_t> key_id;
-    std::vector<uint8_t> key;
-  };
-  /// Defines the KeyPair for the streams. An empty `StreamLabel` indicates the
-  /// default `KeyPair`, which applies to all the `StreamLabels` not present in
-  /// `key_map`.
-  std::map<StreamLabel, KeyPair> key_map;
-};
-
 /// Decryption parameters.
 struct DecryptionParams {
   /// Specifies the key provider, which determines which key provider is used
@@ -203,7 +192,7 @@ struct DecryptionParams {
   KeyProvider key_provider = KeyProvider::kNone;
   // Only one of the two fields is valid.
   WidevineDecryptionParams widevine;
-  RawKeyDecryptionParams raw_key;
+  RawKeyParams raw_key;
 };
 
 }  // namespace shaka
diff --git a/packager/packager.cc b/packager/packager.cc
index af1908b51d..3837ac1347 100644
--- a/packager/packager.cc
+++ b/packager/packager.cc
@@ -425,7 +425,13 @@ std::shared_ptr<MediaHandler> CreateEncryptionHandler(
     encryption_params.protection_scheme = kAppleSampleAesProtectionScheme;
   }
 
-  if (!encryption_params.stream_label_func) {
+  if (!stream.drm_label.empty()) {
+    const std::string& drm_label = stream.drm_label;
+    encryption_params.stream_label_func =
+        [drm_label](const EncryptionParams::EncryptedStreamAttributes&) {
+          return drm_label;
+        };
+  } else if (!encryption_params.stream_label_func) {
     const int kDefaultMaxSdPixels = 768 * 576;
     const int kDefaultMaxHdPixels = 1920 * 1080;
     const int kDefaultMaxUhd1Pixels = 4096 * 2160;
diff --git a/packager/packager.h b/packager/packager.h
index c9df36b983..5d7a9901b2 100644
--- a/packager/packager.h
+++ b/packager/packager.h
@@ -82,6 +82,11 @@ struct StreamDescriptor {
   /// If set to true, the stream will not be encrypted. This is useful, e.g. to
   /// encrypt only video streams.
   bool skip_encryption = false;
+  /// Specifies a custom DRM stream label, which can be a DRM label defined by
+  /// the DRM system. Typically values include AUDIO, SD, HD, UHD1, UHD2. If not
+  /// provided, the DRM stream label is derived from stream type (video, audio),
+  /// resolutions etc.
+  std::string drm_label;
   /// If set to a non-zero value, will generate a trick play / trick mode
   /// stream with frames sampled from the key frames in the original stream.
   /// `trick_play_factor` defines the sampling rate.