Upload 2 files

router_proxy.py

@@ -10,8 +10,6 @@ import ipaddress
 import httpx
 import os
 import re
-import asyncio
-import socket
 import urllib.request
 import urllib.error
 import 数据库连接 as json_db
@@ -22,7 +20,7 @@ from proxy_access import authorized_proxy_item, ProxyAccessError
 from github_source import repository_parts, select_github_token
 from github_archive import ARCHIVE_MAX_BYTES
 from 安全认证 import require_auth, JWT_SECRET
-from safe_http_transport import ValidatedAsyncHTTPTransport
+from safe_http_transport import ValidatedAsyncHTTPTransport, _resolve_public_ips, PublicResolutionError
 from workflow_download_service import prepare_workflow_download
 from zip_stream import ZIP_DOWNLOAD_MAX_BYTES
 from github_zip_service import prepare_github_zip
@@ -47,20 +45,10 @@ async def _validate_public_http_target(url: str) -> str:
 
     host = parsed.hostname.rstrip(".").lower()
     try:
-        addresses = [ipaddress.ip_address(host)]
-    except ValueError:
-        try:
-            infos = await asyncio.to_thread(
-                socket.getaddrinfo, host, port or (443 if parsed.scheme == "https" else 80),
-                type=socket.SOCK_STREAM,
-            )
-        except OSError:
-            raise HTTPException(status_code=400, detail="下载地址无法解析")
-        addresses = [ipaddress.ip_address(info[4][0].split("%", 1)[0]) for info in infos]
-    if not addresses or any(not address.is_global or address.is_multicast
-                            or (getattr(address, 'ipv4_mapped', None) is not None
-                                and address.ipv4_mapped.is_multicast) for address in addresses):
-        raise HTTPException(status_code=400, detail="下载地址解析到非公网地址")
+        await _resolve_public_ips(host, port or (443 if parsed.scheme == "https" else 80))
+    except PublicResolutionError as error:
+        detail = "下载地址无法解析" if error.reason == "dns_unavailable" else "下载地址解析到非公网地址"
+        raise HTTPException(status_code=400, detail=detail) from None
     return host
 
 

safe_http_transport.py

@@ -2,13 +2,114 @@
 from contextvars import ContextVar
 import asyncio
 import ipaddress
+import json
 import socket
 
 import httpcore
 import httpx
+from proxy_resources import ProxyResources
 
 
 _validated_ip = ContextVar("ranking_validated_ip", default=None)
