Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/anyio/abc/_sockets.py: 45%

1from __future__ import annotations 

2 

3import errno 

4import socket 

5from abc import abstractmethod 

6from collections.abc import Callable, Collection, Mapping 

7from contextlib import AsyncExitStack 

8from io import IOBase 

9from ipaddress import IPv4Address, IPv6Address 

10from socket import AddressFamily 

11from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar 

12 

13from .._core._eventloop import get_async_backend 

14from .._core._typedattr import ( 

15 TypedAttributeProvider, 

16 TypedAttributeSet, 

17 typed_attribute, 

18) 

19from ._streams import ByteStream, Listener, UnreliableObjectStream 

20 

21if TYPE_CHECKING: 

22 from ._tasks import TaskGroup 

23 

24IPAddressType: TypeAlias = str | IPv4Address | IPv6Address 

25IPSockAddrType: TypeAlias = tuple[str, int] 

26SockAddrType: TypeAlias = IPSockAddrType | str 

27UDPPacketType: TypeAlias = tuple[bytes, IPSockAddrType] 

28UNIXDatagramPacketType: TypeAlias = tuple[bytes, str] 

29T_Retval = TypeVar("T_Retval") 

30 

31 

32def _validate_socket( 

33 sock_or_fd: socket.socket | int, 

34 sock_type: socket.SocketKind, 

35 addr_family: socket.AddressFamily = socket.AF_UNSPEC, 

36 *, 

37 require_connected: bool = False, 

38 require_bound: bool = False, 

39) -> socket.socket: 

40 if isinstance(sock_or_fd, int): 

41 try: 

42 sock = socket.socket(fileno=sock_or_fd) 

43 except OSError as exc: 

44 if exc.errno == errno.ENOTSOCK: 

45 raise ValueError( 

46 "the file descriptor does not refer to a socket" 

47 ) from exc 

48 elif require_connected: 

49 raise ValueError("the socket must be connected") from exc 

50 elif require_bound: 

51 raise ValueError("the socket must be bound to a local address") from exc 

52 else: 

53 raise 

54 elif isinstance(sock_or_fd, socket.socket): 

55 sock = sock_or_fd 

56 else: 

57 raise TypeError( 

58 f"expected an int or socket, got {type(sock_or_fd).__qualname__} instead" 

59 ) 

60 

61 try: 

62 if require_connected: 

63 try: 

64 sock.getpeername() 

65 except OSError as exc: 

66 raise ValueError("the socket must be connected") from exc 

67 

68 if require_bound: 

69 try: 

70 if sock.family in (socket.AF_INET, socket.AF_INET6): 

71 bound_addr = sock.getsockname()[1] 

72 else: 

73 bound_addr = sock.getsockname() 

74 except OSError: 

75 bound_addr = None 

76 

77 if not bound_addr: 

78 raise ValueError("the socket must be bound to a local address") 

79 

80 if addr_family != socket.AF_UNSPEC and sock.family != addr_family: 

81 raise ValueError( 

82 f"address family mismatch: expected {addr_family.name}, got " 

83 f"{sock.family.name}" 

84 ) 

85 

86 if sock.type != sock_type: 

87 raise ValueError( 

88 f"socket type mismatch: expected {sock_type.name}, got {sock.type.name}" 

89 ) 

90 except BaseException: 

91 # Avoid ResourceWarning from the locally constructed socket object 

92 if isinstance(sock_or_fd, int): 

93 sock.detach() 

94 

95 raise 

96 

97 sock.setblocking(False) 

98 return sock 

99 

100 

101class SocketAttribute(TypedAttributeSet): 

102 """ 

103 .. attribute:: family 

104 :type: socket.AddressFamily 

105 

106 the address family of the underlying socket 

107 

108 .. attribute:: local_address 

109 :type: tuple[str, int] | str 

110 

111 the local address the underlying socket is connected to 

112 

113 .. attribute:: local_port 

114 :type: int 

115 

116 for IP based sockets, the local port the underlying socket is bound to 

117 

118 .. attribute:: raw_socket 

119 :type: socket.socket 

120 

121 the underlying stdlib socket object 

122 

123 .. attribute:: remote_address 

124 :type: tuple[str, int] | str 

125 

126 the remote address the underlying socket is connected to 

127 

128 .. attribute:: remote_port 

129 :type: int 

130 

131 for IP based sockets, the remote port the underlying socket is connected to 

132 """ 

133 

134 family: AddressFamily = typed_attribute() 

135 local_address: SockAddrType = typed_attribute() 

136 local_port: int = typed_attribute() 

137 raw_socket: socket.socket = typed_attribute() 

138 remote_address: SockAddrType = typed_attribute() 

139 remote_port: int = typed_attribute() 

140 

141 

142class _SocketProvider(TypedAttributeProvider): 

