Fix HLS & HTTPS proxy stream with streaming response and proper headers & Yield raw bytes without handling the decoding

mediaflow_proxy/const.py

@@ -7,8 +7,8 @@ SUPPORTED_RESPONSE_HEADERS = [
     "transfer-encoding",
     "last-modified",
     "etag",
-    "server",
-    "date",
+    "cache-control",
+    "expires",
 ]
 
 SUPPORTED_REQUEST_HEADERS = [
@@ -16,8 +16,9 @@ SUPPORTED_REQUEST_HEADERS = [
     "accept-encoding",
     "accept-language",
     "connection",
-    "transfer-encoding",
     "range",
     "if-range",
     "user-agent",
+    "referer",
+    "origin",
 ]

mediaflow_proxy/handlers.py

@@ -1,18 +1,23 @@
 import base64
 import logging
-from ipaddress import ip_address
 
 import httpx
 from fastapi import Request, Response, HTTPException
-from fastapi.responses import StreamingResponse
 from pydantic import HttpUrl
 from starlette.background import BackgroundTask
+from starlette.status import HTTP_206_PARTIAL_CONTENT
 
 from .configs import settings
 from .const import SUPPORTED_RESPONSE_HEADERS
 from .mpd_processor import process_manifest, process_playlist, process_segment
 from .utils.cache_utils import get_cached_mpd, get_cached_init_segment
-from .utils.http_utils import Streamer, DownloadError, download_file_with_retry, request_with_retry
+from .utils.http_utils import (
+    Streamer,
+    DownloadError,
+    download_file_with_retry,
+    request_with_retry,
+    EnhancedStreamingResponse,
+)
 from .utils.m3u8_processor import M3U8Processor
 from .utils.mpd_utils import pad_base64
 
@@ -49,13 +54,18 @@ async def handle_hls_stream_proxy(request: Request, destination: str, headers: d
 
         headers.update({"range": headers.get("range", "bytes=0-")})
         # clean up the headers to only include the necessary headers and remove acl headers
-        response_headers = {
-            k.title(): v for k, v in response.headers.items() if k.lower() in SUPPORTED_RESPONSE_HEADERS
-        }
+        response_headers = {k: v for k, v in response.headers.multi_items() if k in SUPPORTED_RESPONSE_HEADERS}
 
-        return StreamingResponse(
+        if transfer_encoding := response_headers.get("transfer-encoding"):
+            if "chunked" not in transfer_encoding:
+                transfer_encoding += ", chunked"
+        else:
+            transfer_encoding = "chunked"
+        response_headers["transfer-encoding"] = transfer_encoding
+
+        return EnhancedStreamingResponse(
             streamer.stream_content(destination, headers),
-            status_code=206,
+            status_code=HTTP_206_PARTIAL_CONTENT,
             headers=response_headers,
             background=BackgroundTask(streamer.close),
         )
@@ -110,18 +120,22 @@ async def handle_stream_request(method: str, video_url: str, headers: dict):
     try:
         response = await streamer.head(video_url, headers)
         # clean up the headers to only include the necessary headers and remove acl headers
-        response_headers = {
-            k.title(): v for k, v in response.headers.items() if k.lower() in SUPPORTED_RESPONSE_HEADERS
-        }
+        response_headers = {k: v for k, v in response.headers.multi_items() if k in SUPPORTED_RESPONSE_HEADERS}
+        if transfer_encoding := response_headers.get("transfer-encoding"):
+            if "chunked" not in transfer_encoding:
+                transfer_encoding += ", chunked"
+        else:
+            transfer_encoding = "chunked"
+        response_headers["transfer-encoding"] = transfer_encoding
 
         if method == "HEAD":
             await streamer.close()
-            return Response(headers=response_headers, status_code=response.status_code)
+            return Response(headers=response_headers, status_code=HTTP_206_PARTIAL_CONTENT)
         else:
-            return StreamingResponse(
+            return EnhancedStreamingResponse(
                 streamer.stream_content(video_url, headers),
                 headers=response_headers,
-                status_code=206,
+                status_code=HTTP_206_PARTIAL_CONTENT,
                 background=BackgroundTask(streamer.close),
             )
     except httpx.HTTPStatusError as e:
@@ -162,8 +176,9 @@ async def fetch_and_process_m3u8(
             content=processed_content,
             media_type="application/vnd.apple.mpegurl",
             headers={
-                "Content-Disposition": "inline",
-                "Accept-Ranges": "none",
+                "Cache-Control": "no-cache, no-store, must-revalidate",
+                "Pragma": "no-cache",
+                "Expires": "0",
             },
         )
     except httpx.HTTPStatusError as e:

mediaflow_proxy/utils/http_utils.py

@@ -1,9 +1,16 @@
 import logging
+import typing
+from functools import partial
 from urllib import parse
 
+import anyio
 import httpx
 import tenacity
+from fastapi import Response
+from starlette.background import BackgroundTask
+from starlette.concurrency import iterate_in_threadpool
 from starlette.requests import Request
+from starlette.types import Receive, Send, Scope
 from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
 
 from mediaflow_proxy.configs import settings
@@ -84,7 +91,7 @@ class Streamer:
         """
         async with self.client.stream("GET", url, headers=headers, follow_redirects=True) as self.response:
             self.response.raise_for_status()
-            async for chunk in self.response.aiter_bytes():
+            async for chunk in self.response.aiter_raw():
                 yield chunk
 
     async def head(self, url: str, headers: dict):
@@ -266,3 +273,80 @@ def get_proxy_headers(request: Request) -> dict:
     request_headers = {k: v for k, v in request.headers.items() if k in SUPPORTED_REQUEST_HEADERS}
     request_headers.update({k[2:].lower(): v for k, v in request.query_params.items() if k.startswith("h_")})
     return request_headers
+
+
+class EnhancedStreamingResponse(Response):
+    body_iterator: typing.AsyncIterable[typing.Any]
+
+    def __init__(
+        self,
+        content: typing.Union[typing.AsyncIterable[typing.Any], typing.Iterable[typing.Any]],
+        status_code: int = 200,
+        headers: typing.Optional[typing.Mapping[str, str]] = None,
+        media_type: typing.Optional[str] = None,
+        background: typing.Optional[BackgroundTask] = None,
+    ) -> None:
+        if isinstance(content, typing.AsyncIterable):
+            self.body_iterator = content
+        else:
+            self.body_iterator = iterate_in_threadpool(content)
+        self.status_code = status_code
+        self.media_type = self.media_type if media_type is None else media_type
+        self.background = background
+        self.init_headers(headers)
+
+    @staticmethod
+    async def listen_for_disconnect(receive: Receive) -> None:
+        try:
+            while True:
+                message = await receive()
+                if message["type"] == "http.disconnect":
+                    logger.debug("Client disconnected")
+                    break
+        except Exception as e:
+            logger.error(f"Error in listen_for_disconnect: {str(e)}")
+
+    async def stream_response(self, send: Send) -> None:
+        try:
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": self.status_code,
+                    "headers": self.raw_headers,
+                }
+            )
+            async for chunk in self.body_iterator:
+                if not isinstance(chunk, (bytes, memoryview)):
+                    chunk = chunk.encode(self.charset)
+                try:
+                    await send({"type": "http.response.body", "body": chunk, "more_body": True})
+                except (ConnectionResetError, anyio.BrokenResourceError):
+                    logger.info("Client disconnected during streaming")
+                    return
+
+            await send({"type": "http.response.body", "body": b"", "more_body": False})
+        except Exception as e:
+            logger.error(f"Error in stream_response: {str(e)}")
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        async with anyio.create_task_group() as task_group:
+
+            async def wrap(func: typing.Callable[[], typing.Awaitable[None]]) -> None:
+                try:
+                    await func()
+                except ExceptionGroup as e:
+                    if not any(isinstance(exc, anyio.get_cancelled_exc_class()) for exc in e.exceptions):
+                        logger.exception("Error in streaming task")
+                    raise
+                except Exception as e:
+                    if not isinstance(e, anyio.get_cancelled_exc_class()):
+                        logger.exception("Error in streaming task")
+                    raise
+                finally:
+                    task_group.cancel_scope.cancel()
+
+            task_group.start_soon(wrap, partial(self.stream_response, send))
+            await wrap(partial(self.listen_for_disconnect, receive))
+
+        if self.background is not None:
+            await self.background()