Upload 11 files

app.py

@@ -4,7 +4,7 @@ install_uvicorn_privacy_filters()
 
 from fastapi import FastAPI, File, UploadFile, Form, Depends, HTTPException, Request
 from rate_limit import shared_limit
-from response_privacy import requires_no_store
+from response_privacy import PrivateResponseCacheControl
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import Response, JSONResponse, FileResponse 
 from sqlalchemy.orm import Session
@@ -74,6 +74,7 @@ from router_messages import (router as messages_router, start_message_io, stop_m
 from router_wallet import (router as wallet_router, start_wallet_sql, stop_wallet_sql, close_wallet_sql)     # 💰 钱包/提现
 from proxy_access import start_proxy_access, stop_proxy_access, close_proxy_access
 from proxy_download_work import start_proxy_downloads, stop_proxy_downloads, close_proxy_downloads
+from media_proxy_service import start_public_media, stop_public_media, close_public_media, public_media_endpoint
 from router_proxy import router as proxy_router       # 🔗 代理下载
 from router_posts import (router as posts_router, start_posts_io, stop_posts_io, close_posts_io)       # 💬 讨论区
 from router_prompts import (router as prompts_router, start_prompts_io, stop_prompts_io, close_prompts_io)   # ✨ 提示词市场
@@ -170,16 +171,7 @@ app.add_middleware(UploadBodyGuard, capacity=_UPLOAD_CONCURRENCY, error_response
 from public_list_etag import PublicListETag
 app.add_middleware(PublicListETag)
 
-@app.middleware("http")
-async def private_response_cache_control(request: Request, call_next):
-    response = await call_next(request)
-    if requires_no_store(request.url.path, request.headers.get("Authorization")):
-        response.headers["Cache-Control"] = "no-store"
-        vary = [name.strip() for name in response.headers.get("Vary", "").split(",") if name.strip()]
-        if not any(name.lower() in {"authorization", "*"} for name in vary):
-            vary.append("Authorization")
-        response.headers["Vary"] = ", ".join(vary)
-    return response
+app.add_middleware(PrivateResponseCacheControl)
 
 
 @app.exception_handler(RequestValidationError)
@@ -327,6 +319,7 @@ async def on_startup():
             start_tasks_io()
             start_proxy_access()
             start_proxy_downloads()
+            start_public_media()
             start_comment_reads()
             start_admin_scripts()
 
@@ -387,6 +380,7 @@ async def _shutdown_resources():
     for stage, stop in (
         ('stop_wallet_sql', stop_wallet_sql),
         ('stop_proxy_downloads', stop_proxy_downloads),
+        ('stop_public_media', stop_public_media),
         ('stop_resource_reads', stop_resource_reads), ('stop_comment_writes', stop_comment_writes),
         ('stop_posts', stop_posts_io), ('stop_prompts', stop_prompts_io), ('stop_tasks', stop_tasks_io), ('stop_proxy_access', stop_proxy_access),
         ('stop_messages', stop_message_io), ('stop_resources', stop_resource_work), ('stop_open_sql', stop_open_sql), ('stop_config_sql', stop_config_sql), ('stop_shared', stop_shared_work),
@@ -404,6 +398,7 @@ async def _shutdown_resources():
     for stage, close, args in (
         ('wait_warmup', wait_startup_warmup, (app.state,)),
         ('close_proxy_downloads', close_proxy_downloads, ()),
+        ('close_public_media', close_public_media, ()),
         ('close_wallet_sql', close_wallet_sql, ()),
         ('close_resource_reads', close_resource_reads, ()), ('close_comment_writes', close_comment_writes, ()),
         ('close_posts', close_posts_io, ()), ('close_prompts', close_prompts_io, ()), ('close_tasks', close_tasks_io, ()), ('close_proxy_access', close_proxy_access, ()),
@@ -465,11 +460,13 @@ app.include_router(admin_plugins_router)  # 🔒 管理后台 - 第三方插件
 # 解决 Private 仓库下，本地客户端报 401 Unauthorized 的终极方案
 # ==========================================
 from media_proxy_service import read_public_media
+from media_file_response import OwnedMediaFileResponse
 
 @app.get("/api/image_proxy")
-def proxy_hf_image(url: str = None, path: str = None):
-    return read_public_media(url, path, store=db, download=hf_hub_download,
-                             file_response=FileResponse, error_response=create_error_response)
+@public_media_endpoint
+async def proxy_hf_image(url: str = None, path: str = None):
+    return await run_owned_thread(read_public_media, url, path, store=db, download=hf_hub_download,
+                             file_response=OwnedMediaFileResponse, error_response=create_error_response)
 
 # ==========================================
 # 上传接口 (将返回的 URL 替换为 Proxy 代理链接)

media_file_response.py

@@ -0,0 +1,115 @@
+"""媒体缓存文件的有界读取、Range交付与真实线程关闭。"""
+import asyncio
+import os
+from secrets import token_hex
+import stat
+
+from fastapi.responses import FileResponse
+from starlette.datastructures import MutableHeaders
+from owned_thread import run_owned_thread
+from workflow_stream_response import OwnedDelivery
+
+
+class OwnedMediaFileResponse(OwnedDelivery, FileResponse):
+    total_timeout = 300
+    operation_timeout = 15
+    eof_timeout = 5
+
+    def __init__(self, path, **kwargs):
+        snapshot = os.stat(path)
+        if not stat.S_ISREG(snapshot.st_mode):
+            raise ValueError('invalid media cache file')
+        super().__init__(path, stat_result=snapshot, **kwargs)
+        self._file = None
+        self._init_delivery(self._close_file)
+
+    async def send_owned(self, scope, receive, send):
+        async def bounded_send(message):
+            final = message['type'] == 'http.response.body' and not message.get('more_body', False)
+            async with asyncio.timeout(self.eof_timeout if final else self.operation_timeout):
+                await send(message)
+        async with asyncio.timeout(self.total_timeout):
+            return await super().send_owned(scope, receive, bounded_send)
+
+    async def _file_work(self, function, *args):
+        async with asyncio.timeout(self.operation_timeout):
+            return await run_owned_thread(function, *args)
+
+    async def _open_file(self):
+        def open_file():
+            self._file = open(self.path, 'rb')
+            self._check_snapshot(os.fstat(self._file.fileno()))
+        await self._file_work(open_file)
+
+    def _check_snapshot(self, current):
+        fields = ('st_dev', 'st_ino', 'st_size', 'st_mtime_ns')
+        if os.name != 'nt':
+            fields += ('st_ctime_ns',)
+        if (not stat.S_ISREG(current.st_mode) or not current.st_ino or not self.stat_result.st_ino
+                or any(getattr(current, key) != getattr(self.stat_result, key) for key in fields)):
+            raise RuntimeError('media cache file changed')
+
+    async def _prepare_file(self, header_only):
+        if header_only:
+            current = await self._file_work(os.stat, self.path)
+            self._check_snapshot(current)
+        else:
+            await self._open_file()
+
+    async def _close_file(self):
+        if self._file is None:
+            return
+        def close_file():
+            try:
+                self._file.close()
+            finally:
+                self._file = None
+        await run_owned_thread(close_file)
+
+    async def _send_part(self, send, start, end, *, final):
+        await self._file_work(self._file.seek, start)
+        remaining = end - start
+        if remaining == 0:
+            await send({'type': 'http.response.body', 'body': b'', 'more_body': not final})
+        while remaining:
+            chunk = await self._file_work(self._file.read, min(self.chunk_size, remaining))
+            if not chunk:
+                raise RuntimeError('media cache file truncated')
+            remaining -= len(chunk)
+            await send({'type': 'http.response.body', 'body': chunk, 'more_body': remaining > 0 or not final})
+
+    async def _handle_simple(self, send, send_header_only, send_pathsend):
+        await self._prepare_file(send_header_only)
+        await send({'type': 'http.response.start', 'status': self.status_code, 'headers': self.raw_headers})
+        if send_header_only:
+            await send({'type': 'http.response.body', 'body': b'', 'more_body': False})
+            return
+        await self._send_part(send, 0, self.stat_result.st_size, final=True)
+
+    async def _handle_single_range(self, send, start, end, file_size, send_header_only):
+        await self._prepare_file(send_header_only)
+        headers = MutableHeaders(raw=list(self.raw_headers))
+        headers['content-range'] = f'bytes {start}-{end - 1}/{file_size}'
+        headers['content-length'] = str(end - start)
+        await send({'type': 'http.response.start', 'status': 206, 'headers': headers.raw})
+        if send_header_only:
+            await send({'type': 'http.response.body', 'body': b'', 'more_body': False})
+            return
+        await self._send_part(send, start, end, final=True)
+
+    async def _handle_multiple_ranges(self, send, ranges, file_size, send_header_only):
+        await self._prepare_file(send_header_only)
+        boundary = token_hex(13)
+        length, header = self.generate_multipart(ranges, boundary, file_size, self.headers['content-type'])
+        headers = MutableHeaders(raw=list(self.raw_headers))
+        headers['content-type'] = f'multipart/byteranges; boundary={boundary}'
+        headers['content-length'] = str(length)
+        await send({'type': 'http.response.start', 'status': 206, 'headers': headers.raw})
+        if send_header_only:
+            await send({'type': 'http.response.body', 'body': b'', 'more_body': False})
+            return
+        for start, end in ranges:
+            await send({'type': 'http.response.body', 'body': header(start, end), 'more_body': True})
+            await self._send_part(send, start, end, final=False)
+            await send({'type': 'http.response.body', 'body': b'\r\n', 'more_body': True})
+        await send({'type': 'http.response.body', 'body': f'--{boundary}--'.encode('latin-1'), 'more_body': False})

media_proxy_service.py

@@ -7,6 +7,13 @@ import urllib.parse
 import uuid
 from fastapi import HTTPException
 from fastapi.responses import JSONResponse
+from proxy_download_work import ProxyDownloadWork
+
+_public_media_work = ProxyDownloadWork(16)
+public_media_endpoint = _public_media_work.endpoint
+start_public_media = _public_media_work.start
+stop_public_media = _public_media_work.stop
+close_public_media = _public_media_work.close
 
 logger = logging.getLogger("ComfyUI-Ranking")
 

proxy_download_work.py

@@ -3,7 +3,7 @@ import asyncio
 from functools import wraps
 
 from fastapi import HTTPException
-from workflow_stream_response import OwnedProxyResponse
+from workflow_stream_response import OwnedDelivery
 
 
 def _spawn(function):
@@ -58,7 +58,7 @@ class _Download:
 
     async def prepare(self, function, args, kwargs):
         result = await function(*args, **kwargs)
-        if isinstance(result, OwnedProxyResponse):
+        if isinstance(result, OwnedDelivery):
             result.bind_download(self)
             self.response = result
         return result

response_privacy.py

@@ -1,5 +1,6 @@
 """浏览器/共享 HTTP 缓存的私有响应边界，无数据库或应用依赖。"""
 import re
+from starlette.datastructures import Headers, MutableHeaders
 
 _PRIVATE_PATH = re.compile(
     r"^/api/(wallet|messages|chats|admin|disputes|users)(?:/|$)"
@@ -12,3 +13,26 @@ def requires_no_store(path, authorization=None):
     # 含 Authorization 的响应即使路径公开也不交给共享 HTTP 缓存。
     # 匿名访问私有路径的错误响应及登录返回的 Token 同样不能缓存。
     return bool(authorization or _PRIVATE_PATH.search(path))
+
+
+class PrivateResponseCacheControl:
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        if (scope['type'] != 'http'
+                or not requires_no_store(scope.get('path', ''), Headers(scope=scope).get('authorization'))):
+            return await self.app(scope, receive, send)
+
+        async def private_send(message):
+            if message['type'] == 'http.response.start':
+                headers = MutableHeaders(raw=list(message.get('headers', ())))
+                headers['Cache-Control'] = 'no-store'
+                vary = [name.strip() for name in headers.get('Vary', '').split(',') if name.strip()]
+                if not any(name.lower() in {'authorization', '*'} for name in vary):
+                    vary.append('Authorization')
+                headers['Vary'] = ', '.join(vary)
+                message = {**message, 'headers': headers.raw}
+            await send(message)
+
+        return await self.app(scope, receive, private_send)

router_proxy.py

@@ -48,6 +48,8 @@ async def _validate_public_http_target(url: str) -> str:
     try:
         await _resolve_public_ips(host, port or (443 if parsed.scheme == "https" else 80))
     except PublicResolutionError as error:
+        if error.reason == 'dns_busy':
+            raise HTTPException(status_code=503, detail='service_unavailable') from None
         detail = "下载地址无法解析" if error.reason == "dns_unavailable" else "下载地址解析到非公网地址"
         raise HTTPException(status_code=400, detail=detail) from None
     return host

safe_http_transport.py

@@ -8,6 +8,7 @@ import socket
 import httpcore
 import httpx
 from proxy_resources import ProxyResources
+from owned_thread import OwnedThreadGroup
 
 
 _validated_ip = ContextVar("ranking_validated_ip", default=None)
@@ -18,7 +19,12 @@ _DNS_TIMEOUT = 8
 class PublicResolutionError(httpx.ConnectError):
     def __init__(self, reason):
         self.reason = reason
-        super().__init__("目标域名无法解析" if reason == "dns_unavailable" else "目标域名解析到非公网地址")
+        super().__init__("目标域名无法解析" if reason in ('dns_unavailable', 'dns_busy') else "目标域名解析到非公网地址")
+
+
+_DNS_SYSTEM_CAPACITY = 8
+_DNS_MAX_RECORDS = 32
+_system_dns = OwnedThreadGroup(_DNS_SYSTEM_CAPACITY, lambda: PublicResolutionError('dns_busy'))
 
 
 def _public(address):
@@ -113,23 +119,42 @@ async def _resolve_hf_public_ips(host):
 
 
 async def _resolve_public_ips(host: str, port: int) -> tuple[str, ...]:
-    normalized_host = host.rstrip(".").lower()
+    try:
+        normalized_host = host.rstrip('.').lower()
+        if (not normalized_host or any(ord(char) <= 32 or ord(char) == 127 for char in normalized_host)
+                or len(normalized_host.encode('idna')) > 253):
+            raise ValueError()
+    except (AttributeError, TypeError, ValueError):
+        raise PublicResolutionError('dns_unavailable') from None
+    try:
+        async with asyncio.timeout(_DNS_TIMEOUT):
+            return await _resolve_public_ips_once(normalized_host, port)
+    except TimeoutError:
+        raise PublicResolutionError('dns_unavailable') from None
+
+
+async def _resolve_public_ips_once(normalized_host: str, port: int) -> tuple[str, ...]:
     try:
         address = ipaddress.ip_address(normalized_host)
         addresses = [address]
     except ValueError:
         try:
-            records = await asyncio.to_thread(
+            records = await _system_dns.run(
                 socket.getaddrinfo, normalized_host, port, type=socket.SOCK_STREAM
             )
         except OSError as exc:
             if normalized_host == "huggingface.co":
                 return await _resolve_hf_public_ips(normalized_host)
             raise PublicResolutionError("dns_unavailable") from exc
-        addresses = [
-            ipaddress.ip_address(record[4][0].split("%", 1)[0])
-            for record in records
-        ]
+        if not isinstance(records, list) or len(records) > _DNS_MAX_RECORDS:
+            raise PublicResolutionError('dns_unavailable')
+        try:
+            addresses = [
+                ipaddress.ip_address(record[4][0].split("%", 1)[0])
+                for record in records
+            ]
+        except (IndexError, AttributeError, TypeError, ValueError):
+            raise PublicResolutionError('dns_unavailable') from None
 
     if not addresses or any(not _public(address) for address in addresses):
         if normalized_host == "huggingface.co":

workflow_download_service.py

@@ -119,6 +119,7 @@ async def prepare_workflow_download(item, item_id, *, store, source_authorizer,
             "下载源重定向无效或次数过多": "invalid_redirect",
             "下载源重定向失败": "redirect_failed",
             "下载文件超过大小限制": "size_exceeded",
+            "service_unavailable": "dns_busy",
         }
         reason = reasons.get(error.detail, "http_rejected") if type(error.detail) is str else "http_rejected"
         status = error.status_code if type(error.status_code) is int else 0

workflow_stream_response.py

@@ -3,9 +3,8 @@ import asyncio
 from fastapi.responses import StreamingResponse
 
 
-class OwnedProxyResponse(StreamingResponse):
-    def __init__(self, content, *, close_callback, **kwargs):
-        super().__init__(content, **kwargs)
+class OwnedDelivery:
+    def _init_delivery(self, close_callback):
         self._close_callback = close_callback
         self._download = None
         self._close_task = None
@@ -19,7 +18,7 @@ class OwnedProxyResponse(StreamingResponse):
         if self._close_task is None:
             async def close():
                 try:
-                    close_iterator = getattr(self.body_iterator, "aclose", None)
+                    close_iterator = getattr(getattr(self, 'body_iterator', None), "aclose", None)
                     if close_iterator is not None:
                         await close_iterator()
                 finally:
@@ -36,6 +35,8 @@ class OwnedProxyResponse(StreamingResponse):
             raise cancelled
 
     async def __call__(self, scope, receive, send):
+        if self._close_task is not None:
+            raise asyncio.CancelledError()
         if self._download is not None:
             return await self._download.deliver(self, scope, receive, send)
         return await self.send_owned(scope, receive, send)
@@ -46,5 +47,11 @@ class OwnedProxyResponse(StreamingResponse):
         finally:
             await self.aclose()
 
+
+class OwnedProxyResponse(OwnedDelivery, StreamingResponse):
+    def __init__(self, content, *, close_callback, **kwargs):
+        super().__init__(content, **kwargs)
+        self._init_delivery(close_callback)
+
 # 保留既有工作流调用方的导出名称。
 OwnedWorkflowResponse = OwnedProxyResponse