143 @property 

144 def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]: 

145 from .._core._sockets import convert_ipv6_sockaddr as convert 

146 

147 attributes: dict[Any, Callable[[], Any]] = { 

148 SocketAttribute.family: lambda: self._raw_socket.family, 

149 SocketAttribute.local_address: lambda: convert( 

150 self._raw_socket.getsockname() 

151 ), 

152 SocketAttribute.raw_socket: lambda: self._raw_socket, 

153 } 

154 try: 

155 peername: tuple[str, int] | None = convert(self._raw_socket.getpeername()) 

156 except OSError: 

157 peername = None 

158 

159 # Provide the remote address for connected sockets 

160 if peername is not None: 

161 attributes[SocketAttribute.remote_address] = lambda: peername 

162 

163 # Provide local and remote ports for IP based sockets 

164 if self._raw_socket.family in (AddressFamily.AF_INET, AddressFamily.AF_INET6): 

165 attributes[SocketAttribute.local_port] = lambda: ( 

166 self._raw_socket.getsockname()[1] 

167 ) 

168 if peername is not None: 

169 remote_port = peername[1] 

170 attributes[SocketAttribute.remote_port] = lambda: remote_port 

171 

172 return attributes 

173 

174 @property 

175 @abstractmethod 

176 def _raw_socket(self) -> socket.socket: 

177 pass 

178 

179 

180class SocketStream(ByteStream, _SocketProvider): 

181 """ 

182 Transports bytes over a socket. 

183 

184 Supports all relevant extra attributes from :class:`~SocketAttribute`. 

185 """ 

186 

187 @classmethod 

188 async def from_socket(cls, sock_or_fd: socket.socket | int) -> SocketStream: 

189 """ 

190 Wrap an existing socket object or file descriptor as a socket stream. 

191 

192 The newly created socket wrapper takes ownership of the socket being passed in. 

193 The existing socket must already be connected. 

194 

195 :param sock_or_fd: a socket object or file descriptor 

196 :return: a socket stream 

197 

198 """ 

199 sock = _validate_socket(sock_or_fd, socket.SOCK_STREAM, require_connected=True) 

200 return await get_async_backend().wrap_stream_socket(sock) 

201 

202 

203class UNIXSocketStream(SocketStream): 

204 @classmethod 

205 async def from_socket(cls, sock_or_fd: socket.socket | int) -> UNIXSocketStream: 

206 """ 

207 Wrap an existing socket object or file descriptor as a UNIX socket stream. 

208 

209 The newly created socket wrapper takes ownership of the socket being passed in. 

210 The existing socket must already be connected. 

211 

212 :param sock_or_fd: a socket object or file descriptor 

213 :return: a UNIX socket stream 

214 

215 """ 

216 sock = _validate_socket( 

217 sock_or_fd, socket.SOCK_STREAM, socket.AF_UNIX, require_connected=True 

218 ) 

219 return await get_async_backend().wrap_unix_stream_socket(sock) 

220 

221 @abstractmethod 

222 async def send_fds(self, message: bytes, fds: Collection[int | IOBase]) -> None: 

223 """ 

224 Send file descriptors along with a message to the peer. 

225 

226 :param message: a non-empty bytestring 

227 :param fds: a collection of files (either numeric file descriptors or open file 

228 or socket objects) 

229 """ 

230 

231 @abstractmethod 

232 async def receive_fds(self, msglen: int, maxfds: int) -> tuple[bytes, list[int]]: 

233 """ 

234 Receive file descriptors along with a message from the peer. 

235 

236 :param msglen: length of the message to expect from the peer 

237 :param maxfds: maximum number of file descriptors to expect from the peer 

238 :return: a tuple of (message, file descriptors) 

239 """ 

240 

241 

242class SocketListener(Listener[SocketStream], _SocketProvider): 

243 """ 

244 Listens to incoming socket connections. 

245 

246 Supports all relevant extra attributes from :class:`~SocketAttribute`. 

247 """ 

248 

249 @classmethod 

250 async def from_socket( 

251 cls, 

252 sock_or_fd: socket.socket | int, 

253 ) -> SocketListener: 

254 """ 

255 Wrap an existing socket object or file descriptor as a socket listener. 

256 

257 The newly created listener takes ownership of the socket being passed in. 

258 

259 :param sock_or_fd: a socket object or file descriptor 

260 :return: a socket listener 

261 

262 """ 

263 sock = _validate_socket(sock_or_fd, socket.SOCK_STREAM, require_bound=True) 

264 return await get_async_backend().wrap_listener_socket(sock) 

265 

266 @abstractmethod 

267 async def accept(self) -> SocketStream: 

268 """Accept an incoming connection.""" 

269 

270 async def serve( 

271 self, 

272 handler: Callable[[SocketStream], Any], 

273 task_group: TaskGroup | None = None, 

274 ) -> None: 

