resolver.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149
  1. import asyncio
  2. import socket
  3. from typing import Any, Dict, List, Optional
  4. from .abc import AbstractResolver
  5. from .helpers import get_running_loop
  6. __all__ = ("ThreadedResolver", "AsyncResolver", "DefaultResolver")
  7. try:
  8. import aiodns
  9. # aiodns_default = hasattr(aiodns.DNSResolver, 'gethostbyname')
  10. except ImportError: # pragma: no cover
  11. aiodns = None
  12. aiodns_default = False
  13. class ThreadedResolver(AbstractResolver):
  14. """Use Executor for synchronous getaddrinfo() calls, which defaults to
  15. concurrent.futures.ThreadPoolExecutor.
  16. """
  17. def __init__(self, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
  18. self._loop = get_running_loop(loop)
  19. async def resolve(
  20. self, hostname: str, port: int = 0, family: int = socket.AF_INET
  21. ) -> List[Dict[str, Any]]:
  22. infos = await self._loop.getaddrinfo(
  23. hostname,
  24. port,
  25. type=socket.SOCK_STREAM,
  26. family=family,
  27. flags=socket.AI_ADDRCONFIG,
  28. )
  29. hosts = []
  30. for family, _, proto, _, address in infos:
  31. if family == socket.AF_INET6 and address[3]: # type: ignore
  32. # This is essential for link-local IPv6 addresses.
  33. # LL IPv6 is a VERY rare case. Strictly speaking, we should use
  34. # getnameinfo() unconditionally, but performance makes sense.
  35. host, _port = socket.getnameinfo(
  36. address, socket.NI_NUMERICHOST | socket.NI_NUMERICSERV
  37. )
  38. port = int(_port)
  39. else:
  40. host, port = address[:2]
  41. hosts.append(
  42. {
  43. "hostname": hostname,
  44. "host": host,
  45. "port": port,
  46. "family": family,
  47. "proto": proto,
  48. "flags": socket.AI_NUMERICHOST | socket.AI_NUMERICSERV,
  49. }
  50. )
  51. return hosts
  52. async def close(self) -> None:
  53. pass
  54. class AsyncResolver(AbstractResolver):
  55. """Use the `aiodns` package to make asynchronous DNS lookups"""
  56. def __init__(
  57. self,
  58. loop: Optional[asyncio.AbstractEventLoop] = None,
  59. *args: Any,
  60. **kwargs: Any
  61. ) -> None:
  62. if aiodns is None:
  63. raise RuntimeError("Resolver requires aiodns library")
  64. self._loop = get_running_loop(loop)
  65. self._resolver = aiodns.DNSResolver(*args, loop=loop, **kwargs)
  66. if not hasattr(self._resolver, "gethostbyname"):
  67. # aiodns 1.1 is not available, fallback to DNSResolver.query
  68. self.resolve = self._resolve_with_query # type: ignore
  69. async def resolve(
  70. self, host: str, port: int = 0, family: int = socket.AF_INET
  71. ) -> List[Dict[str, Any]]:
  72. try:
  73. resp = await self._resolver.gethostbyname(host, family)
  74. except aiodns.error.DNSError as exc:
  75. msg = exc.args[1] if len(exc.args) >= 1 else "DNS lookup failed"
  76. raise OSError(msg) from exc
  77. hosts = []
  78. for address in resp.addresses:
  79. hosts.append(
  80. {
  81. "hostname": host,
  82. "host": address,
  83. "port": port,
  84. "family": family,
  85. "proto": 0,
  86. "flags": socket.AI_NUMERICHOST | socket.AI_NUMERICSERV,
  87. }
  88. )
  89. if not hosts:
  90. raise OSError("DNS lookup failed")
  91. return hosts
  92. async def _resolve_with_query(
  93. self, host: str, port: int = 0, family: int = socket.AF_INET
  94. ) -> List[Dict[str, Any]]:
  95. if family == socket.AF_INET6:
  96. qtype = "AAAA"
  97. else:
  98. qtype = "A"
  99. try:
  100. resp = await self._resolver.query(host, qtype)
  101. except aiodns.error.DNSError as exc:
  102. msg = exc.args[1] if len(exc.args) >= 1 else "DNS lookup failed"
  103. raise OSError(msg) from exc
  104. hosts = []
  105. for rr in resp:
  106. hosts.append(
  107. {
  108. "hostname": host,
  109. "host": rr.host,
  110. "port": port,
  111. "family": family,
  112. "proto": 0,
  113. "flags": socket.AI_NUMERICHOST,
  114. }
  115. )
  116. if not hosts:
  117. raise OSError("DNS lookup failed")
  118. return hosts
  119. async def close(self) -> None:
  120. return self._resolver.cancel()
  121. DefaultResolver = AsyncResolver if aiodns_default else ThreadedResolver