"""请求转发 + base64→URL 改写。 - 用 httpx.AsyncClient 把客户端请求原样转发到选中的 worker - 附加 X-Internal-Token 头 - 失败重试（换 worker） - 响应中含 *_base64 图片字段 → 解码落盘 → 改写为 *_url """ import base64 import logging import re import uuid from pathlib import Path from typing import Any, Optional import httpx from fastapi import Request from fastapi.responses import JSONResponse from gateway.pool import ( NoWorkerAvailable, acquire_worker, mark_worker_unhealthy, release_worker, ) logger = logging.getLogger("gateway.forward") # --------------------------------------------------------------------------- # 全局 httpx client（连接复用） # --------------------------------------------------------------------------- _client: Optional[httpx.AsyncClient] = None def get_client() -> httpx.AsyncClient: global _client if _client is None: _client = httpx.AsyncClient() return _client async def close_client(): global _client if _client: await _client.aclose() _client = None # --------------------------------------------------------------------------- # base64 → URL 改写 # --------------------------------------------------------------------------- # data URI 正则：data:image/png;base64,xxxx _DATA_URI_RE = re.compile(r"^data:(image/\w+);base64,(.+)$", re.IGNORECASE) def _decode_base64_value(value: str) -> Optional[bytes]: """解码 base64 值（支持 data URI 和原始 base64 两种格式）。 格式1（data URI）: data:image/png;base64,xxxx 格式2（原始base64）: xxxx（无前缀，自动尝试解码） """ # 优先匹配 data URI m = _DATA_URI_RE.match(value) if m: try: return base64.b64decode(m.group(2)) except Exception: logger.warning("data URI base64 解码失败") return None # 尝试当作原始 base64 解码（排除明显不是 base64 的短字符串） stripped = value.strip() if len(stripped) < 20: return None # 太短，不可能是图片 try: decoded = base64.b64decode(stripped) if len(decoded) >= 50: # 最小合法图片大小 return decoded except Exception: pass return None def rewrite_base64_to_url( obj: Any, public_base_url: str, static_dir: str, ) -> Any: """递归遍历响应 JSON，将所有 *_base64 字段改写为 *_url。 - 识别 key 以 _base64 结尾的字段 - 支持 data: URI 和原始 base64 两种格式 - 解码 base64 → 保存到 static_dir/{uuid}.png - 删除 *_base64 字段，新增 *_url 字段指向公网 URL """ if isinstance(obj, dict): new_dict: Dict[str, Any] = {} for key, value in obj.items(): if key.endswith("_base64") and isinstance(value, str): img_bytes = _decode_base64_value(value) if img_bytes is not None: # 按内容嗅探扩展名：PNG(接口1标注图，含透明) / JPEG(接口2/3/5 照片) ext = "png" if img_bytes[:8] == b"\x89PNG\r\n\x1a\n" else "jpg" filename = f"{uuid.uuid4().hex}.{ext}" filepath = Path(static_dir) / filename filepath.write_bytes(img_bytes) # 构造对外 URL url_key = key[:-7] + "_url" # "xxx_base64" → "xxx_url" new_dict[url_key] = f"{public_base_url}/static/annotations/{filename}" logger.info("base64→URL: %s → %s (%d bytes)", key, new_dict[url_key], len(img_bytes)) continue # 跳过原 _base64 key else: logger.warning("字段 %s 的值无法解码为图片，保留", key) new_dict[key] = value else: new_dict[key] = rewrite_base64_to_url(value, public_base_url, static_dir) return new_dict elif isinstance(obj, list): return [rewrite_base64_to_url(item, public_base_url, static_dir) for item in obj] else: return obj # --------------------------------------------------------------------------- # 请求转发 # --------------------------------------------------------------------------- async def proxy_request(request: Request, path: str) -> JSONResponse: """代理一次请求到 worker，处理重试与 base64 改写。 流程： 1. acquire worker（排队等待空闲） 2. 重构 multipart/form 请求，加 X-Internal-Token 3. 转发到 worker 4. 成功 → release worker → 改写 base64 → 返回 JSONResponse 5. 失败（连接/超时/5xx/401） → mark unhealthy → retry 6. retry 耗尽 / acquire 失败 → 返回 1007 """ from gateway.config import get_config cfg = get_config() dispatch_cfg = cfg["dispatch"] token = cfg["shared_password"] request_timeout = dispatch_cfg["request_timeout_seconds"] max_retries = dispatch_cfg.get("max_retries", 1) retry_on_failure = dispatch_cfg.get("retry_on_failure", True) public_base_url = cfg["public_base_url"] static_dir = cfg["static_dir"] # --- 1. 读取客户端请求体（原始字节，不做解析） --- body = await request.body() # --- 2. 获取 worker（最多重试 max_retries+1 次） --- attempts = max_retries + 1 last_error_response = None for attempt in range(attempts): worker = None try: worker = await acquire_worker(cfg) except NoWorkerAvailable: logger.warning("无可用 worker，返回 1007") return JSONResponse( status_code=503, content={ "code": 1007, "message": "后端服务暂不可用，请稍后重试", "request_id": f"gw-{uuid.uuid4().hex[:8]}", "data": None, }, ) try: client = get_client() # 直接转发原始请求体（不解析、不重建，保证 multipart 原样透传） headers = {"X-Internal-Token": token} ct = request.headers.get("content-type", "") if ct: headers["Content-Type"] = ct resp = await client.request( method="POST", url=f"{worker.url}{path}", content=body, headers=headers, timeout=request_timeout, ) # --- 判断响应 --- if resp.status_code == 401: # worker 鉴权失败 → 视为 worker 异常 logger.warning("Worker %s 返回 401（鉴权失败），标记不健康", worker.url) await mark_worker_unhealthy(worker) await release_worker(worker) if attempt < attempts - 1: continue # retry last_error_response = JSONResponse( status_code=503, content={ "code": 1007, "message": "后端服务暂不可用，请稍后重试", "request_id": f"gw-{uuid.uuid4().hex[:8]}", "data": None, }, ) break if resp.status_code >= 500: # worker 内部错误 → 视为该 worker 异常 logger.warning("Worker %s 返回 %d，标记不健康", worker.url, resp.status_code) await mark_worker_unhealthy(worker) await release_worker(worker) if attempt < attempts - 1: continue last_error_response = JSONResponse( status_code=503, content={ "code": 1007, "message": "后端服务暂不可用，请稍后重试", "request_id": f"gw-{uuid.uuid4().hex[:8]}", "data": None, }, ) break # --- 成功：改写 base64 → URL --- await release_worker(worker) try: worker_json = resp.json() except Exception: logger.warning("Worker %s 返回非 JSON 响应", worker.url) return JSONResponse( status_code=502, content={ "code": 1007, "message": "后端服务响应异常", "request_id": f"gw-{uuid.uuid4().hex[:8]}", "data": None, }, ) # 递归改写 base64 图片字段 rewritten = rewrite_base64_to_url(worker_json, public_base_url, static_dir) return JSONResponse( status_code=resp.status_code, content=rewritten, ) except (httpx.ConnectError, httpx.ConnectTimeout, httpx.ReadTimeout, httpx.RemoteProtocolError) as exc: logger.warning("Worker %s 连接失败: %s", worker.url, exc) await mark_worker_unhealthy(worker) await release_worker(worker) if attempt < attempts - 1: continue last_error_response = JSONResponse( status_code=503, content={ "code": 1007, "message": "后端服务暂不可用，请稍后重试", "request_id": f"gw-{uuid.uuid4().hex[:8]}", "data": None, }, ) break except Exception as exc: logger.error("转发异常: %s", exc) await release_worker(worker) if attempt < attempts - 1: continue last_error_response = JSONResponse( status_code=503, content={ "code": 1007, "message": "后端服务暂不可用，请稍后重试", "request_id": f"gw-{uuid.uuid4().hex[:8]}", "data": None, }, ) break # --- 所有尝试耗尽 --- return last_error_response or JSONResponse( status_code=503, content={ "code": 1007, "message": "后端服务暂不可用，请稍后重试", "request_id": f"gw-{uuid.uuid4().hex[:8]}", "data": None, }, )