Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/resolver.py: 26%

1import asyncio 

2import socket 

3import sys 

4import weakref 

5from typing import Any, Optional 

6 

7from .abc import AbstractResolver, ResolveResult 

8 

9__all__ = ("ThreadedResolver", "AsyncResolver", "DefaultResolver") 

10 

11 

12try: 

13 import aiodns 

14 

15 aiodns_default = hasattr(aiodns.DNSResolver, "getaddrinfo") 

16except ImportError: 

17 aiodns = None # type: ignore[assignment] 

18 aiodns_default = False 

19 

20 

21_NUMERIC_SOCKET_FLAGS = socket.AI_NUMERICHOST | socket.AI_NUMERICSERV 

22_NAME_SOCKET_FLAGS = socket.NI_NUMERICHOST | socket.NI_NUMERICSERV 

23_AI_ADDRCONFIG = socket.AI_ADDRCONFIG 

24if hasattr(socket, "AI_MASK"): 

25 _AI_ADDRCONFIG &= socket.AI_MASK 

26_IS_WINDOWS = sys.platform == "win32" 

27 

28 

29def _is_windows_localhost(host: str) -> bool: 

30 return _IS_WINDOWS and host.rstrip(".").casefold() == "localhost" 

31 

32 

33class ThreadedResolver(AbstractResolver): 

34 """Threaded resolver. 

35 

36 Uses an Executor for synchronous getaddrinfo() calls. 

37 concurrent.futures.ThreadPoolExecutor is used by default. 

38 """ 

39 

40 def __init__(self) -> None: 

41 self._loop = asyncio.get_running_loop() 

42 

43 async def resolve( 

44 self, host: str, port: int = 0, family: socket.AddressFamily = socket.AF_INET 

45 ) -> list[ResolveResult]: 

46 try: 

47 infos = await self._loop.getaddrinfo( 

48 host, 

49 port, 

50 type=socket.SOCK_STREAM, 

51 family=family, 

52 flags=_AI_ADDRCONFIG, 

53 ) 

54 except socket.gaierror: 

55 if not _is_windows_localhost(host): 

56 raise 

57 infos = await self._loop.getaddrinfo( 

58 host, 

59 port, 

60 type=socket.SOCK_STREAM, 

61 family=family, 

62 flags=0, 

63 ) 

64 

65 hosts: list[ResolveResult] = [] 

66 for family, _, proto, _, address in infos: 

67 if family == socket.AF_INET6: 

68 if len(address) < 3: 

69 # IPv6 is not supported by Python build, 

70 # or IPv6 is not enabled in the host 

71 continue 

72 if address[3]: 

73 # This is essential for link-local IPv6 addresses. 

74 # LL IPv6 is a VERY rare case. Strictly speaking, we should use 

75 # getnameinfo() unconditionally, but performance makes sense. 

76 resolved_host, _port = await self._loop.getnameinfo( 

77 address, _NAME_SOCKET_FLAGS 

78 ) 

79 port = int(_port) 

80 else: 

81 resolved_host, port = address[:2] 

82 else: # IPv4 

83 assert family == socket.AF_INET 

84 resolved_host, port = address # type: ignore[misc] 

85 hosts.append( 

86 ResolveResult( 

87 hostname=host, 

88 host=resolved_host, 

89 port=port, 

90 family=family, 

91 proto=proto, 

92 flags=_NUMERIC_SOCKET_FLAGS, 

93 ) 

94 ) 

95 

96 return hosts 

97 

98 async def close(self) -> None: 

99 pass 

100 

101 

102class AsyncResolver(AbstractResolver): 

103 """Use the `aiodns` package to make asynchronous DNS lookups""" 

104 

105 def __init__(self, *args: Any, **kwargs: Any) -> None: 

106 if aiodns is None: 

107 raise RuntimeError("Resolver requires aiodns library") 

108 

109 self._loop = asyncio.get_running_loop() 

110 self._manager: _DNSResolverManager | None = None 

111 # If custom args are provided, create a dedicated resolver instance 

112 # This means each AsyncResolver with custom args gets its own 

113 # aiodns.DNSResolver instance 

114 if args or kwargs: 

115 self._resolver = aiodns.DNSResolver(*args, **kwargs) 

116 return 

117 # Use the shared resolver from the manager for default arguments 

118 self._manager = _DNSResolverManager() 

119 self._resolver = self._manager.get_resolver(self, self._loop) 

120 

121 async def resolve( 

122 self, host: str, port: int = 0, family: socket.AddressFamily = socket.AF_INET 

123 ) -> list[ResolveResult]: 

124 try: 

125 try: 

126 resp = await self._resolver.getaddrinfo( 

127 host, 

128 port=port, 

129 type=socket.SOCK_STREAM, 

130 family=family, 

131 flags=_AI_ADDRCONFIG, 

132 ) 

133 except aiodns.error.DNSError: 

134 if not _is_windows_localhost(host): 

135 raise 

