From 022f6b96fb6620446fe14e0f6e25b5d81c92aefa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81lvaro=20Velad=20Galv=C3=A1n?= <ladvan91@hotmail.com>
Date: Tue, 20 Jun 2023 22:41:52 +0200
Subject: [PATCH] fix: Get the correct timescale when there are two trak boxes
 (#5327)

Stream used for test it:
https://devstreaming-cdn.apple.com/videos/streaming/examples/bipbop_adv_example_hevc/master.m3u8
---
 lib/media/segment_reference.js |  6 ++---
 lib/media/streaming_engine.js  | 43 ++++++++++++++++++++++++----------
 lib/util/mp4_box_parsers.js    | 32 +++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 16 deletions(-)

diff --git a/lib/media/segment_reference.js b/lib/media/segment_reference.js
index 9b020f3590..924c030df4 100644
--- a/lib/media/segment_reference.js
+++ b/lib/media/segment_reference.js
@@ -30,10 +30,10 @@ shaka.media.InitSegmentReference = class {
    *   segment extends to the end of the resource.
    * @param {null|shaka.extern.MediaQualityInfo=} mediaQuality Information about
    *   the quality of the media associated with this init segment.
-   * @param {number=} timescale
+   * @param {(null|number)=} timescale
    * @param {(null|BufferSource)=} segmentData
    */
-  constructor(uris, startByte, endByte, mediaQuality = null, timescale,
+  constructor(uris, startByte, endByte, mediaQuality = null, timescale = null,
       segmentData = null) {
     /** @type {function():!Array.<string>} */
     this.getUris = uris;
@@ -47,7 +47,7 @@ shaka.media.InitSegmentReference = class {
     /** @const {shaka.extern.MediaQualityInfo|null} */
     this.mediaQuality = mediaQuality;
 
-    /** @type {number|undefined} */
+    /** @type {number|null} */
     this.timescale = timescale;
 
     /** @type {BufferSource|null} */
diff --git a/lib/media/streaming_engine.js b/lib/media/streaming_engine.js
index 865139b252..7cfbe1f659 100644
--- a/lib/media/streaming_engine.js
+++ b/lib/media/streaming_engine.js
@@ -1758,16 +1758,45 @@ shaka.media.StreamingEngine = class {
             const initSegment = await fetchInit;
             this.destroyer_.ensureNotDestroyed();
 
+            let lastTimescale = null;
+            const timescaleMap = new Map();
+
             const parser = new shaka.util.Mp4Parser();
             const Mp4Parser = shaka.util.Mp4Parser;
             parser.box('moov', Mp4Parser.children)
                 .box('trak', Mp4Parser.children)
                 .box('mdia', Mp4Parser.children)
                 .fullBox('mdhd', (box) => {
-                  this.parseMDHD_(reference, box);
+                  goog.asserts.assert(
+                      box.version != null,
+                      'MDHD is a full box and should have a valid version.');
+                  const parsedMDHDBox = shaka.util.Mp4BoxParsers.parseMDHD(
+                      box.reader, box.version);
+                  lastTimescale = parsedMDHDBox.timescale;
+                })
+                .box('hdlr', (box) => {
+                  const parsedHDLR = shaka.util.Mp4BoxParsers.parseHDLR(
+                      box.reader);
+                  switch (parsedHDLR.handlerType) {
+                    case 'soun':
+                      timescaleMap.set(ContentType.AUDIO, lastTimescale);
+                      break;
+                    case 'vide':
+                      timescaleMap.set(ContentType.VIDEO, lastTimescale);
+                      break;
+                  }
+                  lastTimescale = null;
                 })
                 .parse(initSegment);
 
+            if (timescaleMap.has(mediaState.type)) {
+              reference.initSegmentReference.timescale =
+                  timescaleMap.get(mediaState.type);
+            } else if (lastTimescale != null) {
+              // Fallback for segments without HDLR box
+              reference.initSegmentReference.timescale = lastTimescale;
+            }
+
             shaka.log.v1(logPrefix, 'appending init segment');
             const hasClosedCaptions = mediaState.stream.closedCaptions &&
                 mediaState.stream.closedCaptions.size > 0;
@@ -1997,18 +2026,6 @@ shaka.media.StreamingEngine = class {
     }
   }
 
-  /**
-   * Parse MDHD box.
-   * @param {!shaka.media.SegmentReference} reference
-   * @param {!shaka.extern.ParsedBox} box
-   * @private
-   */
-  parseMDHD_(reference, box) {
-    const parsedMDHDBox = shaka.util.Mp4BoxParsers.parseMDHD(
-        box.reader || 0, box.version || 0);
-    reference.initSegmentReference.timescale = parsedMDHDBox.timescale;
-  }
-
   /**
    * Parse PRFT box.
    * @param {!shaka.media.SegmentReference} reference
diff --git a/lib/util/mp4_box_parsers.js b/lib/util/mp4_box_parsers.js
index 49b2d2e318..3deeb5f8ef 100644
--- a/lib/util/mp4_box_parsers.js
+++ b/lib/util/mp4_box_parsers.js
@@ -251,6 +251,24 @@ shaka.util.Mp4BoxParsers = class {
     const codec = shaka.util.Mp4Parser.typeToString(fourcc);
     return {codec};
   }
+
+  /**
+   * Parses a HDLR box.
+   * @param {!shaka.util.DataViewReader} reader
+   * @return {!shaka.util.ParsedHDLRBox}
+   */
+  static parseHDLR(reader) {
+    reader.skip(8); // Skip "pre_defined"
+
+    const data = reader.readBytes(4);
+    let handlerType = '';
+    handlerType += String.fromCharCode(data[0]);
+    handlerType += String.fromCharCode(data[1]);
+    handlerType += String.fromCharCode(data[2]);
+    handlerType += String.fromCharCode(data[3]);
+
+    return {handlerType};
+  }
 };
 
 
@@ -403,3 +421,17 @@ shaka.util.ParsedFRMABox;
  * @exportDoc
  */
 shaka.util.ParsedMP4ABox;
+
+/**
+ * @typedef {{
+ *    handlerType: string
+ * }}
+ *
+ * @property {string} handlerType
+ *   A four-character code that identifies the type of the media handler or
+ *   data handler. For media handlers, this field defines the type of
+ *   data—for example, 'vide' for video data, 'soun' for sound data.
+ *
+ * @exportDoc
+ */
+shaka.util.ParsedHDLRBox;