+_DNS_MAX_BYTES = 16384
+_DNS_TIMEOUT = 8
+
+
+class PublicResolutionError(httpx.ConnectError):
+    def __init__(self, reason):
+        self.reason = reason
+        super().__init__("目标域名无法解析" if reason == "dns_unavailable" else "目标域名解析到非公网地址")
+
+
+def _public(address):
+    mapped = getattr(address, "ipv4_mapped", None)
+    return (address.is_global and not address.is_multicast
+            and (mapped is None or (mapped.is_global and not mapped.is_multicast)))
+
+
+def _dns_object(pairs):
+    result = {}
+    for key, value in pairs:
+        if key in result:
+            raise ValueError()
+        result[key] = value
+    return result
+
+
+def _dns_addresses(body, host, record_type):
+    data = json.loads(body.decode("utf-8"), object_pairs_hook=_dns_object,
+                      parse_constant=lambda _: (_ for _ in ()).throw(ValueError()))
+    if (not isinstance(data, dict) or type(data.get("Status")) is not int
+            or data["Status"] != 0 or data.get("TC") is not False):
+        raise ValueError()
+    question = data.get("Question")
+    if (not isinstance(question, list) or len(question) != 1
+            or not isinstance(question[0], dict)
+            or question[0].get("name", "").rstrip(".").lower() != host
+            or type(question[0].get("type")) is not int or question[0]["type"] != record_type):
+        raise ValueError()
+    answers = data.get("Answer", [])
+    if not isinstance(answers, list) or len(answers) > 32:
+        raise ValueError()
+    names = {host}
+    for _ in answers:
+        for answer in answers:
+            if not isinstance(answer, dict) or type(answer.get("type")) is not int:
+                raise ValueError()
+            if answer["type"] == 5 and answer.get("name", "").rstrip(".").lower() in names:
+                alias = answer.get("data")
+                if not isinstance(alias, str) or not alias or len(alias) > 253:
+                    raise ValueError()
+                names.add(alias.rstrip(".").lower())
+    addresses = []
+    for answer in answers:
+        if answer["type"] in (1, 28):
+            address = ipaddress.ip_address(answer.get("data"))
+            if not _public(address) or address.version != (4 if answer["type"] == 1 else 6):
+                raise ValueError()
+            if (answer["type"] == record_type
+                    and answer.get("name", "").rstrip(".").lower() in names):
+                addresses.append(str(address))
+    return tuple(dict.fromkeys(addresses))
+
+
+async def _resolve_hf_public_ips(host):
+    client = httpx.AsyncClient(transport=_DNSHTTPTransport(), verify=True, trust_env=False,
+                               follow_redirects=False, timeout=_DNS_TIMEOUT)
+    resources = ProxyResources(client)
+    try:
+        async with asyncio.timeout(_DNS_TIMEOUT):
+            for record_type in (1, 28):
+                request = client.build_request("GET", "https://cloudflare-dns.com/dns-query",
+                    params={"name": host, "type": str(record_type)},
+                    headers={"Accept": "application/dns-json", "Accept-Encoding": "identity"})
+                response = resources.track(await client.send(request, stream=True, follow_redirects=False))
+                if (response.status_code != 200
+                        or response.headers.get("content-encoding", "identity").lower() != "identity"
+                        or response.headers.get("content-type", "").split(";", 1)[0].lower()
+                        not in ("application/dns-json", "application/json")):
+                    raise ValueError()
+                length = response.headers.get("content-length")
+                if length is not None and (not length.isascii() or not length.isdecimal()
+                                           or len(length) > 6 or int(length) > _DNS_MAX_BYTES):
+                    raise ValueError()
+                body = bytearray()
+                async for chunk in response.aiter_raw(4096):
+                    if len(body) + len(chunk) > _DNS_MAX_BYTES:
+                        raise ValueError()
+                    body.extend(chunk)
+                if length is not None and len(body) != int(length):
+                    raise ValueError()
+                addresses = _dns_addresses(body, host, record_type)
+                await resources.close_response(response)
+                if addresses:
+                    return addresses
+        raise ValueError()
+    except (ValueError, TypeError, AttributeError, UnicodeError, RecursionError,
+            httpx.HTTPError, TimeoutError):
+        raise PublicResolutionError("dns_unavailable") from None
+    finally:
+        await resources.close()
 
 
 async def _resolve_public_ips(host: str, port: int) -> tuple[str, ...]:
@@ -22,16 +123,18 @@ async def _resolve_public_ips(host: str, port: int) -> tuple[str, ...]:
                 socket.getaddrinfo, normalized_host, port, type=socket.SOCK_STREAM
             )
         except OSError as exc:
-            raise httpx.ConnectError("目标域名无法解析") from 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 addresses or any(not address.is_global or address.is_multicast
-                            or (getattr(address, 'ipv4_mapped', None) is not None
-                                and address.ipv4_mapped.is_multicast) for address in addresses):
-        raise httpx.ConnectError("目标域名解析到非公网地址")
+    if not addresses or any(not _public(address) for address in addresses):
+        if normalized_host == "huggingface.co":
+            return await _resolve_hf_public_ips(normalized_host)
+        raise PublicResolutionError("dns_non_public")
     return tuple(dict.fromkeys(str(address) for address in addresses))
 
 
@@ -83,10 +186,13 @@ class ValidatedAsyncHTTPTransport(httpx.AsyncHTTPTransport):
             network_backend=self._network_backend,
         )
 
+    async def _resolve(self, host, port):
+        return await _resolve_public_ips(host, port)
+
     async def handle_async_request(self, request):
         host = request.url.host.rstrip(".").lower()
         port = request.url.port or (443 if request.url.scheme == "https" else 80)
-        pinned_ips = await _resolve_public_ips(host, port)
+        pinned_ips = await self._resolve(host, port)
         token = _validated_ip.set(pinned_ips)
         try:
             # HTTP Core passes the original origin hostname to start_tls(), preserving
@@ -94,3 +200,10 @@ class ValidatedAsyncHTTPTransport(httpx.AsyncHTTPTransport):
             return await super().handle_async_request(request)
         finally:
             _validated_ip.reset(token)
+
+
+class _DNSHTTPTransport(ValidatedAsyncHTTPTransport):
+    async def _resolve(self, host, port):
+        if host != "cloudflare-dns.com" or port != 443:
+            raise PublicResolutionError("dns_unavailable")
+        return ("1.1.1.1",)