136 resp = await self._resolver.getaddrinfo( 

137 host, 

138 port=port, 

139 type=socket.SOCK_STREAM, 

140 family=family, 

141 flags=0, 

142 ) 

143 except aiodns.error.DNSError as exc: 

144 msg = exc.args[1] if len(exc.args) >= 1 else "DNS lookup failed" 

145 raise OSError(None, msg) from exc 

146 hosts: list[ResolveResult] = [] 

147 for node in resp.nodes: 

148 address: tuple[bytes, int] | tuple[bytes, int, int, int] = node.addr 

149 if node.family == socket.AF_INET6: 

150 if len(address) > 3 and address[3]: 

151 # This is essential for link-local IPv6 addresses. 

152 # LL IPv6 is a VERY rare case. Strictly speaking, we should use 

153 # getnameinfo() unconditionally, but performance makes sense. 

154 result = await self._resolver.getnameinfo( 

155 (address[0].decode("ascii"), *address[1:]), 

156 _NAME_SOCKET_FLAGS, 

157 ) 

158 resolved_host = result.node 

159 else: 

160 resolved_host = address[0].decode("ascii") 

161 port = address[1] 

162 else: # IPv4 

163 assert node.family == socket.AF_INET 

164 resolved_host = address[0].decode("ascii") 

165 port = address[1] 

166 hosts.append( 

167 ResolveResult( 

168 hostname=host, 

169 host=resolved_host, 

170 port=port, 

171 family=node.family, 

172 proto=0, 

173 flags=_NUMERIC_SOCKET_FLAGS, 

174 ) 

175 ) 

176 

177 if not hosts: 

178 raise OSError(None, "DNS lookup failed") 

179 

180 return hosts 

181 

182 async def close(self) -> None: 

183 if self._manager: 

184 # Release the resolver from the manager if using the shared resolver 

185 self._manager.release_resolver(self, self._loop) 

186 self._manager = None # Clear reference to manager 

187 self._resolver = None # type: ignore[assignment] # Clear reference to resolver 

188 return 

189 # Otherwise cancel our dedicated resolver 

190 if self._resolver is not None: 

191 self._resolver.cancel() 

192 self._resolver = None # type: ignore[assignment] # Clear reference 

193 

194 

195class _DNSResolverManager: 

196 """Manager for aiodns.DNSResolver objects. 

197 

198 This class manages shared aiodns.DNSResolver instances 

199 with no custom arguments across different event loops. 

200 """ 

201 

202 _instance: Optional["_DNSResolverManager"] = None 

203 

204 def __new__(cls) -> "_DNSResolverManager": 

205 if cls._instance is None: 

206 cls._instance = super().__new__(cls) 

207 cls._instance._init() 

208 return cls._instance 

209 

210 def _init(self) -> None: 

211 # Use WeakKeyDictionary to allow event loops to be garbage collected 

212 self._loop_data: weakref.WeakKeyDictionary[ 

213 asyncio.AbstractEventLoop, 

214 tuple[aiodns.DNSResolver, weakref.WeakSet[AsyncResolver]], 

215 ] = weakref.WeakKeyDictionary() 

216 

217 def get_resolver( 

218 self, client: "AsyncResolver", loop: asyncio.AbstractEventLoop 

219 ) -> "aiodns.DNSResolver": 

220 """Get or create the shared aiodns.DNSResolver instance for a specific event loop. 

221 

222 Args: 

223 client: The AsyncResolver instance requesting the resolver. 

224 This is required to track resolver usage. 

225 loop: The event loop to use for the resolver. 

226 """ 

227 # Create a new resolver and client set for this loop if it doesn't exist 

228 if loop not in self._loop_data: 

229 resolver = aiodns.DNSResolver(loop=loop) 

230 client_set: weakref.WeakSet[AsyncResolver] = weakref.WeakSet() 

231 self._loop_data[loop] = (resolver, client_set) 

232 else: 

233 # Get the existing resolver and client set 

234 resolver, client_set = self._loop_data[loop] 

235 

236 # Register this client with the loop 

237 client_set.add(client) 

238 return resolver 

239 

240 def release_resolver( 

241 self, client: "AsyncResolver", loop: asyncio.AbstractEventLoop 

242 ) -> None: 

243 """Release the resolver for an AsyncResolver client when it's closed. 

244 

245 Args: 

246 client: The AsyncResolver instance to release. 

247 loop: The event loop the resolver was using. 

248 """ 

249 # Remove client from its loop's tracking 

250 current_loop_data = self._loop_data.get(loop) 

251 if current_loop_data is None: 

252 return 

253 resolver, client_set = current_loop_data 

254 client_set.discard(client) 

255 # If no more clients for this loop, cancel and remove its resolver 

256 if not client_set: 

257 if resolver is not None: 

258 resolver.cancel() 

259 del self._loop_data[loop] 

260 

261 

262_DefaultType = type[AsyncResolver | ThreadedResolver] 

263DefaultResolver: _DefaultType = AsyncResolver if aiodns_default else ThreadedResolver