301 lines
12 KiB
Python
301 lines
12 KiB
Python
"""数据源代理：服务端替浏览器请求第三方接口
|
||
|
||
两个入口共用 fetch_json：编辑态由登录用户传 URL；线上态只传数据源 ID，配置取自已发布快照
|
||
防护：只允许 http/https；DNS 解析后每个地址都必须是公网地址，并把连接钉在已检查过的地址上
|
||
（不再让 httpx 二次解析，防 DNS 重绑定）；不读环境代理；重定向逐跳重新检查，跨域时丢掉调用方的密钥头；
|
||
总超时 5 秒，响应体上限 2 MB，按解压后的字节计，解压时带输出上限（防压缩炸弹）
|
||
"""
|
||
|
||
import asyncio
|
||
import ipaddress
|
||
import json
|
||
import re
|
||
import socket
|
||
import time
|
||
import zlib
|
||
from collections.abc import Awaitable, Callable
|
||
from typing import Any, Literal
|
||
from urllib.parse import urljoin
|
||
|
||
import httpx
|
||
|
||
from app.core.exceptions import BizException
|
||
|
||
TIMEOUT = 5.0
|
||
MAX_BYTES = 2 * 1024 * 1024
|
||
MAX_REDIRECTS = 3
|
||
REDIRECT_STATUS = {301, 302, 303, 307, 308}
|
||
# 这些头由 httpx 或连接本身决定，不让调用方覆盖
|
||
FORBIDDEN_HEADERS = {
|
||
"host",
|
||
"content-length",
|
||
"connection",
|
||
"transfer-encoding",
|
||
"upgrade",
|
||
"te",
|
||
"keep-alive",
|
||
"proxy-authorization",
|
||
"proxy-connection",
|
||
"accept-encoding",
|
||
}
|
||
# 只接受自己能带上限解压的编码；httpx 默认还会协商 br / zstd，那些解码器没有输出上限
|
||
ACCEPT_ENCODING = "gzip, deflate"
|
||
# is_global 会放过的几段：NAT64（可以嵌入任意 IPv4）、IPv4 兼容地址、已废弃的站点本地地址
|
||
EXTRA_BLOCKED = [
|
||
ipaddress.ip_network(net)
|
||
for net in ("64:ff9b::/96", "64:ff9b:1::/48", "::/96", "fec0::/10")
|
||
]
|
||
# 十进制、十六进制、省略写法的 IPv4（2130706433、0x7f.1、127.1）：不是标准写法，
|
||
# 各平台的解析器解释不一，一律当作内网地址拒绝，不交给 DNS
|
||
NUMERIC_LABEL = re.compile(r"0x[0-9a-f]+|[0-9]+", re.IGNORECASE)
|
||
|
||
Method = Literal["GET", "POST"]
|
||
|
||
|
||
def _blocked() -> BizException:
|
||
return BizException(422, "不允许访问内网或本机地址")
|
||
|
||
|
||
def _is_public(address: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
|
||
# IPv4 映射的 IPv6 地址按内嵌的 IPv4 判断
|
||
if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped is not None:
|
||
address = address.ipv4_mapped
|
||
return (
|
||
address.is_global
|
||
and not address.is_multicast
|
||
and not any(address in net for net in EXTRA_BLOCKED)
|
||
)
|
||
|
||
|
||
def _checked(
|
||
address: ipaddress.IPv4Address | ipaddress.IPv6Address,
|
||
) -> ipaddress.IPv4Address | ipaddress.IPv6Address:
|
||
if not _is_public(address):
|
||
raise _blocked()
|
||
return address
|
||
|
||
|
||
async def _resolve(host: str, port: int) -> str:
|
||
"""解析主机名，所有地址都是公网才放行，返回第一个地址"""
|
||
host = host.rstrip(".")
|
||
try:
|
||
literal = ipaddress.ip_address(host)
|
||
except ValueError:
|
||
pass
|
||
else:
|
||
return str(_checked(literal))
|
||
if host.lower() == "localhost" or host.lower().endswith(".localhost"):
|
||
raise _blocked()
|
||
if all(NUMERIC_LABEL.fullmatch(label) for label in host.split(".")):
|
||
raise _blocked()
|
||
try:
|
||
infos = await asyncio.get_running_loop().getaddrinfo(
|
||
host, port, type=socket.SOCK_STREAM
|
||
)
|
||
except socket.gaierror:
|
||
raise BizException(502, "数据源域名无法解析") from None
|
||
addresses = [ipaddress.ip_address(info[4][0].split("%")[0]) for info in infos]
|
||
if not addresses:
|
||
raise BizException(502, "数据源域名无法解析")
|
||
for address in addresses:
|
||
_checked(address)
|
||
return str(addresses[0])
|
||
|
||
|
||
def _clean_headers(headers: dict[str, Any] | None) -> dict[str, str]:
|
||
cleaned = {"Accept": "application/json", "Accept-Encoding": ACCEPT_ENCODING}
|
||
for name, value in (headers or {}).items():
|
||
if name.lower() in FORBIDDEN_HEADERS or not isinstance(
|
||
value, (str, int, float)
|
||
):
|
||
continue
|
||
cleaned[name] = str(value)
|
||
return cleaned
|
||
|
||
|
||
def _query_value(value: Any) -> Any:
|
||
return value if isinstance(value, (str, int, float, bool)) else json.dumps(value)
|
||
|
||
|
||
def _parse_url(raw: str) -> httpx.URL:
|
||
try:
|
||
url = httpx.URL(raw)
|
||
except httpx.InvalidURL:
|
||
raise BizException(422, "数据源地址不合法") from None
|
||
if url.scheme not in ("http", "https") or not url.host or url.userinfo:
|
||
raise BizException(422, "数据源地址只支持 http / https，且不能带账号密码")
|
||
return url
|
||
|
||
|
||
def _too_large() -> BizException:
|
||
return BizException(502, "数据源响应超过 2 MB")
|
||
|
||
|
||
async def _read_body(response: httpx.Response) -> bytes:
|
||
"""读响应体：原始字节和解压后的字节都不超过 MAX_BYTES
|
||
|
||
不用 aiter_bytes：它对每个 64 KB 的块整块解压，没有输出上限，一个块就能膨胀到上 GB
|
||
"""
|
||
encoding = response.headers.get("content-encoding", "identity").strip().lower()
|
||
if encoding in ("", "identity"):
|
||
decoder = None
|
||
elif encoding in ("gzip", "x-gzip"):
|
||
decoder = zlib.decompressobj(16 + zlib.MAX_WBITS)
|
||
elif encoding == "deflate":
|
||
decoder = zlib.decompressobj()
|
||
else:
|
||
raise BizException(502, "数据源响应的压缩格式不支持")
|
||
declared = response.headers.get("content-length", "")
|
||
if declared.isdigit() and int(declared) > MAX_BYTES:
|
||
raise _too_large()
|
||
chunks: list[bytes] = []
|
||
raw_size = size = 0
|
||
try:
|
||
async for chunk in response.aiter_raw():
|
||
raw_size += len(chunk)
|
||
if raw_size > MAX_BYTES:
|
||
raise _too_large()
|
||
if decoder is None:
|
||
chunks.append(chunk)
|
||
size += len(chunk)
|
||
continue
|
||
pending = chunk
|
||
while pending:
|
||
# max_length 为 0 表示不限，这里恒大于 0
|
||
part = decoder.decompress(pending, MAX_BYTES - size + 1)
|
||
size += len(part)
|
||
if size > MAX_BYTES:
|
||
raise _too_large()
|
||
chunks.append(part)
|
||
pending = decoder.unconsumed_tail
|
||
if not part:
|
||
break
|
||
# 压缩流没有正常结束（被截断）也算失败，别把半截数据当成功
|
||
if decoder is not None and not decoder.eof:
|
||
raise BizException(502, "数据源响应解压失败")
|
||
except zlib.error:
|
||
raise BizException(502, "数据源响应解压失败") from None
|
||
return b"".join(chunks)
|
||
|
||
|
||
async def _request(
|
||
client: httpx.AsyncClient,
|
||
method: str,
|
||
url: httpx.URL,
|
||
headers: dict[str, str],
|
||
body: Any,
|
||
) -> tuple[int, httpx.Headers, bytes]:
|
||
port = url.port or (443 if url.scheme == "https" else 80)
|
||
pinned = url.copy_with(host=await _resolve(url.host, port))
|
||
request = client.build_request(
|
||
method,
|
||
pinned,
|
||
headers={**headers, "Host": url.netloc.decode()},
|
||
json=body if method == "POST" else None,
|
||
# 连接用 IP，TLS 的 SNI 和证书校验仍然用原主机名
|
||
extensions={"sni_hostname": url.host},
|
||
)
|
||
response = await client.send(request, stream=True)
|
||
try:
|
||
return response.status_code, response.headers, await _read_body(response)
|
||
finally:
|
||
await response.aclose()
|
||
|
||
|
||
def _origin(url: httpx.URL) -> tuple[str, str, int | None]:
|
||
return url.scheme, url.host, url.port or (443 if url.scheme == "https" else 80)
|
||
|
||
|
||
async def fetch_json(
|
||
url: str,
|
||
method: Method = "GET",
|
||
params: dict[str, Any] | None = None,
|
||
headers: dict[str, Any] | None = None,
|
||
) -> dict[str, Any] | list[Any]:
|
||
"""GET 的 params 作为查询参数，POST 的 params 作为 JSON 请求体，和前端 fetchData 一致"""
|
||
target = _parse_url(url)
|
||
request_headers = _clean_headers(headers)
|
||
body = None
|
||
if method == "GET":
|
||
if params:
|
||
target = target.copy_merge_params(
|
||
{key: _query_value(value) for key, value in params.items()}
|
||
)
|
||
else:
|
||
body = params or {}
|
||
try:
|
||
async with asyncio.timeout(TIMEOUT):
|
||
# trust_env=False：环境里的代理会绕过上面的地址检查
|
||
# max_keepalive_connections=0：连接按 IP 复用，不关的话同一 IP 上的下一跳会沿用上一跳的 TLS 校验结果
|
||
async with httpx.AsyncClient(
|
||
trust_env=False,
|
||
timeout=TIMEOUT,
|
||
limits=httpx.Limits(max_keepalive_connections=0),
|
||
) as client:
|
||
for _ in range(MAX_REDIRECTS + 1):
|
||
status, response_headers, content = await _request(
|
||
client, method, target, request_headers, body
|
||
)
|
||
location = response_headers.get("location")
|
||
if status not in REDIRECT_STATUS or not location:
|
||
break
|
||
previous = target
|
||
target = _parse_url(urljoin(str(target), location))
|
||
if _origin(target) != _origin(previous):
|
||
# 密钥头只发给调用方指定的那个源
|
||
request_headers = _clean_headers(None)
|
||
if status not in (307, 308):
|
||
method, body = "GET", None
|
||
else:
|
||
raise BizException(502, "数据源重定向次数过多")
|
||
except (TimeoutError, httpx.TimeoutException):
|
||
raise BizException(504, "数据源请求超时") from None
|
||
except httpx.HTTPError:
|
||
raise BizException(502, "数据源请求失败") from None
|
||
if not 200 <= status < 300:
|
||
raise BizException(502, f"数据源返回 HTTP {status}")
|
||
try:
|
||
data = json.loads(content)
|
||
except (ValueError, RecursionError):
|
||
# 2 MB 以内的深层嵌套会打满递归，同样当作非法 JSON
|
||
raise BizException(502, "数据源返回的不是 JSON") from None
|
||
if not isinstance(data, (dict, list)):
|
||
raise BizException(502, "数据源返回的不是 JSON 对象或数组")
|
||
return data
|
||
|
||
|
||
class FetchCache:
|
||
"""线上大屏多人同时开、按 interval 轮询：同一个 key 几秒内只真正请求一次，并发的请求共用同一次
|
||
|
||
失败不缓存；单实例部署，缓存放进程内存。单条最大 2 MB，条数上限乘出来就是内存上限，所以条数要小
|
||
"""
|
||
|
||
def __init__(self, ttl: float = 3.0, max_entries: int = 64):
|
||
self._ttl = ttl
|
||
self._max_entries = max_entries
|
||
self._entries: dict[Any, tuple[float, Any]] = {}
|
||
self._inflight: dict[Any, asyncio.Future] = {}
|
||
|
||
async def get(self, key: Any, fetch: Callable[[], Awaitable[Any]]) -> Any:
|
||
now = time.monotonic()
|
||
hit = self._entries.get(key)
|
||
if hit is not None and hit[0] > now:
|
||
return hit[1]
|
||
future = self._inflight.get(key)
|
||
if future is None:
|
||
future = asyncio.ensure_future(fetch())
|
||
self._inflight[key] = future
|
||
future.add_done_callback(lambda done: self._finish(key, done))
|
||
# shield：某个请求者断开不能取消共用的那次请求
|
||
return await asyncio.shield(future)
|
||
|
||
def _finish(self, key: Any, future: asyncio.Future) -> None:
|
||
self._inflight.pop(key, None)
|
||
if future.cancelled() or future.exception() is not None:
|
||
return
|
||
now = time.monotonic()
|
||
# 每次写入先丢掉过期项；清完还满就不缓存这一条
|
||
self._entries = {k: v for k, v in self._entries.items() if v[0] > now}
|
||
if len(self._entries) < self._max_entries:
|
||
self._entries[key] = (now + self._ttl, future.result())
|