fix(HLS): Avoid duplicate AES request when using the same info (#6118)

shaka-project · Jan 17, 2024 · 1671a3e · 1671a3e
1 parent af12b0b
commit 1671a3e
Showing 1 changed file with 53 additions and 45 deletions.
diff --git a/lib/hls/hls_parser.js b/lib/hls/hls_parser.js
@@ -216,6 +216,9 @@ shaka.hls.HlsParser = class {
      * {!Map.<string, !shaka.media.InitSegmentReference>} */
     this.mapTagToInitSegmentRefMap_ = new Map();
 
+    /** @private {Map.<string, !shaka.extern.aesKey>} */
+    this.aesKeyInfoMap_ = new Map();
+
     /** @private {boolean} */
     this.lowLatencyMode_ = false;
 
@@ -296,6 +299,7 @@ shaka.hls.HlsParser = class {
     this.groupIdToCodecsMap_.clear();
     this.globalVariables_.clear();
     this.mapTagToInitSegmentRefMap_.clear();
+    this.aesKeyInfoMap_.clear();
 
     if (this.contentSteeringManager_) {
       this.contentSteeringManager_.destroy();
@@ -2826,55 +2830,59 @@ shaka.hls.HlsParser = class {
       }
     }
 
-    // Default AES-128
-    const keyInfo = {
-      bitsKey: 128,
-      blockCipherMode: 'CBC',
-      iv,
-      firstMediaSequenceNumber,
-    };
+    const aesKeyInfoKey = `${drmTag.toString()}-${firstMediaSequenceNumber}`;
+    if (!this.aesKeyInfoMap_.has(aesKeyInfoKey)) {
+      // Default AES-128
+      const keyInfo = {
+        bitsKey: 128,
+        blockCipherMode: 'CBC',
+        iv,
+        firstMediaSequenceNumber,
+      };
 
-    const method = drmTag.getRequiredAttrValue('METHOD');
-    switch (method) {
-      case 'AES-256':
-        keyInfo.bitsKey = 256;
-        break;
-      case 'AES-256-CTR':
-        keyInfo.bitsKey = 256;
-        keyInfo.blockCipherMode = 'CTR';
-        break;
-    }
-
-    // Don't download the key object until the segment is parsed, to avoid a
-    // startup delay for long manifests with lots of keys.
-    keyInfo.fetchKey = async () => {
-      const keyUri = shaka.hls.Utils.constructSegmentUris(
-          getUris(), drmTag.getRequiredAttrValue('URI'), variables);
-
-      const requestType = shaka.net.NetworkingEngine.RequestType.KEY;
-      const request = shaka.net.NetworkingEngine.makeRequest(
-          [keyUri], this.config_.retryParameters);
-      const keyResponse = await this.makeNetworkRequest_(request, requestType);
-
-      // keyResponse.status is undefined when URI is "data:text/plain;base64,"
-      if (!keyResponse.data ||
-          keyResponse.data.byteLength != (keyInfo.bitsKey / 8)) {
-        throw new shaka.util.Error(
-            shaka.util.Error.Severity.CRITICAL,
-            shaka.util.Error.Category.MANIFEST,
-            shaka.util.Error.Code.AES_128_INVALID_KEY_LENGTH);
+      const method = drmTag.getRequiredAttrValue('METHOD');
+      switch (method) {
+        case 'AES-256':
+          keyInfo.bitsKey = 256;
+          break;
+        case 'AES-256-CTR':
+          keyInfo.bitsKey = 256;
+          keyInfo.blockCipherMode = 'CTR';
+          break;
       }
 
-      const algorithm = {
-        name: keyInfo.blockCipherMode == 'CTR' ? 'AES-CTR' : 'AES-CBC',
-        length: keyInfo.bitsKey,
-      };
-      keyInfo.cryptoKey = await window.crypto.subtle.importKey(
-          'raw', keyResponse.data, algorithm, true, ['decrypt']);
-      keyInfo.fetchKey = undefined; // No longer needed.
-    };
+      // Don't download the key object until the segment is parsed, to avoid a
+      // startup delay for long manifests with lots of keys.
+      keyInfo.fetchKey = async () => {
+        const keyUri = shaka.hls.Utils.constructSegmentUris(
+            getUris(), drmTag.getRequiredAttrValue('URI'), variables);
+
+        const requestType = shaka.net.NetworkingEngine.RequestType.KEY;
+        const request = shaka.net.NetworkingEngine.makeRequest(
+            [keyUri], this.config_.retryParameters);
+        const keyResponse =
+            await this.makeNetworkRequest_(request, requestType);
+
+        // keyResponse.status is undefined when URI is "data:text/plain;base64,"
+        if (!keyResponse.data ||
+            keyResponse.data.byteLength != (keyInfo.bitsKey / 8)) {
+          throw new shaka.util.Error(
+              shaka.util.Error.Severity.CRITICAL,
+              shaka.util.Error.Category.MANIFEST,
+              shaka.util.Error.Code.AES_128_INVALID_KEY_LENGTH);
+        }
 
-    return keyInfo;
+        const algorithm = {
+          name: keyInfo.blockCipherMode == 'CTR' ? 'AES-CTR' : 'AES-CBC',
+          length: keyInfo.bitsKey,
+        };
+        keyInfo.cryptoKey = await window.crypto.subtle.importKey(
+            'raw', keyResponse.data, algorithm, true, ['decrypt']);
+        keyInfo.fetchKey = undefined; // No longer needed.
+      };
+      this.aesKeyInfoMap_.set(aesKeyInfoKey, keyInfo);
+    }
+    return this.aesKeyInfoMap_.get(aesKeyInfoKey);
   }