HLSPlaylistStructure Key Index Tags Refactor

mambaSharedFramework/HLS Models/Playlist Structure/HLSPlaylistStructure.swift

@@ -244,9 +244,11 @@ final class HLSPlaylistStructure: HLSPlaylistStructureInterface {
         do {
             let result = try HLSPlaylistStructureConstructor.generateMediaGroups(fromTags: _tags)
 
-            let mediaSpans = try HLSPlaylistStructureConstructor.generateMediaSpans(fromTags: _tags,
-                                                                                    header: result.header,
-                                                                                    mediaSegmentGroups: result.mediaSegmentGroups)
+            let mediaSpans = (try? HLSPlaylistStructureConstructor.generateMediaSpans(
+                fromTags: _tags,
+                header: result.header,
+                mediaSegmentGroups: result.mediaSegmentGroups
+            )) ?? []
 
             self._header = result.header
             self._mediaSegmentGroups = result.mediaSegmentGroups
@@ -508,6 +510,11 @@ fileprivate struct HLSPlaylistStructureConstructor {
                                                header: TagGroup?,
                                                mediaSegmentGroups: [MediaSegmentTagGroup]) throws -> [TagSpan] {
 
+        // If the playlist contains no segments then there are no spans
+        if mediaSegmentGroups.isEmpty {
+            return []
+        }
+
         var mediaSpans = [TagSpan]()
 
         // handle our only known spannable tag, `EXT-X-KEY`
@@ -546,7 +553,13 @@ fileprivate struct HLSPlaylistStructureConstructor {
 
                 if let startKeyIndex = startKeyIndex, let startKeyTag = startKeyTag {
                     // we are closing out our last key
-                    mediaSpans.append(TagSpan(parentTag: startKeyTag, tagMediaSpan: startKeyIndex...currentIndex - 1))
+                    let spanEnd = currentIndex - 1
+                    if startKeyIndex <= spanEnd {
+                        mediaSpans.append(TagSpan(parentTag: startKeyTag, tagMediaSpan: startKeyIndex...spanEnd))
+                    } else {
+                        assertionFailure("Invalid media span range: \(startKeyIndex)...\(spanEnd)")
+                        throw ParseError.invalidMediaSpanRange(start: startKeyIndex, end: spanEnd)
+                    }
                 }
 
                 startKeyIndex = currentIndex
@@ -559,16 +572,26 @@ fileprivate struct HLSPlaylistStructureConstructor {
 
         // close out our last tag
         if let startKeyIndex = startKeyIndex, let startKeyTag = startKeyTag {
-            mediaSpans.append(TagSpan(parentTag: startKeyTag, tagMediaSpan: startKeyIndex...(currentIndex - 1)))
+            let spanEnd = currentIndex - 1
+            if startKeyIndex <= spanEnd {
+                mediaSpans.append(TagSpan(parentTag: startKeyTag, tagMediaSpan: startKeyIndex...spanEnd))
+            } else {
+                assertionFailure("Invalid final media span range: \(startKeyIndex)...\(spanEnd)")
+                throw ParseError.invalidMediaSpanRange(start: startKeyIndex, end: spanEnd)
+            }
+        }
+
+        // assert if key counts mismatch (footer keys or malformed playlists)
+        if keyCount != keyTags.count {
+            assert(keyCount == keyTags.count, "Warning: generateMediaSpans counted \(keyCount) EXT-X-KEY tags, but found \(keyTags.count). Possibly due to footer-only key tags.")
         }
-
-        assert(keyCount == keyTags.count, "we missed a key tag")
 
         return mediaSpans
     }
 
     private enum ParseError: Error {
         case foundMediaSegmentWithoutDuration(inMediaSequence: MediaSequence)
+        case invalidMediaSpanRange(start: Int, end: Int)
     }
 }
 

mambaTests/HLSMediaSpanTests.swift

@@ -114,4 +114,52 @@ class HLSMediaSpanTests: XCTestCase {
         let hlsString = hlsArray.joined()
         runTest(hlsString: hlsString, expectedSpans: [0...1, 2...4, 5...8])
     }
+
+    // This validates the early return logic in generateMediaSpans() for empty mediaSegmentGroups.
+    func testNoMediaSegmentsScenario() {
+        let hlsArray = [
+            """
+            #EXTM3U\n,
+            #EXT-X-TARGETDURATION:6\n,
+            #EXT-X-VERSION:3\n,
+            #EXT-X-MEDIA-SEQUENCE:0\n,
+            #EXT-X-PLAYLIST-TYPE:VOD\n,
+            #EXT-X-KEY:METHOD=NONE\n,
+            #EXT-X-MAP:URI=\"test.mp4\",BYTERANGE=\"610@0\"\n
+            """
+        ]
+        let hlsString = hlsArray.joined()
+        runTest(hlsString: hlsString, expectedSpans: [])
+    }
+
+    // Covers an edge case crash where an EXT-X-KEY appears in the header, but the playlist has no media segments.
+    func testKeyInHeaderWithNoMediaSegmentsDoesNotCrash() {
+        let hlsArray = [
+            """
+            #EXTM3U\n,
+            #EXT-X-VERSION:3\n,
+            #EXT-X-KEY:METHOD=AES-128,URI=\"enc.key\"\n,
+            #EXT-X-ENDLIST\n
+            """
+        ]
+        let hlsString = hlsArray.joined()
+        runTest(hlsString: hlsString, expectedSpans: [])
+    }
+
+    // Validates that an EXT-X-KEY tag appearing after the last media segment (in the footer) does not crash generateMediaSpans() or create invalid spans. This seems to occur with DAI
+    func testFooterOnlyKeyDoesNotCrashOrAppend() {
+        let hlsArray = [
+            """
+            #EXTM3U\n,
+            #EXT-X-VERSION:3\n,
+            #EXT-X-TARGETDURATION:6\n,
+            #EXTINF:6.0,\n,
+            segment1.ts\n,
+            #EXT-X-KEY:METHOD=AES-128,URI=\"footer.key\"\n,
+            #EXT-X-ENDLIST\n
+            """
+        ]
+        let hlsString = hlsArray.joined()
+        runTest(hlsString: hlsString, expectedSpans: [])
+    }
 }