feat: improve debrid stream proxy

.env-sample

@@ -15,5 +15,4 @@ GET_TORRENT_TIMEOUT=5 # maximum time to obtain the torrent info hash in seconds
 ZILEAN_URL=None # for DMM search - https://github.com/iPromKnight/zilean
 CUSTOM_HEADER_HTML=None # only set it if you know what it is
 PROXY_DEBRID_STREAM=False # Proxy Debrid Streams (very useful to use your debrid service on multiple IPs at same time)
-PROXY_DEBRID_STREAM_PASSWORD=CHANGE_ME # Secret password to enter on configuration page to prevent people from abusing your debrid stream proxy
-PROXY_DEBRID_STREAM_BYTES_PER_CHUNK=102400 # 10MB per chunks
+PROXY_DEBRID_STREAM_PASSWORD=CHANGE_ME # Secret password to enter on configuration page to prevent people from abusing your debrid stream proxy

comet/api/stream.py

@@ -3,9 +3,11 @@
 import json
 import time
 import aiohttp
+import httpx
 
 from fastapi import APIRouter, Request
 from fastapi.responses import RedirectResponse, StreamingResponse
+from starlette.background import BackgroundTask
 from RTN import Torrent, parse, sort_torrents, title_match
 
 from comet.debrid.manager import getDebrid
@@ -310,43 +312,41 @@ async def playback(b64config: str, hash: str, index: str):
 
     return RedirectResponse(download_link, status_code=302)
 
-
 @streams.get("/{b64config}/playback/{hash}/{index}")
 async def playback(request: Request, b64config: str, hash: str, index: str):
     config = config_check(b64config)
     if not config:
         return
-
+    
     async with aiohttp.ClientSession() as session:
         debrid = getDebrid(session, config)
         download_link = await debrid.generate_download_link(hash, index)
-
         if download_link is None:
             return
-
+        
         proxy = (
             debrid.proxy if config["debridService"] == "alldebrid" else None
         )  # proxy is not needed to proxy realdebrid stream
-
+        
         if (
             settings.PROXY_DEBRID_STREAM
             and settings.PROXY_DEBRID_STREAM_PASSWORD
             == config["debridStreamProxyPassword"]
         ):
-
-            async def stream_content(headers: dict):
-                async with aiohttp.ClientSession() as session:
-                    async with session.get(
-                        download_link, headers=headers, proxy=proxy
-                    ) as response:
-                        while True:
-                            chunk = await response.content.read(
-                                settings.PROXY_DEBRID_STREAM_BYTES_PER_CHUNK
-                            )
-                            if not chunk:
-                                break
-                            yield chunk
-
+            class Streamer:
+                def __init__(self):
+                    self.response = None
+
+                async def stream_content(self, headers: dict):
+                    async with httpx.AsyncClient(proxy=proxy) as client:
+                        async with client.stream("GET", download_link, headers=headers) as self.response:
+                            async for chunk in self.response.aiter_raw():
+                                yield chunk
+
+                async def close(self):
+                    if self.response is not None:
+                        await self.response.aclose()
+            
             range = None
             range_header = request.headers.get("range")
             if range_header:
@@ -357,21 +357,20 @@ async def stream_content(headers: dict):
                 range = f"bytes={start}-{end}"
 
             async with await session.get(
-                download_link, headers={"Range": f"bytes={start}-{end}"}, proxy=proxy
+                download_link, headers={"Range": range}, proxy=proxy
             ) as response:
-                await session.close()
-
                 if response.status == 206:
+                    streamer = Streamer()
+
                     return StreamingResponse(
-                        stream_content({"Range": range}),
+                        streamer.stream_content({"Range": range}),
                         status_code=206,
                         headers={
                             "Content-Range": response.headers["Content-Range"],
                             "Content-Length": response.headers["Content-Length"],
                             "Accept-Ranges": "bytes",
-                        },
+                        }, background=BackgroundTask(await streamer.close())
                     )
-
             return
-
-        return RedirectResponse(download_link, status_code=302)
+        
+        return RedirectResponse(download_link, status_code=302)

comet/utils/models.py

@@ -28,7 +28,6 @@ class AppSettings(BaseSettings):
     CUSTOM_HEADER_HTML: Optional[str] = None
     PROXY_DEBRID_STREAM: Optional[bool] = False
     PROXY_DEBRID_STREAM_PASSWORD: Optional[str] = "CHANGE_ME"
-    PROXY_DEBRID_STREAM_BYTES_PER_CHUNK: Optional[int] = 102400
 
 
 settings = AppSettings()

pyproject.toml

@@ -19,6 +19,7 @@ databases = "*"
 aiosqlite = "*"
 pydantic-settings = "*"
 bencode-py = "*"
+httpx = "*"
 
 [tool.poetry.group.dev.dependencies]
 isort = "*"