Spaces:
Configuration error
Configuration error
| import asyncio | |
| import logging | |
| import typing | |
| from dataclasses import dataclass | |
| from functools import partial | |
| from urllib import parse | |
| from urllib.parse import urlencode | |
| import aiohttp | |
| from aiohttp import ClientSession, ClientTimeout, ClientResponse | |
| import anyio | |
| 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 tqdm.asyncio import tqdm as tqdm_asyncio | |
| from mediaflow_proxy.configs import settings | |
| from mediaflow_proxy.const import SUPPORTED_REQUEST_HEADERS | |
| from mediaflow_proxy.utils.crypto_utils import EncryptionHandler | |
| from mediaflow_proxy.utils.stream_transformers import StreamTransformer | |
| from mediaflow_proxy.utils.http_client import ( | |
| create_aiohttp_session, | |
| get_routing_config, | |
| _ensure_routing_initialized, | |
| _create_connector, | |
| ) | |
| logger = logging.getLogger(__name__) | |
| class DownloadError(Exception): | |
| def __init__(self, status_code, message): | |
| self.status_code = status_code | |
| self.message = message | |
| super().__init__(message) | |
| async def fetch_with_retry( | |
| session: ClientSession, | |
| method: str, | |
| url: str, | |
| headers: dict, | |
| proxy: typing.Optional[str] = None, | |
| **kwargs, | |
| ) -> ClientResponse: | |
| """ | |
| Fetches a URL with retry logic using native aiohttp. | |
| Args: | |
| session: The aiohttp ClientSession to use for the request. | |
| method: The HTTP method to use (e.g., GET, POST). | |
| url: The URL to fetch. | |
| headers: The headers to include in the request. | |
| proxy: Optional proxy URL for HTTP proxies. | |
| **kwargs: Additional arguments to pass to the request. | |
| Returns: | |
| ClientResponse: The HTTP response. | |
| Raises: | |
| DownloadError: If the request fails after retries. | |
| """ | |
| try: | |
| response = await session.request(method, url, headers=headers, proxy=proxy, **kwargs) | |
| response.raise_for_status() | |
| return response | |
| except asyncio.TimeoutError: | |
| logger.warning(f"Timeout while downloading {url}") | |
| raise DownloadError(409, f"Timeout while downloading {url}") | |
| except aiohttp.ClientResponseError as e: | |
| if e.status == 404: | |
| logger.debug(f"Segment not found (404): {url}") | |
| raise | |
| logger.error(f"HTTP error {e.status} while downloading {url}") | |
| raise DownloadError(e.status, f"HTTP error {e.status} while downloading {url}") | |
| except aiohttp.ClientError as e: | |
| logger.error(f"Client error downloading {url}: {e}") | |
| raise DownloadError(502, f"Client error downloading {url}: {e}") | |
| except Exception as e: | |
| logger.error(f"Error downloading {url}: {e}") | |
| raise | |
| class Streamer: | |
| """Handles streaming HTTP responses using aiohttp.""" | |
| def __init__(self, session: ClientSession, proxy_url: typing.Optional[str] = None): | |
| """ | |
| Initializes the Streamer with an aiohttp session. | |
| Args: | |
| session: The aiohttp ClientSession to use for streaming. | |
| proxy_url: Optional proxy URL for HTTP proxies. | |
| """ | |
| self.session = session | |
| self.proxy_url = proxy_url | |
| self.response: typing.Optional[ClientResponse] = None | |
| self.progress_bar = None | |
| self.bytes_transferred = 0 | |
| self.start_byte = 0 | |
| self.end_byte = 0 | |
| self.total_size = 0 | |
| async def create_streaming_response(self, url: str, headers: dict): | |
| """ | |
| Creates and sends a streaming request. | |
| Args: | |
| url: The URL to stream from. | |
| headers: The headers to include in the request. | |
| """ | |
| try: | |
| self.response = await self.session.get(url, headers=headers, proxy=self.proxy_url) | |
| self.response.raise_for_status() | |
| except asyncio.TimeoutError: | |
| logger.warning("Timeout while creating streaming response") | |
| raise DownloadError(409, "Timeout while creating streaming response") | |
| except aiohttp.ClientResponseError as e: | |
| if e.status == 404: | |
| logger.debug(f"Segment not found (404): {url}") | |
| raise | |
| logger.error(f"HTTP error {e.status} while creating streaming response") | |
| raise DownloadError(e.status, f"HTTP error {e.status} while creating streaming response") | |
| except aiohttp.ClientError as e: | |
| logger.error(f"Error creating streaming response: {e}") | |
| raise DownloadError(502, f"Error creating streaming response: {e}") | |
| except Exception as e: | |
| logger.error(f"Error creating streaming response: {e}") | |
| raise RuntimeError(f"Error creating streaming response: {e}") | |
| async def stream_content( | |
| self, transformer: typing.Optional[StreamTransformer] = None | |
| ) -> typing.AsyncGenerator[bytes, None]: | |
| """ | |
| Stream content from the response, optionally applying a transformer. | |
| Args: | |
| transformer: Optional StreamTransformer to apply host-specific | |
| content manipulation (e.g., PNG stripping, TS detection). | |
| If None, content is streamed directly without modification. | |
| Yields: | |
| Bytes chunks from the upstream response. | |
| """ | |
| if not self.response: | |
| raise RuntimeError("No response available for streaming") | |
| try: | |
| self.parse_content_range() | |
| # Create async generator from response content | |
| async def raw_chunks(): | |
| async for chunk in self.response.content.iter_any(): | |
| yield chunk | |
| # Choose the chunk source based on whether we have a transformer | |
| if transformer: | |
| chunk_source = transformer.transform(raw_chunks()) | |
| else: | |
| chunk_source = raw_chunks() | |
| if settings.enable_streaming_progress: | |
| with tqdm_asyncio( | |
| total=self.total_size, | |
| initial=self.start_byte, | |
| unit="B", | |
| unit_scale=True, | |
| unit_divisor=1024, | |
| desc="Streaming", | |
| ncols=100, | |
| mininterval=1, | |
| ) as self.progress_bar: | |
| async for chunk in chunk_source: | |
| yield chunk | |
| self.bytes_transferred += len(chunk) | |
| self.progress_bar.update(len(chunk)) | |
| else: | |
| async for chunk in chunk_source: | |
| yield chunk | |
| self.bytes_transferred += len(chunk) | |
| except asyncio.TimeoutError: | |
| logger.warning("Timeout while streaming") | |
| raise DownloadError(409, "Timeout while streaming") | |
| except aiohttp.ServerDisconnectedError as e: | |
| # Special handling for connection closed errors | |
| logger.warning(f"Remote server closed connection prematurely: {e}") | |
| if self.bytes_transferred > 0: | |
| logger.info( | |
| f"Partial content received ({self.bytes_transferred} bytes). Continuing with available data." | |
| ) | |
| return | |
| else: | |
| raise DownloadError(502, f"Remote server closed connection without sending any data: {e}") | |
| except GeneratorExit: | |
| logger.info("Streaming session stopped by the user") | |
| except aiohttp.ClientPayloadError as e: | |
| # Handle network read errors gracefully | |
| logger.warning(f"ClientPayloadError while streaming: {e}") | |
| if self.bytes_transferred > 0: | |
| logger.info( | |
| f"Partial content received ({self.bytes_transferred} bytes) before error. Graceful termination." | |
| ) | |
| return | |
| else: | |
| raise DownloadError(502, f"ClientPayloadError while streaming: {e}") | |
| except aiohttp.ClientError as e: | |
| logger.warning(f"ClientError while streaming: {e}") | |
| if self.bytes_transferred > 0: | |
| logger.info( | |
| f"Partial content received ({self.bytes_transferred} bytes) before error. Graceful termination." | |
| ) | |
| return | |
| else: | |
| raise DownloadError(502, f"ClientError while streaming: {e}") | |
| def format_bytes(size) -> str: | |
| power = 2**10 | |
| n = 0 | |
| units = {0: "B", 1: "KB", 2: "MB", 3: "GB", 4: "TB"} | |
| while size > power: | |
| size /= power | |
| n += 1 | |
| return f"{size:.2f} {units[n]}" | |
| def parse_content_range(self): | |
| content_range = self.response.headers.get("Content-Range", "") | |
| if content_range: | |
| range_info = content_range.split()[-1] | |
| self.start_byte, self.end_byte, self.total_size = map(int, range_info.replace("/", "-").split("-")) | |
| else: | |
| self.start_byte = 0 | |
| self.total_size = int(self.response.headers.get("Content-Length", 0)) | |
| self.end_byte = self.total_size - 1 if self.total_size > 0 else 0 | |
| async def get_text(self, url: str, headers: dict) -> str: | |
| """ | |
| Sends a GET request to a URL and returns the response text. | |
| Args: | |
| url: The URL to send the GET request to. | |
| headers: The headers to include in the request. | |
| Returns: | |
| str: The response text. | |
| """ | |
| try: | |
| self.response = await fetch_with_retry(self.session, "GET", url, headers, proxy=self.proxy_url) | |
| return await self.response.text() | |
| except tenacity.RetryError as e: | |
| raise e.last_attempt.result() | |
| async def close(self): | |
| """ | |
| Closes the HTTP response and session. | |
| """ | |
| if self.response: | |
| self.response.close() | |
| if self.progress_bar: | |
| self.progress_bar.close() | |
| await self.session.close() | |
| async def download_file_with_retry(url: str, headers: dict) -> bytes: | |
| """ | |
| Downloads a file with retry logic. | |
| Args: | |
| url: The URL of the file to download. | |
| headers: The headers to include in the request. | |
| Returns: | |
| bytes: The downloaded file content. | |
| Raises: | |
| DownloadError: If the download fails after retries. | |
| """ | |
| async with create_aiohttp_session(url) as (session, proxy_url): | |
| try: | |
| response = await fetch_with_retry(session, "GET", url, headers, proxy=proxy_url) | |
| return await response.read() | |
| except DownloadError as e: | |
| logger.error(f"Failed to download file: {e}") | |
| raise e | |
| except tenacity.RetryError as e: | |
| raise DownloadError(502, f"Failed to download file: {e.last_attempt.result()}") | |
| async def request_with_retry(method: str, url: str, headers: dict, **kwargs) -> ClientResponse: | |
| """ | |
| Sends an HTTP request with retry logic. | |
| Args: | |
| method: The HTTP method to use (e.g., GET, POST). | |
| url: The URL to send the request to. | |
| headers: The headers to include in the request. | |
| **kwargs: Additional arguments to pass to the request. | |
| Returns: | |
| ClientResponse: The HTTP response. | |
| Raises: | |
| DownloadError: If the request fails after retries. | |
| """ | |
| async with create_aiohttp_session(url) as (session, proxy_url): | |
| try: | |
| response = await fetch_with_retry(session, method, url, headers, proxy=proxy_url, **kwargs) | |
| # Read the content so it's available after session closes | |
| await response.read() | |
| return response | |
| except DownloadError as e: | |
| logger.error(f"Failed to make request: {e}") | |
| raise | |
| async def create_streamer(url: str = None) -> Streamer: | |
| """ | |
| Create a Streamer configured for the given URL. | |
| The Streamer manages its own session lifecycle. Call streamer.close() | |
| when done to release resources. | |
| Args: | |
| url: Optional URL for routing configuration (SSL/proxy settings). | |
| Returns: | |
| Streamer: A configured Streamer instance. | |
| """ | |
| _ensure_routing_initialized() | |
| routing_config = get_routing_config() | |
| route_match = routing_config.match_url(url) | |
| # Use sock_read timeout: no total timeout, but timeout if no data received | |
| # for sock_read seconds. This correctly handles: | |
| # - Live streams (indefinite duration) | |
| # - Large file downloads (total time depends on file size) | |
| # - Seek operations (upstream may take time to seek) | |
| # - Dead connection detection (timeout if no data flows) | |
| timeout_config = ClientTimeout( | |
| total=None, | |
| sock_read=settings.transport_config.timeout, | |
| ) | |
| connector, proxy_url = _create_connector(route_match.proxy_url, route_match.verify_ssl) | |
| session = ClientSession(connector=connector, timeout=timeout_config) | |
| return Streamer(session, proxy_url) | |
| # Keep setup_streamer as alias for backward compatibility during transition | |
| async def setup_streamer(url: str = None) -> typing.Tuple[ClientSession, str, Streamer]: | |
| """ | |
| Set up an aiohttp session and streamer. | |
| DEPRECATED: Use create_streamer() instead which returns only the Streamer. | |
| Args: | |
| url: Optional URL for routing configuration. | |
| Returns: | |
| Tuple of (session, proxy_url, streamer) | |
| """ | |
| streamer = await create_streamer(url) | |
| return streamer.session, streamer.proxy_url, streamer | |
| def encode_mediaflow_proxy_url( | |
| mediaflow_proxy_url: str, | |
| endpoint: typing.Optional[str] = None, | |
| destination_url: typing.Optional[str] = None, | |
| query_params: typing.Optional[dict] = None, | |
| request_headers: typing.Optional[dict] = None, | |
| response_headers: typing.Optional[dict] = None, | |
| propagate_response_headers: typing.Optional[dict] = None, | |
| remove_response_headers: typing.Optional[list[str]] = None, | |
| encryption_handler: EncryptionHandler = None, | |
| expiration: int = None, | |
| ip: str = None, | |
| filename: typing.Optional[str] = None, | |
| stream_transformer: typing.Optional[str] = None, | |
| ) -> str: | |
| """ | |
| Encodes & Encrypt (Optional) a MediaFlow proxy URL with query parameters and headers. | |
| Args: | |
| mediaflow_proxy_url: The base MediaFlow proxy URL. | |
| endpoint: The endpoint to append to the base URL. Defaults to None. | |
| destination_url: The destination URL to include in the query parameters. Defaults to None. | |
| query_params: Additional query parameters to include. Defaults to None. | |
| request_headers: Headers to include as query parameters. Defaults to None. | |
| response_headers: Headers to include as query parameters (r_ prefix). Defaults to None. | |
| propagate_response_headers: Response headers that propagate to segments (rp_ prefix). Defaults to None. | |
| remove_response_headers: List of response header names to remove. Defaults to None. | |
| encryption_handler: The encryption handler to use. Defaults to None. | |
| expiration: The expiration time for the encrypted token. Defaults to None. | |
| ip: The public IP address to include in the query parameters. Defaults to None. | |
| filename: Filename to be preserved for media players like Infuse. Defaults to None. | |
| stream_transformer: ID of the stream transformer to apply. Defaults to None. | |
| Returns: | |
| str: The encoded MediaFlow proxy URL. | |
| """ | |
| # Prepare query parameters | |
| query_params = query_params or {} | |
| if destination_url is not None: | |
| query_params["d"] = destination_url | |
| # Add headers if provided (always use lowercase prefix for consistency) | |
| # Filter out empty values to avoid URLs like &h_if-range=&h_referer=... | |
| if request_headers: | |
| query_params.update( | |
| { | |
| key if key.lower().startswith("h_") else f"h_{key}": value | |
| for key, value in request_headers.items() | |
| if value # Skip empty/None values | |
| } | |
| ) | |
| if response_headers: | |
| query_params.update( | |
| { | |
| key if key.lower().startswith("r_") else f"r_{key}": value | |
| for key, value in response_headers.items() | |
| if value # Skip empty/None values | |
| } | |
| ) | |
| # Add propagate response headers (rp_ prefix - these propagate to segments) | |
| if propagate_response_headers: | |
| query_params.update( | |
| { | |
| key if key.lower().startswith("rp_") else f"rp_{key}": value | |
| for key, value in propagate_response_headers.items() | |
| if value # Skip empty/None values | |
| } | |
| ) | |
| # Add remove headers if provided (x_ prefix for "exclude") | |
| if remove_response_headers: | |
| query_params["x_headers"] = ",".join(remove_response_headers) | |
| # Add stream transformer if provided | |
| if stream_transformer: | |
| query_params["transformer"] = stream_transformer | |
| # Construct the base URL | |
| if endpoint is None: | |
| base_url = mediaflow_proxy_url | |
| else: | |
| base_url = parse.urljoin(mediaflow_proxy_url, endpoint) | |
| # Ensure base_url doesn't end with a slash for consistent handling | |
| if base_url.endswith("/"): | |
| base_url = base_url[:-1] | |
| # Handle encryption if needed | |
| if encryption_handler: | |
| encrypted_token = encryption_handler.encrypt_data(query_params, expiration, ip) | |
| # Parse the base URL to get its components | |
| parsed_url = parse.urlparse(base_url) | |
| # Insert the token at the beginning of the path | |
| new_path = f"/_token_{encrypted_token}{parsed_url.path}" | |
| # Reconstruct the URL with the token at the beginning of the path | |
| url_parts = list(parsed_url) | |
| url_parts[2] = new_path # Update the path component | |
| # Build the URL | |
| url = parse.urlunparse(url_parts) | |
| # Add filename at the end if provided | |
| if filename: | |
| url = f"{url}/{parse.quote(filename)}" | |
| return url | |
| else: | |
| # No encryption, use regular query parameters | |
| url = base_url | |
| if filename: | |
| url = f"{url}/{parse.quote(filename)}" | |
| if query_params: | |
| return f"{url}?{urlencode(query_params)}" | |
| return url | |
| def encode_stremio_proxy_url( | |
| stremio_proxy_url: str, | |
| destination_url: str, | |
| request_headers: typing.Optional[dict] = None, | |
| response_headers: typing.Optional[dict] = None, | |
| ) -> str: | |
| """ | |
| Encodes a Stremio proxy URL with destination URL and headers. | |
| Format: http://127.0.0.1:11470/proxy/d=<encoded_origin>&h=<headers>&r=<response_headers>/<path><query> | |
| Args: | |
| stremio_proxy_url: The base Stremio proxy URL. | |
| destination_url: The destination URL to proxy. | |
| request_headers: Headers to include as query parameters. Defaults to None. | |
| response_headers: Response headers to include as query parameters. Defaults to None. | |
| Returns: | |
| str: The encoded Stremio proxy URL. | |
| """ | |
| # Parse the destination URL to separate origin, path, and query | |
| parsed_dest = parse.urlparse(destination_url) | |
| dest_origin = f"{parsed_dest.scheme}://{parsed_dest.netloc}" | |
| dest_path = parsed_dest.path.lstrip("/") | |
| dest_query = parsed_dest.query | |
| # Prepare query parameters list for proper handling of multiple headers | |
| query_parts = [] | |
| # Add destination origin (scheme + netloc only) with proper encoding | |
| query_parts.append(f"d={parse.quote_plus(dest_origin)}") | |
| # Add request headers | |
| if request_headers: | |
| for key, value in request_headers.items(): | |
| header_string = f"{key}:{value}" | |
| query_parts.append(f"h={parse.quote_plus(header_string)}") | |
| # Add response headers | |
| if response_headers: | |
| for key, value in response_headers.items(): | |
| header_string = f"{key}:{value}" | |
| query_parts.append(f"r={parse.quote_plus(header_string)}") | |
| # Ensure base_url doesn't end with a slash for consistent handling | |
| base_url = stremio_proxy_url.rstrip("/") | |
| # Construct the URL path with query string | |
| query_string = "&".join(query_parts) | |
| # Build the final URL: /proxy/{opts}/{pathname}{search} | |
| url_path = f"/proxy/{query_string}" | |
| # Append the path from destination URL | |
| if dest_path: | |
| url_path = f"{url_path}/{dest_path}" | |
| # Append the query string from destination URL | |
| if dest_query: | |
| url_path = f"{url_path}?{dest_query}" | |
| return f"{base_url}{url_path}" | |
| def get_original_scheme(request: Request) -> str: | |
| """ | |
| Determines the original scheme (http or https) of the request. | |
| Args: | |
| request: The incoming HTTP request. | |
| Returns: | |
| str: The original scheme ('http' or 'https') | |
| """ | |
| # Check the X-Forwarded-Proto header first | |
| forwarded_proto = request.headers.get("X-Forwarded-Proto") | |
| if forwarded_proto: | |
| return forwarded_proto | |
| # Check if the request is secure | |
| if request.url.scheme == "https" or request.headers.get("X-Forwarded-Ssl") == "on": | |
| return "https" | |
| # Check for other common headers that might indicate HTTPS | |
| if ( | |
| request.headers.get("X-Forwarded-Ssl") == "on" | |
| or request.headers.get("X-Forwarded-Protocol") == "https" | |
| or request.headers.get("X-Url-Scheme") == "https" | |
| ): | |
| return "https" | |
| # Default to http if no indicators of https are found | |
| return "http" | |
| class ProxyRequestHeaders: | |
| request: dict | |
| response: dict | |
| remove: list # headers to remove from response | |
| propagate: dict # response headers to propagate to segments (rp_ prefix) | |
| def apply_header_manipulation( | |
| base_headers: dict, proxy_headers: ProxyRequestHeaders, include_propagate: bool = True | |
| ) -> dict: | |
| """ | |
| Apply response header additions and removals. | |
| This function filters out headers specified in proxy_headers.remove, | |
| then merges in headers from proxy_headers.response and optionally proxy_headers.propagate. | |
| Args: | |
| base_headers: The base headers to start with. | |
| proxy_headers: The proxy headers containing response additions and removals. | |
| include_propagate: Whether to include propagate headers (rp_). | |
| Set to False for manifests, True for segments. Defaults to True. | |
| Returns: | |
| dict: The manipulated headers. | |
| """ | |
| remove_set = set(h.lower() for h in proxy_headers.remove) | |
| result = {k: v for k, v in base_headers.items() if k.lower() not in remove_set} | |
| # Apply propagate headers first (for segments), then response headers (response takes precedence) | |
| if include_propagate: | |
| result.update(proxy_headers.propagate) | |
| result.update(proxy_headers.response) | |
| return result | |
| def get_proxy_headers(request: Request) -> ProxyRequestHeaders: | |
| """ | |
| Extracts proxy headers from the request query parameters. | |
| Args: | |
| request: The incoming HTTP request. | |
| Returns: | |
| ProxyRequest: A named tuple containing the request headers, response headers, and headers to remove. | |
| """ | |
| request_headers = {k: v for k, v in request.headers.items() if k in SUPPORTED_REQUEST_HEADERS and v} | |
| # Extract h_ prefixed headers from query params, filtering out empty values | |
| for k, v in request.query_params.items(): | |
| if k.lower().startswith("h_") and v: # Skip empty values | |
| request_headers[k[2:].lower()] = v | |
| request_headers.setdefault("user-agent", settings.user_agent) | |
| # Handle common misspelling of referer | |
| if "referrer" in request_headers: | |
| if "referer" not in request_headers: | |
| request_headers["referer"] = request_headers.pop("referrer") | |
| # r_ prefix: response headers (manifest only, not propagated to segments) | |
| # Filter out empty values | |
| response_headers = { | |
| k[2:].lower(): v | |
| for k, v in request.query_params.items() | |
| if k.lower().startswith("r_") and not k.lower().startswith("rp_") and v | |
| } | |
| # rp_ prefix: response headers that propagate to segments | |
| # Filter out empty values | |
| propagate_headers = {k[3:].lower(): v for k, v in request.query_params.items() if k.lower().startswith("rp_") and v} | |
| # Parse headers to remove from response (x_headers parameter) | |
| x_headers_param = request.query_params.get("x_headers", "") | |
| remove_headers = [h.strip().lower() for h in x_headers_param.split(",") if h.strip()] if x_headers_param else [] | |
| return ProxyRequestHeaders(request_headers, response_headers, remove_headers, propagate_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) | |
| self.actual_content_length = 0 | |
| 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: | |
| # Track if response headers have been sent to prevent duplicate headers | |
| response_started = False | |
| # Track if response finalization (more_body: False) has been sent to prevent ASGI protocol violation | |
| finalization_sent = False | |
| try: | |
| # Initialize headers | |
| headers = list(self.raw_headers) | |
| # Start the response | |
| await send( | |
| { | |
| "type": "http.response.start", | |
| "status": self.status_code, | |
| "headers": headers, | |
| } | |
| ) | |
| response_started = True | |
| # Track if we've sent any data | |
| data_sent = False | |
| try: | |
| 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}) | |
| data_sent = True | |
| self.actual_content_length += len(chunk) | |
| except (ConnectionResetError, anyio.BrokenResourceError): | |
| logger.info("Client disconnected during streaming") | |
| return | |
| # Successfully streamed all content | |
| await send({"type": "http.response.body", "body": b"", "more_body": False}) | |
| finalization_sent = True | |
| except (aiohttp.ServerDisconnectedError, aiohttp.ClientPayloadError, aiohttp.ClientError) as e: | |
| # Handle connection closed / read errors gracefully | |
| if data_sent: | |
| # We've sent some data to the client. With Content-Length set, we cannot | |
| # gracefully finalize a partial response - h11 will raise LocalProtocolError | |
| # if we try to send more_body: False without delivering all promised bytes. | |
| # The best we can do is log and return silently, letting the client handle | |
| # the incomplete response (most players will just stop or retry). | |
| logger.warning( | |
| f"Upstream connection error after partial streaming ({self.actual_content_length} bytes transferred): {e}" | |
| ) | |
| # Don't try to finalize - just return and let the connection close naturally | |
| return | |
| else: | |
| # No data was sent, re-raise the error | |
| logger.error(f"Upstream error before any data was streamed: {e}") | |
| raise | |
| except Exception as e: | |
| logger.exception(f"Error in stream_response: {str(e)}") | |
| if not isinstance(e, (ConnectionResetError, anyio.BrokenResourceError)) and not response_started: | |
| # Only attempt to send error response if headers haven't been sent yet | |
| try: | |
| await send( | |
| { | |
| "type": "http.response.start", | |
| "status": 502, | |
| "headers": [(b"content-type", b"text/plain")], | |
| } | |
| ) | |
| error_message = f"Streaming error: {str(e)}".encode("utf-8") | |
| await send({"type": "http.response.body", "body": error_message, "more_body": False}) | |
| finalization_sent = True | |
| except Exception: | |
| # If we can't send an error response, just log it | |
| pass | |
| elif response_started and not finalization_sent and not data_sent: | |
| # Response started but no data sent yet - we can safely finalize | |
| # (If data was sent with Content-Length, we can't finalize without h11 error) | |
| try: | |
| await send({"type": "http.response.body", "body": b"", "more_body": False}) | |
| finalization_sent = True | |
| except Exception: | |
| pass | |
| # If data was sent but streaming failed, just return silently | |
| # The client will see an incomplete response which is unavoidable with Content-Length | |
| async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: | |
| async with anyio.create_task_group() as task_group: | |
| stream_func = partial(self.stream_response, send) | |
| listen_func = partial(self.listen_for_disconnect, receive) | |
| async def wrap(func: typing.Callable[[], typing.Awaitable[None]]) -> None: | |
| try: | |
| await func() | |
| except Exception as e: | |
| # Note: stream_response and listen_for_disconnect handle their own exceptions | |
| # internally. This is a safety net for any unexpected exceptions that might | |
| # escape due to future code changes. | |
| if not isinstance(e, anyio.get_cancelled_exc_class()): | |
| logger.exception(f"Unexpected error in streaming task: {type(e).__name__}: {e}") | |
| # Re-raise unexpected errors to surface bugs rather than silently swallowing them | |
| raise | |
| finally: | |
| # Cancel task group when either task completes or fails: | |
| # - stream_func finished (success or failure) -> stop listening for disconnect | |
| # - listen_func finished (client disconnected) -> stop streaming | |
| task_group.cancel_scope.cancel() | |
| # Start the streaming response in a separate task | |
| task_group.start_soon(wrap, stream_func) | |
| # Listen for disconnect events | |
| await wrap(listen_func) | |
| if self.background is not None: | |
| await self.background() | |