shared/video_file_server.py

"""Threaded HTTP server for serving video files on a separate port.

Gradio's SSE connections consume Chrome's 6-per-origin HTTP/1.1 connection
budget. When multiple tabs are open, file requests to Gradio's ``/file=``
route get queued behind active SSE connections, causing videos to appear
"empty" until other tabs finish processing.

This module serves video files on a different port (different origin),
giving the browser a separate connection pool. Gradio's Video component
receives an HTTP URL and passes it directly to ``<video src=...>`` without
any postprocessing (no ffprobe, no caching, no hashing).
"""

from __future__ import annotations

import http.server
import logging
import mimetypes
import os
import threading
from pathlib import Path


class _VideoHandler(http.server.BaseHTTPRequestHandler):
    """Serves files under *root_dir* with range-request and CORS support."""

    # Set by VideoFileServer before binding
    root_dir: str = ""
    file_logger: logging.Logger | None = None

    # ------------------------------------------------------------------
    # helpers
    # ------------------------------------------------------------------

    def _resolve_path(self) -> Path | None:
        """Map the URL path to a file under *root_dir*, rejecting traversal."""
        # Strip leading slash
        rel = self.path.lstrip("/")
        if not rel:
            return None
        try:
            resolved = (Path(self.root_dir) / rel).resolve()
        except (OSError, ValueError):
            return None
        # Must be under root_dir
        root = Path(self.root_dir).resolve()
        try:
            resolved.relative_to(root)
        except ValueError:
            return None
        if not resolved.is_file():
            return None
        return resolved

    def _send_cors_headers(self) -> None:
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, HEAD, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Range")
        self.send_header(
            "Access-Control-Expose-Headers", "Content-Range, Content-Length, Accept-Ranges"
        )

    # ------------------------------------------------------------------
    # HTTP methods
    # ------------------------------------------------------------------

    def do_OPTIONS(self) -> None:  # noqa: N802
        self.send_response(204)
        self._send_cors_headers()
        self.end_headers()

    def do_HEAD(self) -> None:  # noqa: N802
        path = self._resolve_path()
        if path is None:
            self.send_error(404)
            return
        file_size = path.stat().st_size
        content_type = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
        self.send_response(200)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(file_size))
        self.send_header("Accept-Ranges", "bytes")
        self._send_cors_headers()
        self.end_headers()

    def do_GET(self) -> None:  # noqa: N802
        path = self._resolve_path()
        if path is None:
            self.send_error(404)
            return

        file_size = path.stat().st_size
        content_type = mimetypes.guess_type(str(path))[0] or "application/octet-stream"

        range_header = self.headers.get("Range")
        if range_header:
            self._serve_range(path, file_size, content_type, range_header)
        else:
            self._serve_full(path, file_size, content_type)

    # ------------------------------------------------------------------
    # response helpers
    # ------------------------------------------------------------------

    def _serve_full(self, path: Path, file_size: int, content_type: str) -> None:
        self.send_response(200)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(file_size))
        self.send_header("Accept-Ranges", "bytes")
        self._send_cors_headers()
        self.end_headers()
        self._stream_bytes(path, 0, file_size - 1)

    def _serve_range(
        self, path: Path, file_size: int, content_type: str, range_header: str
    ) -> None:
        try:
            start, end = self._parse_range(range_header, file_size)
        except ValueError:
            self.send_response(416)
            self.send_header("Content-Range", f"bytes */{file_size}")
            self._send_cors_headers()
            self.end_headers()
            return

        length = end - start + 1
        self.send_response(206)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(length))
        self.send_header("Content-Range", f"bytes {start}-{end}/{file_size}")
        self.send_header("Accept-Ranges", "bytes")
        self._send_cors_headers()
        self.end_headers()
        self._stream_bytes(path, start, end)

    def _stream_bytes(self, path: Path, start: int, end: int) -> None:
        chunk_size = 64 * 1024
        remaining = end - start + 1
        try:
            with open(path, "rb") as f:
                f.seek(start)
                while remaining > 0:
                    chunk = f.read(min(chunk_size, remaining))
                    if not chunk:
                        break
                    self.wfile.write(chunk)
                    remaining -= len(chunk)
        except (BrokenPipeError, ConnectionResetError):
            pass

    @staticmethod
    def _parse_range(header: str, file_size: int) -> tuple[int, int]:
        """Parse ``Range: bytes=start-end`` and return (start, end) inclusive."""
        if not header.startswith("bytes="):
            raise ValueError(header)
        spec = header[len("bytes=") :]
        parts = spec.split("-", 1)
        if len(parts) != 2:
            raise ValueError(header)

        start_s, end_s = parts
        if start_s == "":
            # Suffix range: bytes=-500 means last 500 bytes
            suffix = int(end_s)
            if suffix <= 0 or suffix > file_size:
                raise ValueError(header)
            return file_size - suffix, file_size - 1

        start = int(start_s)
        end = int(end_s) if end_s else file_size - 1
        if start < 0 or start >= file_size or end < start or end >= file_size:
            raise ValueError(header)
        return start, end

    # Suppress default stderr logging — route through our logger
    def log_message(self, format: str, *args: object) -> None:
        if self.file_logger:
            self.file_logger.debug(format, *args)


