From 6b036b9bb17170b67276ee3f7e5595af553a8796 Mon Sep 17 00:00:00 2001
From: Joe Foraci <jforaci@gmail.com>
Date: Thu, 25 Jun 2020 23:37:50 -0400
Subject: [PATCH] Support different IVs for each track (#793)

This is supported by allowing IV to be specified in the "key_info_string".

Fixes #543.
---
 CONTRIBUTORS                                       |  1 +
 docs/source/options/raw_key_encryption_options.rst |  6 +++++-
 docs/source/tutorials/raw_key.rst                  |  5 ++---
 packager/app/packager_main.cc                      | 12 ++++++++++++
 packager/media/base/raw_key_source.cc              |  7 ++++++-
 packager/media/public/crypto_params.h              |  1 +
 6 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/CONTRIBUTORS b/CONTRIBUTORS
index 4571ca25f9..85f29dcd29 100644
--- a/CONTRIBUTORS
+++ b/CONTRIBUTORS
@@ -32,6 +32,7 @@ Evgeny Zajcev <zevlg@yandex.ru>
 Gabe Kopley <gabe@philo.com>
 Haoming Chen <hmchen@google.com>
 Jacob Trimble <modmaker@google.com>
+Joe Foraci <jforaci@gmail.com>
 Joey Parrish <joeyparrish@google.com>
 Kongqun Yang <kqyang@google.com>
 Leandro Moreira <leandro.ribeiro.moreira@gmail.com>
diff --git a/docs/source/options/raw_key_encryption_options.rst b/docs/source/options/raw_key_encryption_options.rst
index 92c66d6823..10050ce2bf 100644
--- a/docs/source/options/raw_key_encryption_options.rst
+++ b/docs/source/options/raw_key_encryption_options.rst
@@ -17,7 +17,7 @@ Raw key encryption options
 
     **key_info_string** is of the form::
 
-        label=<label>:key_id=<key_id>:key=<key>
+        label=<label>:key_id=<key_id>:key=<key>[:iv=<initialization_vector>]
 
     *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
@@ -27,6 +27,10 @@ Raw key encryption options
 
     *key_id* and *key* should be 32-digit hex strings.
 
+    *initialization_vector* is an optional IV with the same format and semantics
+    as the parameter for the *--iv* option below. This is mutually exclusive with
+    that option.
+
 --iv <16-digit or 32-digit hex string>
 
     IV in hex string format. If not specified, a random IV will be generated.
diff --git a/docs/source/tutorials/raw_key.rst b/docs/source/tutorials/raw_key.rst
index 53e9a16584..ac2a20a6a2 100644
--- a/docs/source/tutorials/raw_key.rst
+++ b/docs/source/tutorials/raw_key.rst
@@ -20,7 +20,7 @@ Synopsis
 
 **key_info_string** is of the form::
 
-    label=<label>:key_id=<key_id>:key=<key>
+    label=<label>:key_id=<key_id>:key=<key>[:iv=<initialization_vector>]
 
 Custom PSSH(s) can be provided in *--pssh*. If neither --pssh nor
 --protection_systems is specified, `v1 common PSSH box <https://goo.gl/s8RIhr>`_
@@ -67,9 +67,8 @@ The examples below use the H264 streams created in :doc:`encoding`.
       in=h264_high_1080p_6000.mp4,stream=video,output=h264_1080p.mp4,drm_label=HD \
       --protection_scheme cbcs \
       --enable_raw_key_encryption \
-      --keys label=AUDIO:key_id=f3c5e0361e6654b28f8049c778b23946:key=a4631a153a443df9eed0593043db7519,label=SD:key_id=abba271e8bcf552bbd2e86a434a9a5d9:key=69eaa802a6763af979e8d1940fb88392,label=HD:key_id=6d76f25cb17f5e16b8eaef6bbf582d8e:key=cb541084c99731aef4fff74500c12ead \
+      --keys label=AUDIO:key_id=f3c5e0361e6654b28f8049c778b23946:key=a4631a153a443df9eed0593043db7519:iv=11223344556677889900112233445566,label=SD:key_id=abba271e8bcf552bbd2e86a434a9a5d9:key=69eaa802a6763af979e8d1940fb88392:iv=22334455667788990011223344556677,label=HD:key_id=6d76f25cb17f5e16b8eaef6bbf582d8e:key=cb541084c99731aef4fff74500c12ead:iv=33445566778899001122334455667788 \
       --protection_systems FairPlay \
-      --iv 11223344556677889900112233445566
       --hls_master_playlist_output h264_master.m3u8 \
       --hls_key_uri skd://testAssetID
 
diff --git a/packager/app/packager_main.cc b/packager/app/packager_main.cc
index a6214fe185..9d5c8b7af4 100644
--- a/packager/app/packager_main.cc
+++ b/packager/app/packager_main.cc
@@ -118,6 +118,7 @@ const char kUsage[] =
 const char kDrmLabelLabel[] = "label";
 const char kKeyIdLabel[] = "key_id";
 const char kKeyLabel[] = "key";
+const char kKeyIvLabel[] = "iv";
 
 enum ExitStatus {
   kSuccess = 0,
@@ -207,6 +208,17 @@ bool ParseKeys(const std::string& keys, RawKeyParams* raw_key) {
                  << value_map[kKeyLabel];
       return false;
     }
+    if (!value_map[kKeyIvLabel].empty()) {
+      if (!raw_key->iv.empty()) {
+        LOG(ERROR) << "IV already specified with --iv";
+        return false;
+      }
+      if (!base::HexStringToBytes(value_map[kKeyIvLabel], &key_info.iv)) {
+        LOG(ERROR) << "Empty IV or invalid hex string for IV: "
+                   << value_map[kKeyIvLabel];
+        return false;
+      }
+    }
   }
   return true;
 }
diff --git a/packager/media/base/raw_key_source.cc b/packager/media/base/raw_key_source.cc
index 17262df502..0eafd5a1e6 100644
--- a/packager/media/base/raw_key_source.cc
+++ b/packager/media/base/raw_key_source.cc
@@ -113,12 +113,17 @@ std::unique_ptr<RawKeySource> RawKeySource::Create(
                  << "', must be 16 bytes.";
       return std::unique_ptr<RawKeySource>();
     }
+    if (!key_pair.iv.empty() && key_pair.iv.size() != 8 && key_pair.iv.size() != 16) {
+      LOG(ERROR) << "Invalid IV '" << key_pair.iv.size()
+                 << "', must be 8 or 16 bytes.";
+      return std::unique_ptr<RawKeySource>();
+    }
 
     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->iv = (key_pair.iv.empty()) ? raw_key.iv : key_pair.iv;
     encryption_key->key_system_info = key_system_info;
     encryption_key_map[drm_label] = std::move(encryption_key);
   }
diff --git a/packager/media/public/crypto_params.h b/packager/media/public/crypto_params.h
index e63109fef9..0e66a8e191 100644
--- a/packager/media/public/crypto_params.h
+++ b/packager/media/public/crypto_params.h
@@ -134,6 +134,7 @@ struct RawKeyParams {
   struct KeyInfo {
     std::vector<uint8_t> key_id;
     std::vector<uint8_t> key;
+    std::vector<uint8_t> iv;
   };
   /// Defines the KeyInfo for the streams. An empty `StreamLabel` indicates the
   /// default `KeyInfo`, which applies to all the `StreamLabels` not present in