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

1import asyncio 

2import socket 

3import sys 

4import weakref 

5from typing import Any 

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 info in infos: # Unpacking here would lose the tagged union for typing. 

67 if info[0] == socket.AF_INET6: 

68 address = info[4] 

69 if len(address) < 3: 

70 # IPv6 is not supported by Python build, 

71 # or IPv6 is not enabled in the host 

72 continue 

73 if address[3]: 

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

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

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

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

78 address, _NAME_SOCKET_FLAGS 

79 ) 

80 port = int(_port) 

81 else: 

82 resolved_host, port = address[:2] 

83 else: # IPv4 

84 assert info[0] == socket.AF_INET 

85 resolved_host, port = info[4] 

86 hosts.append( 

87 ResolveResult( 

88 hostname=host, 

89 host=resolved_host, 

90 port=port, 

91 family=info[0], 

92 proto=info[2], 

93 flags=_NUMERIC_SOCKET_FLAGS, 

94 ) 

95 ) 

96 

97 return hosts 

98 

99 async def close(self) -> None: 

100 pass 

101 

102 

103class AsyncResolver(AbstractResolver): 

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

105 

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

107 if aiodns is None: 

108 raise RuntimeError("Resolver requires aiodns library") 

109 

110 self._loop = asyncio.get_running_loop() 

111 self._manager: _DNSResolverManager | None = None 

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

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

114 # aiodns.DNSResolver instance 

115 if args or kwargs: 

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

117 return 

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

119 self._manager = _DNSResolverManager() 

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

121 

122 async def resolve( 

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

124 ) -> list[ResolveResult]: 

125 try: 

126 try: 

127 resp = await self._resolver.getaddrinfo( 

128 host, 

129 port=port, 

130 type=socket.SOCK_STREAM, 

131 family=family, 

132 flags=_AI_ADDRCONFIG, 

133 ) 

134 except aiodns.error.DNSError: 

135 if not _is_windows_localhost(host): 

136 raise 

137 resp = await self._resolver.getaddrinfo( 

138 host, 

139 port=port, 

140 type=socket.SOCK_STREAM, 

141 family=family, 

142 flags=0, 

143 ) 

144 except aiodns.error.DNSError as exc: 

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

146 raise OSError(None, msg) from exc 

147 hosts: list[ResolveResult] = [] 

148 for node in resp.nodes: 

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

150 if node.family == socket.AF_INET6: 

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

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

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

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

155 result = await self._resolver.getnameinfo( 

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

157 _NAME_SOCKET_FLAGS, 

158 ) 

159 resolved_host = result.node 

160 else: 

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

162 port = address[1] 

163 else: # IPv4 

164 assert node.family == socket.AF_INET 

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

166 port = address[1] 

167 hosts.append( 

168 ResolveResult( 

169 hostname=host, 

170 host=resolved_host, 

171 port=port, 

172 family=node.family, 

173 proto=0, 

174 flags=_NUMERIC_SOCKET_FLAGS, 

175 ) 

176 ) 

177 

178 if not hosts: 

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

180 

181 return hosts 

182 

183 async def close(self) -> None: 

184 if self._manager: 

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

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

187 self._manager = None # Clear reference to manager 

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

189 return 

190 # Otherwise cancel our dedicated resolver 

191 if self._resolver is not None: 

192 self._resolver.cancel() 

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

194 

195 

196class _DNSResolverManager: 

197 """Manager for aiodns.DNSResolver objects. 

198 

199 This class manages shared aiodns.DNSResolver instances 

200 with no custom arguments across different event loops. 

201 """ 

202 

203 _instance: "_DNSResolverManager | None" = None 

204 

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

206 if cls._instance is None: 

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

208 cls._instance._init() 

209 return cls._instance 

210 

211 def _init(self) -> None: 

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

213 self._loop_data: weakref.WeakKeyDictionary[ 

214 asyncio.AbstractEventLoop, 

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

216 ] = weakref.WeakKeyDictionary() 

217 

218 def get_resolver( 

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

220 ) -> "aiodns.DNSResolver": 

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

222 

223 Args: 

224 client: The AsyncResolver instance requesting the resolver. 

225 This is required to track resolver usage. 

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

227 """ 

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

229 if loop not in self._loop_data: 

230 resolver = aiodns.DNSResolver(loop=loop) 

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

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

233 else: 

234 # Get the existing resolver and client set 

235 resolver, client_set = self._loop_data[loop] 

236 

237 # Register this client with the loop 

238 client_set.add(client) 

239 return resolver 

240 

241 def release_resolver( 

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

243 ) -> None: 

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

245 

246 Args: 

247 client: The AsyncResolver instance to release. 

248 loop: The event loop the resolver was using. 

249 """ 

250 # Remove client from its loop's tracking 

251 current_loop_data = self._loop_data.get(loop) 

252 if current_loop_data is None: 

253 return 

254 resolver, client_set = current_loop_data 

255 client_set.discard(client) 

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

257 if not client_set: 

258 if resolver is not None: 

259 resolver.cancel() 

260 del self._loop_data[loop] 

261 

262 

263_DefaultType = type[AsyncResolver | ThreadedResolver] 

264DefaultResolver: _DefaultType = AsyncResolver if aiodns_default else ThreadedResolver