From 42784a9fe47b61adaa287275b2636b863bbbccf1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81lvaro=20Velad=20Galv=C3=A1n?= <ladvan91@hotmail.com>
Date: Mon, 25 Sep 2023 12:24:11 +0200
Subject: [PATCH] fix(HLS): Support AES-128 in init segment according the RFC
 (#5677)

Fixes https://github.com/shaka-project/shaka-player/issues/5667

Backported to v4.4.x
---
 lib/hls/hls_parser.js       | 41 +++++++++++++++++++++++++++++--------
 test/hls/hls_parser_unit.js |  2 ++
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/lib/hls/hls_parser.js b/lib/hls/hls_parser.js
index cc526dc3c76..623cc074a00 100644
--- a/lib/hls/hls_parser.js
+++ b/lib/hls/hls_parser.js
@@ -856,7 +856,7 @@ shaka.hls.HlsParser = class {
 
     let segmentUris = [firstSegment.absoluteUri];
     const initSegmentRef = this.getInitSegmentReference_(
-        playlist.absoluteUri, firstSegment.tags, new Map());
+        playlist, firstSegment.tags, new Map());
     if (initSegmentRef) {
       segmentUris = initSegmentRef.getUris();
     }
@@ -2819,13 +2819,13 @@ shaka.hls.HlsParser = class {
 
   /**
    * Get the InitSegmentReference for a segment if it has a EXT-X-MAP tag.
-   * @param {string} playlistUri The absolute uri of the media playlist.
+   * @param {!shaka.hls.Playlist} playlist
    * @param {!Array.<!shaka.hls.Tag>} tags Segment tags
    * @param {!Map.<string, string>} variables
    * @return {shaka.media.InitSegmentReference}
    * @private
    */
-  getInitSegmentReference_(playlistUri, tags, variables) {
+  getInitSegmentReference_(playlist, tags, variables) {
     /** @type {?shaka.hls.Tag} */
     const mapTag = shaka.hls.Utils.getFirstTagWithName(tags, 'EXT-X-MAP');
 
@@ -2836,7 +2836,7 @@ shaka.hls.HlsParser = class {
     const verbatimInitSegmentUri = mapTag.getRequiredAttrValue('URI');
     const absoluteInitSegmentUri = this.variableSubstitution_(
         shaka.hls.Utils.constructAbsoluteUri(
-            playlistUri, verbatimInitSegmentUri),
+            playlist.absoluteUri, verbatimInitSegmentUri),
         variables);
 
     const mapTagKey = [
@@ -2844,8 +2844,18 @@ shaka.hls.HlsParser = class {
       mapTag.getAttributeValue('BYTERANGE', ''),
     ].join('-');
     if (!this.mapTagToInitSegmentRefMap_.has(mapTagKey)) {
+      /** @type {shaka.extern.HlsAes128Key|undefined} */
+      let aes128Key = undefined;
+      for (const tag of tags) {
+        if (tag.name == 'EXT-X-KEY') {
+          if (tag.getRequiredAttrValue('METHOD') == 'AES-128' &&
+              tag.id < mapTag.id) {
+            aes128Key = this.parseAES128DrmTag_(tag, playlist);
+          }
+        }
+      }
       const initSegmentRef = this.createInitSegmentReference_(
-          absoluteInitSegmentUri, mapTag);
+          absoluteInitSegmentUri, mapTag, aes128Key);
       this.mapTagToInitSegmentRefMap_.set(mapTagKey, initSegmentRef);
     }
     return this.mapTagToInitSegmentRefMap_.get(mapTagKey);
@@ -2856,10 +2866,11 @@ shaka.hls.HlsParser = class {
    * playlist.
    * @param {string} absoluteInitSegmentUri
    * @param {!shaka.hls.Tag} mapTag EXT-X-MAP
+   * @param {shaka.extern.HlsAes128Key=} aes128Key
    * @return {!shaka.media.InitSegmentReference}
    * @private
    */
-  createInitSegmentReference_(absoluteInitSegmentUri, mapTag) {
+  createInitSegmentReference_(absoluteInitSegmentUri, mapTag, aes128Key) {
     let startByte = 0;
     let endByte = null;
     const byterange = mapTag.getAttributeValue('BYTERANGE');
@@ -2870,12 +2881,26 @@ shaka.hls.HlsParser = class {
       const byteLength = Number(blocks[0]);
       startByte = Number(blocks[1]);
       endByte = startByte + byteLength - 1;
+
+      if (aes128Key) {
+        // MAP segment encrypted with method 'AES-128', when served with
+        // HTTP Range, has the unencrypted size specified in the range.
+        // See: https://tools.ietf.org/html/draft-pantos-hls-rfc8216bis-08#section-6.3.6
+        const length = (endByte + 1) - startByte;
+        if (length % 16) {
+          endByte += (16 - (length % 16));
+        }
+      }
     }
 
     const initSegmentRef = new shaka.media.InitSegmentReference(
         () => [absoluteInitSegmentUri],
         startByte,
-        endByte);
+        endByte,
+        /* mediaQuality= */ null,
+        /* timescale= */ null,
+        /* segmentData= */ null,
+        aes128Key);
     return initSegmentRef;
   }
 
@@ -3273,7 +3298,7 @@ shaka.hls.HlsParser = class {
 
       mediaSequenceToStartTime.set(position, startTime);
 
-      initSegmentRef = this.getInitSegmentReference_(playlist.absoluteUri,
+      initSegmentRef = this.getInitSegmentReference_(playlist,
           item.tags, variables);
 
       // If the stream is low latency and the user has not configured the
diff --git a/test/hls/hls_parser_unit.js b/test/hls/hls_parser_unit.js
index 77b7813a915..f0143d900e0 100644
--- a/test/hls/hls_parser_unit.js
+++ b/test/hls/hls_parser_unit.js
@@ -3125,8 +3125,10 @@ describe('HlsParser', () => {
 
     const firstMp4Segment = mp4AesEncryptionVideo.segmentIndex.get(0);
     expect(firstMp4Segment.hlsAes128Key).toBeDefined();
+    expect(firstMp4Segment.initSegmentReference.hlsAes128Key).toBeDefined();
     const secondMp4Segment = mp4AesEncryptionVideo.segmentIndex.get(1);
     expect(secondMp4Segment.hlsAes128Key).toBeNull();
+    expect(secondMp4Segment.initSegmentReference.hlsAes128Key).toBeDefined();
 
     const tsAesEncryptionVideo = actual.variants[2].video;
     await tsAesEncryptionVideo.createSegmentIndex();