275 from .. import create_task_group 

276 

277 async with AsyncExitStack() as stack: 

278 if task_group is None: 

279 task_group = await stack.enter_async_context(create_task_group()) 

280 

281 while True: 

282 stream = await self.accept() 

283 task_group.start_soon(handler, stream) 

284 

285 

286class UDPSocket(UnreliableObjectStream[UDPPacketType], _SocketProvider): 

287 """ 

288 Represents an unconnected UDP socket. 

289 

290 Supports all relevant extra attributes from :class:`~SocketAttribute`. 

291 """ 

292 

293 @classmethod 

294 async def from_socket(cls, sock_or_fd: socket.socket | int) -> UDPSocket: 

295 """ 

296 Wrap an existing socket object or file descriptor as a UDP socket. 

297 

298 The newly created socket wrapper takes ownership of the socket being passed in. 

299 The existing socket must be bound to a local address. 

300 

301 :param sock_or_fd: a socket object or file descriptor 

302 :return: a UDP socket 

303 

304 """ 

305 sock = _validate_socket(sock_or_fd, socket.SOCK_DGRAM, require_bound=True) 

306 return await get_async_backend().wrap_udp_socket(sock) 

307 

308 async def sendto(self, data: bytes, host: str, port: int) -> None: 

309 """ 

310 Alias for :meth:`~.UnreliableObjectSendStream.send` ((data, (host, port))). 

311 

312 """ 

313 return await self.send((data, (host, port))) 

314 

315 

316class ConnectedUDPSocket(UnreliableObjectStream[bytes], _SocketProvider): 

317 """ 

318 Represents an connected UDP socket. 

319 

320 Supports all relevant extra attributes from :class:`~SocketAttribute`. 

321 """ 

322 

323 @classmethod 

324 async def from_socket(cls, sock_or_fd: socket.socket | int) -> ConnectedUDPSocket: 

325 """ 

326 Wrap an existing socket object or file descriptor as a connected UDP socket. 

327 

328 The newly created socket wrapper takes ownership of the socket being passed in. 

329 The existing socket must already be connected. 

330 

331 :param sock_or_fd: a socket object or file descriptor 

332 :return: a connected UDP socket 

333 

334 """ 

335 sock = _validate_socket( 

336 sock_or_fd, 

337 socket.SOCK_DGRAM, 

338 require_connected=True, 

339 ) 

340 return await get_async_backend().wrap_connected_udp_socket(sock) 

341 

342 

343class UNIXDatagramSocket( 

344 UnreliableObjectStream[UNIXDatagramPacketType], _SocketProvider 

345): 

346 """ 

347 Represents an unconnected Unix datagram socket. 

348 

349 Supports all relevant extra attributes from :class:`~SocketAttribute`. 

350 """ 

351 

352 @classmethod 

353 async def from_socket( 

354 cls, 

355 sock_or_fd: socket.socket | int, 

356 ) -> UNIXDatagramSocket: 

357 """ 

358 Wrap an existing socket object or file descriptor as a UNIX datagram 

359 socket. 

360 

361 The newly created socket wrapper takes ownership of the socket being passed in. 

362 

363 :param sock_or_fd: a socket object or file descriptor 

364 :return: a UNIX datagram socket 

365 

366 """ 

367 sock = _validate_socket(sock_or_fd, socket.SOCK_DGRAM, socket.AF_UNIX) 

368 return await get_async_backend().wrap_unix_datagram_socket(sock) 

369 

370 async def sendto(self, data: bytes, path: str) -> None: 

371 """Alias for :meth:`~.UnreliableObjectSendStream.send` ((data, path)).""" 

372 return await self.send((data, path)) 

373 

374 

375class ConnectedUNIXDatagramSocket(UnreliableObjectStream[bytes], _SocketProvider): 

376 """ 

377 Represents a connected Unix datagram socket. 

378 

379 Supports all relevant extra attributes from :class:`~SocketAttribute`. 

380 """ 

381 

382 @classmethod 

383 async def from_socket( 

384 cls, 

385 sock_or_fd: socket.socket | int, 

386 ) -> ConnectedUNIXDatagramSocket: 

387 """ 

388 Wrap an existing socket object or file descriptor as a connected UNIX datagram 

389 socket. 

390 

391 The newly created socket wrapper takes ownership of the socket being passed in. 

392 The existing socket must already be connected. 

393 

394 :param sock_or_fd: a socket object or file descriptor 

395 :return: a connected UNIX datagram socket 

396 

397 """ 

398 sock = _validate_socket( 

399 sock_or_fd, socket.SOCK_DGRAM, socket.AF_UNIX, require_connected=True 

400 ) 

401 return await get_async_backend().wrap_connected_unix_datagram_socket(sock)