class VideoFileServer:
    """Threaded HTTP file server for video output delivery.

    Args:
        root_dir: Directory tree to serve files from. Only files under this
            directory (resolved, no symlink traversal) are accessible.
        port: TCP port to listen on. ``0`` picks an available port.
        logger: Optional logger for request logging.
    """

    def __init__(
        self,
        root_dir: str,
        port: int | None = None,
        logger: logging.Logger | None = None,
    ) -> None:
        if port is None:
            port = int(os.environ.get("VIDEO_SERVER_PORT", "7861"))
        self._root_dir = str(Path(root_dir).resolve())
        self._logger = logger

        # Create a handler subclass bound to this server's config
        handler_cls = type(
            "_BoundVideoHandler",
            (_VideoHandler,),
            {"root_dir": self._root_dir, "file_logger": logger},
        )
        self._httpd = http.server.ThreadingHTTPServer(("0.0.0.0", port), handler_cls)
        self._thread: threading.Thread | None = None

    @property
    def port(self) -> int:
        return self._httpd.server_address[1]

    def start(self) -> None:
        """Start serving in a daemon thread."""
        self._thread = threading.Thread(target=self._httpd.serve_forever, daemon=True)
        self._thread.start()
        if self._logger:
            self._logger.info("Video file server started on port %d", self.port)

    def stop(self) -> None:
        """Shut down the server."""
        self._httpd.shutdown()
        if self._thread:
            self._thread.join(timeout=5)
            self._thread = None
        if self._logger:
            self._logger.info("Video file server stopped")

    def build_url(self, relative_path: str, request: object | None = None) -> str:
        """Build a full HTTP URL for *relative_path*.

        Args:
            relative_path: Path relative to *root_dir* (e.g.
                ``eve_abc12345/output_deadbeef.mp4``).
            request: A ``gr.Request`` (or any object with a ``.headers`` dict)
                used to detect the client-facing hostname. Falls back to
                ``localhost`` if unavailable.

        Returns:
            URL like ``http://192.168.1.5:7861/eve_abc/output_dead.mp4``.
        """
        host = "localhost"
        if request is not None:
            headers = getattr(request, "headers", {}) or {}
            raw_host = headers.get("host", "")
            if raw_host:
                # Strip port from Host header (e.g. "192.168.1.5:7860" → "192.168.1.5")
                host = raw_host.rsplit(":", 1)[0]
        # Normalise backslashes for URL
        url_path = relative_path.replace("\\", "/")
        return f"http://{host}:{self.port}/{url_path}"