Spaces:
Build error
Build error
| import asyncio | |
| import socket | |
| import sys | |
| from typing import Any, Dict, List, Optional, Tuple, Type, Union | |
| from .abc import AbstractResolver, ResolveResult | |
| __all__ = ("ThreadedResolver", "AsyncResolver", "DefaultResolver") | |
| try: | |
| import aiodns | |
| aiodns_default = hasattr(aiodns.DNSResolver, "getaddrinfo") | |
| except ImportError: # pragma: no cover | |
| aiodns = None # type: ignore[assignment] | |
| aiodns_default = False | |
| _NUMERIC_SOCKET_FLAGS = socket.AI_NUMERICHOST | socket.AI_NUMERICSERV | |
| _NAME_SOCKET_FLAGS = socket.NI_NUMERICHOST | socket.NI_NUMERICSERV | |
| _SUPPORTS_SCOPE_ID = sys.version_info >= (3, 9, 0) | |
| class ThreadedResolver(AbstractResolver): | |
| """Threaded resolver. | |
| Uses an Executor for synchronous getaddrinfo() calls. | |
| concurrent.futures.ThreadPoolExecutor is used by default. | |
| """ | |
| def __init__(self, loop: Optional[asyncio.AbstractEventLoop] = None) -> None: | |
| self._loop = loop or asyncio.get_running_loop() | |
| async def resolve( | |
| self, host: str, port: int = 0, family: socket.AddressFamily = socket.AF_INET | |
| ) -> List[ResolveResult]: | |
| infos = await self._loop.getaddrinfo( | |
| host, | |
| port, | |
| type=socket.SOCK_STREAM, | |
| family=family, | |
| flags=socket.AI_ADDRCONFIG, | |
| ) | |
| hosts: List[ResolveResult] = [] | |
| for family, _, proto, _, address in infos: | |
| if family == socket.AF_INET6: | |
| if len(address) < 3: | |
| # IPv6 is not supported by Python build, | |
| # or IPv6 is not enabled in the host | |
| continue | |
| if address[3] and _SUPPORTS_SCOPE_ID: | |
| # This is essential for link-local IPv6 addresses. | |
| # LL IPv6 is a VERY rare case. Strictly speaking, we should use | |
| # getnameinfo() unconditionally, but performance makes sense. | |
| resolved_host, _port = await self._loop.getnameinfo( | |
| address, _NAME_SOCKET_FLAGS | |
| ) | |
| port = int(_port) | |
| else: | |
| resolved_host, port = address[:2] | |
| else: # IPv4 | |
| assert family == socket.AF_INET | |
| resolved_host, port = address # type: ignore[misc] | |
| hosts.append( | |
| ResolveResult( | |
| hostname=host, | |
| host=resolved_host, | |
| port=port, | |
| family=family, | |
| proto=proto, | |
| flags=_NUMERIC_SOCKET_FLAGS, | |
| ) | |
| ) | |
| return hosts | |
| async def close(self) -> None: | |
| pass | |
| class AsyncResolver(AbstractResolver): | |
| """Use the `aiodns` package to make asynchronous DNS lookups""" | |
| def __init__( | |
| self, | |
| loop: Optional[asyncio.AbstractEventLoop] = None, | |
| *args: Any, | |
| **kwargs: Any | |
| ) -> None: | |
| if aiodns is None: | |
| raise RuntimeError("Resolver requires aiodns library") | |
| self._resolver = aiodns.DNSResolver(*args, **kwargs) | |
| if not hasattr(self._resolver, "gethostbyname"): | |
| # aiodns 1.1 is not available, fallback to DNSResolver.query | |
| self.resolve = self._resolve_with_query # type: ignore | |
| async def resolve( | |
| self, host: str, port: int = 0, family: socket.AddressFamily = socket.AF_INET | |
| ) -> List[ResolveResult]: | |
| try: | |
| resp = await self._resolver.getaddrinfo( | |
| host, | |
| port=port, | |
| type=socket.SOCK_STREAM, | |
| family=family, | |
| flags=socket.AI_ADDRCONFIG, | |
| ) | |
| except aiodns.error.DNSError as exc: | |
| msg = exc.args[1] if len(exc.args) >= 1 else "DNS lookup failed" | |
| raise OSError(None, msg) from exc | |
| hosts: List[ResolveResult] = [] | |
| for node in resp.nodes: | |
| address: Union[Tuple[bytes, int], Tuple[bytes, int, int, int]] = node.addr | |
| family = node.family | |
| if family == socket.AF_INET6: | |
| if len(address) > 3 and address[3] and _SUPPORTS_SCOPE_ID: | |
| # This is essential for link-local IPv6 addresses. | |
| # LL IPv6 is a VERY rare case. Strictly speaking, we should use | |
| # getnameinfo() unconditionally, but performance makes sense. | |
| result = await self._resolver.getnameinfo( | |
| (address[0].decode("ascii"), *address[1:]), | |
| _NAME_SOCKET_FLAGS, | |
| ) | |
| resolved_host = result.node | |
| else: | |
| resolved_host = address[0].decode("ascii") | |
| port = address[1] | |
| else: # IPv4 | |
| assert family == socket.AF_INET | |
| resolved_host = address[0].decode("ascii") | |
| port = address[1] | |
| hosts.append( | |
| ResolveResult( | |
| hostname=host, | |
| host=resolved_host, | |
| port=port, | |
| family=family, | |
| proto=0, | |
| flags=_NUMERIC_SOCKET_FLAGS, | |
| ) | |
| ) | |
| if not hosts: | |
| raise OSError(None, "DNS lookup failed") | |
| return hosts | |
| async def _resolve_with_query( | |
| self, host: str, port: int = 0, family: int = socket.AF_INET | |
| ) -> List[Dict[str, Any]]: | |
| if family == socket.AF_INET6: | |
| qtype = "AAAA" | |
| else: | |
| qtype = "A" | |
| try: | |
| resp = await self._resolver.query(host, qtype) | |
| except aiodns.error.DNSError as exc: | |
| msg = exc.args[1] if len(exc.args) >= 1 else "DNS lookup failed" | |
| raise OSError(None, msg) from exc | |
| hosts = [] | |
| for rr in resp: | |
| hosts.append( | |
| { | |
| "hostname": host, | |
| "host": rr.host, | |
| "port": port, | |
| "family": family, | |
| "proto": 0, | |
| "flags": socket.AI_NUMERICHOST, | |
| } | |
| ) | |
| if not hosts: | |
| raise OSError(None, "DNS lookup failed") | |
| return hosts | |
| async def close(self) -> None: | |
| self._resolver.cancel() | |
| _DefaultType = Type[Union[AsyncResolver, ThreadedResolver]] | |
| DefaultResolver: _DefaultType = AsyncResolver if aiodns_default else ThreadedResolver | |