fix: main.py - persistent pump + async refresh + hls live opts

main.py

@@ -214,38 +214,32 @@ def proxy_m3u_hls(url: str = Query(...), ua: str = Query(default=None), cookie:
             )
         except Exception:
             return Response(status_code=502, content="proxy error")
-    # Raw MPEG-TS / Stalker stream: emit a live HLS wrapper playlist.
-    # FIX: Emit a LIVE sliding window with multiple #EXTINF segments (all
-    # pointing at the same /api/proxy/seg URL). The old single-segment manifest
-    # caused Hls.js to play ~8s of data and then stall — the playlist had no
-    # new segments to discover. With 3 segments Hls.js keeps the playback
-    # buffer alive and polls for the next window, giving continuous live
-    # playback for MPEG-TS / Stalker streams.
-    _seg = seg_proxy_base
-    manifest = (
-        "#EXTM3U\n"
-        "#EXT-X-VERSION:6\n"
-        "#EXT-X-TARGETDURATION:8\n"
-        "#EXT-X-MEDIA-SEQUENCE:0\n"
-        "#EXT-X-DISCONTINUITY-SEQUENCE:0\n"
-        "#EXT-X-START:TIME-OFFSET=0,PRECISE=YES\n"
-        "#EXTINF:8.0,\n"
-        + _seg + "\n"
-        "#EXTINF:8.0,\n"
-        + _seg + "\n"
-        "#EXTINF:8.0,\n"
-        + _seg + "\n"
-        "#EXTINF:8.0,\n"
-        + _seg + "\n"
-        "#EXTINF:8.0,\n"
-        + _seg + "\n"
-        "#EXTINF:8.0,\n"
-        + _seg + "\n"
-        "#EXTINF:8.0,\n"
-        + _seg + "\n"
-        "#EXTINF:8.0,\n"
-        + _seg + "\n"
-    )
+    # Raw MPEG-TS / Stalker / unknown: serve a SLIDING WINDOW from
+    # ONE persistent upstream pump. The old manifest listed 8
+    # identical /api/proxy/seg URLs -> hls.js treated segments as
+    # replayed after ~64s (buffer drains -> giât) and every segment
+    # open a NEW upstream connection (~3.4s reconnect stalls).
+    from urllib.parse import urlencode as _ue
+    _key = _hls_pump_key(url, ua or "", cookie or "", auth or "")
+    _live = _hls_pump_progress(_key, url, headers)
+    _seq = max(0, (_live // _HLS_WINDOW) - _HLS_SEGS)
+    _man_lines = [
+        "#EXTM3U",
+        "#EXT-X-VERSION:6",
+        "#EXT-X-TARGETDURATION:10",
+        "#EXT-X-MEDIA-SEQUENCE:" + str(_seq),
+        "#EXT-X-DISCONTINUITY-SEQUENCE:0",
+        "#EXT-X-START:TIME-OFFSET=0,PRECISE=YES",
+    ]
+    for _i in range(_HLS_SEGS):
+        _q = _ue({"url": url})
+        if ua: _q += "&ua=" + quote(ua, safe="")
+        if cookie: _q += "&cookie=" + quote(cookie, safe="")
+        if auth: _q += "&auth=" + quote(auth, safe="")
+        _q += "&off=" + str((_seq + _i) * _HLS_WINDOW)
+        _man_lines.append("#EXTINF:8.0,")
+        _man_lines.append("/api/proxy/seg?" + _q)
+    manifest = "\n".join(_man_lines) + "\n"
     return Response(content=manifest.encode("utf-8"),
                     media_type="application/vnd.apple.mpegurl",
                     headers={"Access-Control-Allow-Origin":"*","Cache-Control":"no-cache, no-store, must-revalidate"})
@@ -311,25 +305,172 @@ def proxy_stream(url: str = Query(...), ua: str = Query(default=None), cookie: s
     except Exception:
         return Response(status_code=502, content="proxy error")
 
+_HLS_PUMPS = {}
+_HLS_PUMPS_LOCK = threading.Lock()
+_HLS_WINDOW = 4194280
+_HLS_SEGS = 4
+_HLS_KEEP = 8 * _HLS_WINDOW
+
+
+def _hls_pump_key(url, ua, cookie, auth):
+    return (url or "") + "\x1f" + (ua or "") + "\x1f" + (cookie or "") + "\x1f" + (auth or "")
+
+
+class _UpstreamPump:
+    """Hold ONE connection to a raw MPEG-TS / mp2t / .ts stream and buffer it
+    as a monotonic byte stream. Segment requests stream byte ranges from this
+    buffer, so N consecutive HLS segments share ONE upstream connection —
+    eliminating the per-segment reconnect that caused the giât on restreamed
+    .ts/.mp2t channels."""
+
+    def __init__(self, url, headers):
+        self.url = url
+        self.headers = headers
+        self.buf = bytearray()
+        self.base = 0
+        self.live = 0
+        self.dead = False
+        self.err = ""
+        self.last_use = time.time()
+        self.cv = threading.Condition()
+        self._t = threading.Thread(target=self._run, daemon=True)
+
+    def start(self):
+        self._t.start()
+
+    def _run(self):
+        rr = None
+        try:
+            rr = requests.get(self.url, headers=self.headers, timeout=(10, 120),
+                              stream=True, allow_redirects=True)
+            if rr.status_code != 200:
+                with self.cv:
+                    self.err = "upstream %s" % rr.status_code
+                    self.dead = True
+                    self.cv.notify_all()
+                return
+            for chunk in rr.iter_content(chunk_size=65536):
+                if not chunk:
+                    continue
+                with self.cv:
+                    self.buf.extend(chunk)
+                    self.live += len(chunk)
+                    if len(self.buf) > _HLS_KEEP:
+                        drop = len(self.buf) - _HLS_KEEP
+                        del self.buf[:drop]
+                        self.base += drop
+                    self.cv.notify_all()
+        except Exception as e:
+            self.err = str(e)[:120]
+        finally:
+            if rr is not None:
+                try:
+                    rr.close()
+                except Exception:
+                    pass
+            with self.cv:
+                self.dead = True
+                self.cv.notify_all()
+
+    def stream(self, off, want):
+        """Yield bytes from absolute offset `off`, up to `want` bytes total.
+        Blocks waiting for data; ends at EOF/error or after `want` bytes."""
+        self.last_use = time.time()
+        sent = 0
+        while sent < want:
+            with self.cv:
+                if off < self.base:
+                    off = self.base
+                have = self.live - off
+                if have <= 0:
+                    if self.dead:
+                        return
+                    self.cv.wait(0.5)
+                    continue
+                n = min(want - sent, have)
+                start = off - self.base
+                data = bytes(self.buf[start:start + n])
+                off += n
+            if not data:
+                return
+            sent += len(data)
+            yield data
+
+
+def _hls_get_pump(key, url, headers):
+    with _HLS_PUMPS_LOCK:
+        p = _HLS_PUMPS.get(key)
+        if p is None or p.dead:
+            p = _UpstreamPump(url, headers)
+            p.start()
+            _HLS_PUMPS[key] = p
+        return p
+
+
+def _hls_pump_progress(key, url, headers):
+    with _HLS_PUMPS_LOCK:
+        p = _HLS_PUMPS.get(key)
+    if p is None:
+        p = _hls_get_pump(key, url, headers)
+    return getattr(p, "live", 0)
+
+
+def _hls_pump_janitor():
+    while True:
+        time.sleep(300)
+        try:
+            now = time.time()
+            with _HLS_PUMPS_LOCK:
+                for k in list(_HLS_PUMPS):
+                    p = _HLS_PUMPS[k]
+                    if p.dead or now - p.last_use > 900:
+                        try:
+                            p.buf = bytearray()
+                        except Exception:
+                            pass
+                        del _HLS_PUMPS[k]
+        except Exception:
+            pass
+
+
+try:
+    threading.Thread(target=_hls_pump_janitor, daemon=True).start()
+except Exception:
+    pass
+
+
 @app.get("/api/proxy/seg")
-def proxy_segment(url: str = Query(...), ua: str = Query(default=None), cookie: str = Query(default=None), auth: str = Query(default=None), fallbacks: str = Query(default=None), max_sec: float = Query(default=0)):
-    """Proxy a stream segment. For live MPEG-TS streams, streams the raw data
-    continuously so Hls.js can buffer and play it. If the fetched URL turns
-    out to be an HLS *playlist* (master or media), it is rewritten so the
-    downstream segment(s) also go through this proxy. The max_sec parameter
-    is accepted for backwards compatibility but the stream is left open-ended
-    (the client / CDN closes the read when sufficient data is buffered or the
-    user navigates away)."""
+def proxy_segment(url: str = Query(...), ua: str = Query(default=None), cookie: str = Query(default=None), auth: str = Query(default=None), fallbacks: str = Query(default=None), max_sec: float = Query(default=0), off: int = Query(default=-1)):
+    """Proxy a stream segment.
+
+    - off >= 0: raw-stream path served from the shared pump (ONE persistent
+      upstream connection). Consecutive HLS windows read continuous byte ranges
+      from the same connection -> no reconnect giât on restreamed .ts/.mp2t.
+    - off < 0: legacy path. Probes the upstream; HLS playlists (master/media)
+      are rewritten so their segments also go through this proxy.
+    """
+    seg_headers = dict(HEADERS)
+    if ua: seg_headers["User-Agent"] = ua
+    if cookie: seg_headers["Cookie"] = cookie
+    if auth: seg_headers["Authorization"] = auth
+
+    if off >= 0:
+        key = _hls_pump_key(url, ua, cookie, auth)
+        p = _hls_get_pump(key, url, seg_headers)
+
+        def _gen():
+            try:
+                yield from p.stream(off, _HLS_WINDOW)
+            except Exception:
+                return
+        return StreamingResponse(
+            _gen(),
+            media_type="video/mp2t",
+            headers={"Access-Control-Allow-Origin": "*", "Cache-Control": "no-cache, no-store"},
+        )
+
+    # ---------- legacy sniff path (HLS playlists / fallbacks) ----------
     try:
-        seg_headers = dict(HEADERS)
-        if ua: seg_headers["User-Agent"] = ua
-        if cookie: seg_headers["Cookie"] = cookie
-        if auth: seg_headers["Authorization"] = auth
-        # FIX: try the primary URL first; if it returns non-200 (e.g. an
-        # expired stalker play_token -> 502/403), fall through the "fallbacks"
-        # list (same stream ID on alternate servers) and use the first that
-        # returns 200. This fixes "black screen" on Stalker channels whose
-        # cached token on one server expired but works on another.
         import json as _json
         fb_urls = []
         if fallbacks:
@@ -343,11 +484,6 @@ def proxy_segment(url: str = Query(...), ua: str = Query(default=None), cookie:
             try:
                 r = requests.get(attempt_url, headers=seg_headers, timeout=(8, 300), stream=True)
                 if r.status_code == 200:
-                    # Sniff the first bytes: if it's an HLS playlist (e.g. a
-                    # master playlist from "/api/proxy/m3u_hls" for an .m3u8
-                    # URL, or a media playlist variant), rewrite it so every
-                    # segment goes through this proxy again. MPEG-TS (0x47
-                    # sync bytes) or other media streams pass through raw.
                     lead = b""
                     try:
                         lead = next(r.iter_content(4096))
@@ -362,16 +498,13 @@ def proxy_segment(url: str = Query(...), ua: str = Query(default=None), cookie:
                                 body = rr.text
                         except Exception:
                             pass
-                        # Resolve against the playlist location (some variants
-                        # use relative URIs) — use the attempt URL as base.
                         rewritten, _ = rewrite_hls_playlist(body, attempt_url, ua, cookie, auth, None)
                         return Response(
                             content=rewritten.encode("utf-8"),
                             media_type="application/vnd.apple.mpegurl",
                             headers={"Access-Control-Allow-Origin":"*", "Cache-Control":"no-cache, no-store, must-revalidate"},
                         )
-                    # Real media segment — stream through.
-                    def _stream(src=r):
+                    def _stream(src=r, lead=lead):
                         if lead:
                             yield lead
                         for chunk in src.iter_content(chunk_size=65536):
@@ -390,6 +523,8 @@ def proxy_segment(url: str = Query(...), ua: str = Query(default=None), cookie:
     except Exception:
         return Response(status_code=502, content="proxy error")
 
+
+
 @app.get("/api/proxy/video")
 def proxy_video(url: str = Query(...), request: Request = None):
     try: