"""公開用アプリ。バイト単位の無改変中継。 書き換えるのは ``Authorization``（内部キーへ差し替え）と ``Host`` のみ。 リクエスト本文は受信した ``bytes`` をそのまま送り、JSON として読み直して 再シリアライズすることは一切しない。応答は ``aiter_raw()`` で受け、 ``Content-Encoding`` を保ったまま素通しする。 本文に触れる例外は 2 つだけで、いずれも意図的なもの。 * ``/v1/models`` の応答絞り込み（``FILTER_MODELS_RESPONSE``） * 4xx / 5xx の小さな本文に内部キーが混入した場合の伏字化 """ from __future__ import annotations import asyncio import contextlib import gzip import json import logging import time import zlib from typing import AsyncIterator, Callable, Iterable, Sequence import httpx from starlette.applications import Starlette from starlette.requests import Request from starlette.responses import Response, StreamingResponse from starlette.routing import Route from . import errors, policy from .config import Settings from .db import LogWriter, RequestLogEntry, iso_jst from .keys import KeyCache, KeyRecord from .limits import RateLimiter, Reservation from .logging_utils import contains_secret, register_secret, scrub, scrub_bytes from .upstream import UpstreamClient logger = logging.getLogger(__name__) #: RFC 上の hop-by-hop ヘッダ。``proxy-`` 始まりも併せて落とす。 HOP_BY_HOP: frozenset[str] = frozenset( {"connection", "keep-alive", "te", "trailer", "transfer-encoding", "upgrade"} ) #: 中継時に必ず落とすヘッダ。差し替えるか、httpx が正しい値を組み立てる。 _REPLACED_REQUEST_HEADERS: frozenset[str] = frozenset({"authorization", "host", "content-length"}) #: 内部キー混入を走査する応答本文の上限。 ERROR_SCAN_LIMIT = 64 * 1024 #: ``/v1/models`` を絞り込むために読み込む上限。 MODELS_BUFFER_LIMIT = 4 * 1024 * 1024 #: 非ストリーミング応答から ``usage`` を拾うための上限。 USAGE_BUFFER_LIMIT = 1024 * 1024 #: SSE から末尾の ``usage`` を拾うために保持する量。 SSE_TAIL_LIMIT = 64 * 1024 _IDENTITY_ENCODINGS = frozenset({"", "identity"}) LogCallback = Callable[..., None] # ------------------------------------------------------------------ ヘッダ操作 def _connection_tokens(headers: Iterable[tuple[bytes, bytes]]) -> set[str]: """``Connection`` に列挙されたヘッダ名（RFC 上 hop-by-hop）。""" tokens: set[str] = set() for name, value in headers: if name.decode("latin-1").lower() != "connection": continue for token in value.decode("latin-1").split(","): token = token.strip().lower() if token and token not in ("close", "keep-alive"): tokens.add(token) return tokens def filter_request_headers( raw_headers: Sequence[tuple[bytes, bytes]], *, upstream_api_key: str, ) -> list[tuple[bytes, bytes]]: """クライアントのヘッダから上流向けヘッダを組み立てる。 順序と重複はそのまま保つ。落とすのは hop-by-hop、``proxy-*``、 クライアント由来の ``x-forwarded-*``、および差し替え対象のみ。 """ drop = set(HOP_BY_HOP) | _connection_tokens(raw_headers) out: list[tuple[bytes, bytes]] = [] has_accept_encoding = False for name, value in raw_headers: lowered = name.decode("latin-1").lower() if lowered in drop or lowered.startswith("proxy-"): continue if lowered.startswith("x-forwarded-"): continue if lowered in _REPLACED_REQUEST_HEADERS: continue if lowered == "accept-encoding": has_accept_encoding = True out.append((name, value)) if not has_accept_encoding: # httpx の既定は ``gzip, deflate, br`` であり、クライアントが解釈できない # 応答を招く。明示的に無圧縮を要求する。 out.append((b"accept-encoding", b"identity")) if upstream_api_key: out.append((b"authorization", f"Bearer {upstream_api_key}".encode("latin-1"))) return out def force_identity_encoding(headers: list[tuple[bytes, bytes]]) -> list[tuple[bytes, bytes]]: """``Accept-Encoding`` を ``identity`` に固定する（応答を読む経路専用）。""" out = [(n, v) for n, v in headers if n.decode("latin-1").lower() != "accept-encoding"] out.append((b"accept-encoding", b"identity")) return out def filter_response_headers(response: httpx.Response) -> list[tuple[bytes, bytes]]: """上流の応答ヘッダから hop-by-hop を落とす。それ以外は素通し。""" raw = list(response.headers.raw) drop = set(HOP_BY_HOP) | _connection_tokens(raw) out: list[tuple[bytes, bytes]] = [] for name, value in raw: lowered = name.decode("latin-1").lower() if lowered in drop or lowered.startswith("proxy-"): continue out.append((name, value)) return out def client_ip_from(request: Request) -> str | None: """``X-Forwarded-For`` の最右端。無ければ接続元。 最左端はクライアントが自由に詐称できるため採用しない。 """ values: list[str] = [] for raw in request.headers.getlist("x-forwarded-for"): values.extend(part.strip() for part in raw.split(",")) for candidate in reversed(values): if candidate: return candidate return request.client.host if request.client else None def _raw_path_and_query(request: Request) -> tuple[bytes, bytes]: """受信した生のパスとクエリを取り出す。転送はこの値をそのまま使う。""" scope = request.scope raw_path = scope.get("raw_path") or scope["path"].encode("utf-8", "surrogateescape") if b"?" in raw_path: raw_path = raw_path.split(b"?", 1)[0] return raw_path, scope.get("query_string", b"") or b"" def _content_encoding(response: httpx.Response) -> str: return (response.headers.get("content-encoding") or "").strip().lower() def _decode_body(raw: bytes, encoding: str) -> bytes | None: """伏字走査のために本文を復号する。扱えない符号化は ``None``。""" if encoding in _IDENTITY_ENCODINGS: return raw try: if encoding == "gzip": return gzip.decompress(raw) if encoding == "deflate": return zlib.decompress(raw) except Exception: return None return None # ------------------------------------------------------------------ 本文読み取り async def read_body_capped(request: Request, limit: int) -> bytes: """上限付きでリクエスト本文を読む。超過は 413。""" declared = request.headers.get("content-length") if declared and declared.isdigit() and int(declared) > limit: raise errors.payload_too_large( f"Request body exceeds the maximum allowed size of {limit} bytes." ) chunks: list[bytes] = [] total = 0 async for chunk in request.stream(): if not chunk: continue total += len(chunk) if total > limit: raise errors.payload_too_large( f"Request body exceeds the maximum allowed size of {limit} bytes." ) chunks.append(chunk) return b"".join(chunks) async def drain(source: AsyncIterator[bytes], limit: int) -> tuple[bytes, bool]: """応答本文を上限まで読み込む。戻り値は (読み込んだ内容, 打ち切ったか)。 打ち切った場合、読み込んだ内容は続きへ連結する前置きとして扱う。 """ buffer = bytearray() async for chunk in source: buffer.extend(chunk) if len(buffer) > limit: return bytes(buffer), True return bytes(buffer), False # ------------------------------------------------------------------ 使用量集計 class UsageSniffer: """応答から ``usage`` を拾う。本文は改変しない。""" __slots__ = ("_sse", "_buffer", "_enabled", "prompt_tokens", "completion_tokens") def __init__(self, *, sse: bool, enabled: bool) -> None: self._sse = sse self._enabled = enabled self._buffer = bytearray() self.prompt_tokens: int | None = None self.completion_tokens: int | None = None def feed(self, chunk: bytes) -> None: if not self._enabled: return self._buffer.extend(chunk) if self._sse: if len(self._buffer) > SSE_TAIL_LIMIT: del self._buffer[: len(self._buffer) - SSE_TAIL_LIMIT] elif len(self._buffer) > USAGE_BUFFER_LIMIT: self._enabled = False self._buffer.clear() def finish(self) -> None: if not self._enabled or not self._buffer: return raw = bytes(self._buffer) usage = _find_usage_sse(raw) if self._sse else _find_usage_json(raw) if usage: prompt = usage.get("prompt_tokens") completion = usage.get("completion_tokens") if isinstance(prompt, int): self.prompt_tokens = prompt if isinstance(completion, int): self.completion_tokens = completion self._buffer.clear() @property def total_tokens(self) -> int: return (self.prompt_tokens or 0) + (self.completion_tokens or 0) def _find_usage_json(raw: bytes) -> dict | None: try: payload = json.loads(raw) except (ValueError, UnicodeDecodeError): return None if isinstance(payload, dict): usage = payload.get("usage") if isinstance(usage, dict): return usage return None def _find_usage_sse(raw: bytes) -> dict | None: """SSE の末尾から ``usage`` を探す。クライアントが要求した場合のみ現れる。""" for line in reversed(raw.split(b"\n")): line = line.strip() if not line.startswith(b"data:"): continue payload_raw = line[5:].strip() if not payload_raw or payload_raw == b"[DONE]": continue try: payload = json.loads(payload_raw) except (ValueError, UnicodeDecodeError): continue if isinstance(payload, dict): usage = payload.get("usage") if isinstance(usage, dict): return usage return None # ------------------------------------------------------------------ 応答の組立 def raw_response(body: bytes, status_code: int, headers: Sequence[tuple[bytes, bytes]]) -> Response: """ヘッダの順序と重複を保ったまま応答を組み立てる。""" response = Response(status_code=status_code) response.body = body response.raw_headers = list(headers) return response def streaming_response( iterator: AsyncIterator[bytes], status_code: int, headers: Sequence[tuple[bytes, bytes]] ) -> StreamingResponse: """ヘッダの順序と重複を保ったままストリーミング応答を組み立てる。""" response = StreamingResponse(iterator, status_code=status_code) response.raw_headers = list(headers) return response def replace_body_headers( headers: Sequence[tuple[bytes, bytes]], length: int, *, drop_encoding: bool = False ) -> list[tuple[bytes, bytes]]: """本文を差し替えた場合のヘッダ補正。""" out: list[tuple[bytes, bytes]] = [] for name, value in headers: lowered = name.decode("latin-1").lower() if lowered == "content-length": continue if drop_encoding and lowered == "content-encoding": continue out.append((name, value)) out.append((b"content-length", str(length).encode("latin-1"))) return out # ------------------------------------------------------------------ アプリ本体 def create_proxy_app( *, settings: Settings, keys: KeyCache, limiter: RateLimiter, upstream: UpstreamClient, logs: LogWriter | None = None, ) -> Starlette: """公開ポート用の Starlette アプリを組み立てる。""" # 伏字処理はアプリの一部として常時有効にする。起動経路に依存させない。 register_secret(settings.upstream_api_key, settings.control_token) async def handler(request: Request) -> Response: started = time.monotonic() raw_path, query_string = _raw_path_and_query(request) normalized = policy.normalize_path(raw_path) client_ip = client_ip_from(request) state: dict[str, object] = {"key": None, "model": None} def log( status: int, *, prompt_tokens: int | None = None, completion_tokens: int | None = None, ) -> None: if logs is None: return key = state["key"] logs.enqueue( RequestLogEntry( key_id=key.id if isinstance(key, KeyRecord) else None, ts=iso_jst(), method=request.method, path=normalized, model=state["model"], # type: ignore[arg-type] status=status, duration_ms=int((time.monotonic() - started) * 1000), prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, client_ip=client_ip, ) ) stack = contextlib.AsyncExitStack() try: async with stack: policy.check_path(normalized, settings) key = _authenticate(request, keys) state["key"] = key body = await read_body_capped(request, settings.max_request_body_bytes) policy.check_admission(body, settings) reservation: Reservation = await limiter.acquire(key) stack.callback(reservation.release) if policy.requires_model_check(normalized): state["model"] = policy.enforce_model( body, request.headers.get("content-type", ""), key.models ) filter_models = ( settings.filter_models_response and policy.is_models_path(normalized) and bool(key.models) and request.method != "HEAD" ) headers = filter_request_headers( request.headers.raw, upstream_api_key=settings.upstream_api_key ) if filter_models: # 絞り込みのため復号可能な形で受け取る。この経路のみ本文を読み書きする。 headers = force_identity_encoding(headers) upstream_request = upstream.build_request( method=request.method, raw_path=raw_path, query_string=query_string, headers=headers, content=body, ) deadline = asyncio.get_running_loop().time() + settings.max_request_duration try: async with asyncio.timeout_at(deadline): response = await upstream.send(upstream_request) except (httpx.ReadTimeout, httpx.WriteTimeout, httpx.PoolTimeout) as exc: upstream.mark_unreachable() logger.warning("上流が応答しませんでした: %s", scrub(str(exc))) log(504) return errors.gateway_timeout().to_response() except httpx.HTTPError as exc: upstream.mark_unreachable() logger.warning("上流へ到達できませんでした: %s", scrub(str(exc))) log(502) return errors.bad_gateway().to_response() stack.push_async_callback(response.aclose) out_headers = filter_response_headers(response) encoding = _content_encoding(response) is_sse = (response.headers.get("content-type") or "").lower().startswith( "text/event-stream" ) prefix = b"" if filter_models and response.status_code == 200: async with asyncio.timeout_at(deadline): raw, truncated = await drain(response.aiter_raw(), MODELS_BUFFER_LIMIT) if not truncated: filtered = policy.filter_models_payload(raw, key.models) log(response.status_code) if filtered is None: return raw_response(raw, response.status_code, out_headers) return raw_response( filtered, response.status_code, replace_body_headers(out_headers, len(filtered)), ) logger.warning("/v1/models の応答が大きすぎるため絞り込みを見送りました") prefix = raw elif response.status_code >= 400 and request.method != "HEAD": declared = response.headers.get("content-length") too_big = bool(declared and declared.isdigit() and int(declared) > ERROR_SCAN_LIMIT) if not too_big: async with asyncio.timeout_at(deadline): raw, truncated = await drain(response.aiter_raw(), ERROR_SCAN_LIMIT) if not truncated: log(response.status_code) return _scrubbed_error_response( raw, response.status_code, out_headers, encoding ) prefix = raw released = stack.pop_all() sniffer = UsageSniffer(sse=is_sse, enabled=encoding in _IDENTITY_ENCODINGS) return streaming_response( _stream_body( response=response, released=released, sniffer=sniffer, limiter=limiter, key_id=key.id, deadline=deadline, log=log, prefix=prefix, ), response.status_code, out_headers, ) except errors.GatewayError as exc: log(exc.status_code) return exc.to_response() except (asyncio.TimeoutError, TimeoutError): logger.warning("最大所要時間を超えました") log(504) return errors.gateway_timeout( "The request exceeded the maximum allowed duration.", code="request_timeout" ).to_response() except Exception: logger.exception("中継中に予期しない例外が発生しました") log(502) return errors.bad_gateway( "The gateway encountered an internal error.", code="internal_error" ).to_response() routes = [ Route( "/{path:path}", handler, methods=["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"], ) ] app = Starlette(routes=routes) app.state.settings = settings app.state.keys = keys app.state.limiter = limiter app.state.upstream = upstream return app def _scrubbed_error_response( raw: bytes, status_code: int, out_headers: Sequence[tuple[bytes, bytes]], encoding: str ) -> Response: """エラー本文に内部キーが混入していれば伏字にして返す。 ストリーミング本体には適用しない。混入が無ければ受信したバイト列をそのまま返す。 """ plain = _decode_body(raw, encoding) if plain is not None and contains_secret(plain): logger.error("上流のエラー応答に内部キーが含まれていたため伏字にしました") replaced = scrub_bytes(plain) return raw_response( replaced, status_code, replace_body_headers(out_headers, len(replaced), drop_encoding=True), ) return raw_response(raw, status_code, out_headers) def _authenticate(request: Request, keys: KeyCache) -> KeyRecord: """``Authorization: Bearer`` を照合する。""" header = request.headers.get("authorization", "") scheme, _, token = header.partition(" ") if scheme.lower() != "bearer" or not token.strip(): raise errors.unauthorized( "You didn't provide an API key. You need to provide your API key in an" " Authorization header using Bearer auth (i.e. Authorization: Bearer YOUR_KEY)." ) record = keys.resolve(token.strip()) if record is None or not record.is_active(): raise errors.unauthorized() return record async def _stream_body( *, response: httpx.Response, released: contextlib.AsyncExitStack, sniffer: UsageSniffer, limiter: RateLimiter, key_id: int, deadline: float, log: LogCallback, prefix: bytes = b"", ) -> AsyncIterator[bytes]: """上流の生バイト列をそのまま流す。解放は必ず ``finally`` で行う。""" status = response.status_code try: async with released: try: if prefix: sniffer.feed(prefix) yield prefix async with asyncio.timeout_at(deadline): async for chunk in response.aiter_raw(): sniffer.feed(chunk) yield chunk except (asyncio.TimeoutError, TimeoutError): status = 504 logger.warning("最大所要時間を超えたためストリームを打ち切りました") except (httpx.ReadTimeout, httpx.WriteTimeout): status = 504 logger.warning("無通信時間の上限を超えたためストリームを打ち切りました") except httpx.HTTPError as exc: status = 502 logger.warning("ストリーム中に上流との接続が切れました: %s", scrub(str(exc))) finally: sniffer.finish() if sniffer.total_tokens: limiter.record_tokens(key_id, sniffer.total_tokens) log(status, prompt_tokens=sniffer.prompt_tokens, completion_tokens=sniffer.completion_tokens)