From 9d457d8022aa99a5e78882e39395f6bb754eb35e Mon Sep 17 00:00:00 2001
From: mikhail <mikhail@tskau.team>
Date: Sun, 26 May 2024 23:41:56 +0500
Subject: [PATCH] improvement: HLS support for internal streams

---
 package-lock.json              |  4 ++--
 src/modules/stream/internal.js | 40 ++++++++++++++++++++++++++++++++--
 src/modules/stream/manage.js   |  6 -----
 3 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index 9052067e..88004217 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
     "name": "cobalt",
-    "version": "7.14.1",
+    "version": "7.14.2",
     "lockfileVersion": 3,
     "requires": true,
     "packages": {
         "": {
             "name": "cobalt",
-            "version": "7.14.1",
+            "version": "7.14.2",
             "license": "AGPL-3.0",
             "dependencies": {
                 "content-disposition-header": "0.6.0",
diff --git a/src/modules/stream/internal.js b/src/modules/stream/internal.js
index fae23f57..4204b1a2 100644
--- a/src/modules/stream/internal.js
+++ b/src/modules/stream/internal.js
@@ -1,7 +1,9 @@
 import { request } from 'undici';
 import { Readable } from 'node:stream';
 import { assert } from 'console';
+import { createInternalStream } from './manage.js';
 import { getHeaders } from './shared.js';
+import HLS from 'hls-parser';
 
 const CHUNK_SIZE = BigInt(8e6); // 8 MB
 const min = (a, b) => a < b ? a : b;
@@ -73,6 +75,36 @@ async function handleYoutubeStream(streamInfo, res) {
     }
 }
 
+function transformHLSMediaPlaylist(streamInfo, hlsPlaylist) {
+    function generateInternalStreamsOfSegments(segment) {
+        const fullUri = new URL(segment.uri, streamInfo.url).toString();
+        segment.uri = createInternalStream(fullUri, streamInfo);
+
+        return segment;
+    }
+
+    hlsPlaylist.segments =
+        hlsPlaylist.segments.map(generateInternalStreamsOfSegments);
+    hlsPlaylist.prefetchSegments =
+        hlsPlaylist.prefetchSegments.map(generateInternalStreamsOfSegments);
+
+    return hlsPlaylist;
+}
+
+async function handleHLSPlaylist(streamInfo, req, res) {
+    let hlsPlaylist = await req.body.text();
+    hlsPlaylist = HLS.parse(hlsPlaylist);
+
+    // NOTE no processing module is passing the master playlist
+    assert(!hlsPlaylist.isMasterPlaylist);
+
+    hlsPlaylist = transformHLSMediaPlaylist(streamInfo, hlsPlaylist);
+    hlsPlaylist = HLS.stringify(hlsPlaylist);
+
+    res.write(hlsPlaylist);
+    res.end();
+}
+
 export async function internalStream(streamInfo, res) {
     if (streamInfo.service === 'youtube') {
         return handleYoutubeStream(streamInfo, res);
@@ -97,8 +129,12 @@ export async function internalStream(streamInfo, res) {
         if (req.statusCode < 200 || req.statusCode > 299)
             return res.end();
 
-        req.body.pipe(res);
-        req.body.on('error', () => res.end());
+        if (["application/vnd.apple.mpegurl", "audio/mpegurl"].includes(req.headers['content-type'])) {
+            await handleHLSPlaylist(streamInfo, req, res);
+        } else {
+            req.body.pipe(res);
+            req.body.on('error', () => res.end());
+        }
     } catch {
         streamInfo.controller.abort();
     }
diff --git a/src/modules/stream/manage.js b/src/modules/stream/manage.js
index bccf5c80..d21cdd41 100644
--- a/src/modules/stream/manage.js
+++ b/src/modules/stream/manage.js
@@ -108,12 +108,6 @@ export function destroyInternalStream(url) {
 }
 
 function wrapStream(streamInfo) {
-    /* m3u8 links are currently not supported
-     * for internal streams, skip them */
-    if (M3U_SERVICES.includes(streamInfo.service)) {
-        return streamInfo;
-    }
-
     const url = streamInfo.urls;
 
     if (typeof url === 'string') {