_sockets.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399
  1. from __future__ import annotations
  2. import errno
  3. import socket
  4. from abc import abstractmethod
  5. from collections.abc import Callable, Collection, Mapping
  6. from contextlib import AsyncExitStack
  7. from io import IOBase
  8. from ipaddress import IPv4Address, IPv6Address
  9. from socket import AddressFamily
  10. from typing import Any, TypeAlias, TypeVar
  11. from .._core._eventloop import get_async_backend
  12. from .._core._typedattr import (
  13. TypedAttributeProvider,
  14. TypedAttributeSet,
  15. typed_attribute,
  16. )
  17. from ._streams import ByteStream, Listener, UnreliableObjectStream
  18. from ._tasks import TaskGroup
  19. IPAddressType: TypeAlias = str | IPv4Address | IPv6Address
  20. IPSockAddrType: TypeAlias = tuple[str, int]
  21. SockAddrType: TypeAlias = IPSockAddrType | str
  22. UDPPacketType: TypeAlias = tuple[bytes, IPSockAddrType]
  23. UNIXDatagramPacketType: TypeAlias = tuple[bytes, str]
  24. T_Retval = TypeVar("T_Retval")
  25. def _validate_socket(
  26. sock_or_fd: socket.socket | int,
  27. sock_type: socket.SocketKind,
  28. addr_family: socket.AddressFamily = socket.AF_UNSPEC,
  29. *,
  30. require_connected: bool = False,
  31. require_bound: bool = False,
  32. ) -> socket.socket:
  33. if isinstance(sock_or_fd, int):
  34. try:
  35. sock = socket.socket(fileno=sock_or_fd)
  36. except OSError as exc:
  37. if exc.errno == errno.ENOTSOCK:
  38. raise ValueError(
  39. "the file descriptor does not refer to a socket"
  40. ) from exc
  41. elif require_connected:
  42. raise ValueError("the socket must be connected") from exc
  43. elif require_bound:
  44. raise ValueError("the socket must be bound to a local address") from exc
  45. else:
  46. raise
  47. elif isinstance(sock_or_fd, socket.socket):
  48. sock = sock_or_fd
  49. else:
  50. raise TypeError(
  51. f"expected an int or socket, got {type(sock_or_fd).__qualname__} instead"
  52. )
  53. try:
  54. if require_connected:
  55. try:
  56. sock.getpeername()
  57. except OSError as exc:
  58. raise ValueError("the socket must be connected") from exc
  59. if require_bound:
  60. try:
  61. if sock.family in (socket.AF_INET, socket.AF_INET6):
  62. bound_addr = sock.getsockname()[1]
  63. else:
  64. bound_addr = sock.getsockname()
  65. except OSError:
  66. bound_addr = None
  67. if not bound_addr:
  68. raise ValueError("the socket must be bound to a local address")
  69. if addr_family != socket.AF_UNSPEC and sock.family != addr_family:
  70. raise ValueError(
  71. f"address family mismatch: expected {addr_family.name}, got "
  72. f"{sock.family.name}"
  73. )
  74. if sock.type != sock_type:
  75. raise ValueError(
  76. f"socket type mismatch: expected {sock_type.name}, got {sock.type.name}"
  77. )
  78. except BaseException:
  79. # Avoid ResourceWarning from the locally constructed socket object
  80. if isinstance(sock_or_fd, int):
  81. sock.detach()
  82. raise
  83. sock.setblocking(False)
  84. return sock
  85. class SocketAttribute(TypedAttributeSet):
  86. """
  87. .. attribute:: family
  88. :type: socket.AddressFamily
  89. the address family of the underlying socket
  90. .. attribute:: local_address
  91. :type: tuple[str, int] | str
  92. the local address the underlying socket is connected to
  93. .. attribute:: local_port
  94. :type: int
  95. for IP based sockets, the local port the underlying socket is bound to
  96. .. attribute:: raw_socket
  97. :type: socket.socket
  98. the underlying stdlib socket object
  99. .. attribute:: remote_address
  100. :type: tuple[str, int] | str
  101. the remote address the underlying socket is connected to
  102. .. attribute:: remote_port
  103. :type: int
  104. for IP based sockets, the remote port the underlying socket is connected to
  105. """
  106. family: AddressFamily = typed_attribute()
  107. local_address: SockAddrType = typed_attribute()
  108. local_port: int = typed_attribute()
  109. raw_socket: socket.socket = typed_attribute()
  110. remote_address: SockAddrType = typed_attribute()
  111. remote_port: int = typed_attribute()
  112. class _SocketProvider(TypedAttributeProvider):
  113. @property
  114. def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
  115. from .._core._sockets import convert_ipv6_sockaddr as convert
  116. attributes: dict[Any, Callable[[], Any]] = {
  117. SocketAttribute.family: lambda: self._raw_socket.family,
  118. SocketAttribute.local_address: lambda: convert(
  119. self._raw_socket.getsockname()
  120. ),
  121. SocketAttribute.raw_socket: lambda: self._raw_socket,
  122. }
  123. try:
  124. peername: tuple[str, int] | None = convert(self._raw_socket.getpeername())
  125. except OSError:
  126. peername = None
  127. # Provide the remote address for connected sockets
  128. if peername is not None:
  129. attributes[SocketAttribute.remote_address] = lambda: peername
  130. # Provide local and remote ports for IP based sockets
  131. if self._raw_socket.family in (AddressFamily.AF_INET, AddressFamily.AF_INET6):
  132. attributes[SocketAttribute.local_port] = lambda: (
  133. self._raw_socket.getsockname()[1]
  134. )
  135. if peername is not None:
  136. remote_port = peername[1]
  137. attributes[SocketAttribute.remote_port] = lambda: remote_port
  138. return attributes
  139. @property
  140. @abstractmethod
  141. def _raw_socket(self) -> socket.socket:
  142. pass
  143. class SocketStream(ByteStream, _SocketProvider):
  144. """
  145. Transports bytes over a socket.
  146. Supports all relevant extra attributes from :class:`~SocketAttribute`.
  147. """
  148. @classmethod
  149. async def from_socket(cls, sock_or_fd: socket.socket | int) -> SocketStream:
  150. """
  151. Wrap an existing socket object or file descriptor as a socket stream.
  152. The newly created socket wrapper takes ownership of the socket being passed in.
  153. The existing socket must already be connected.
  154. :param sock_or_fd: a socket object or file descriptor
  155. :return: a socket stream
  156. """
  157. sock = _validate_socket(sock_or_fd, socket.SOCK_STREAM, require_connected=True)
  158. return await get_async_backend().wrap_stream_socket(sock)
  159. class UNIXSocketStream(SocketStream):
  160. @classmethod
  161. async def from_socket(cls, sock_or_fd: socket.socket | int) -> UNIXSocketStream:
  162. """
  163. Wrap an existing socket object or file descriptor as a UNIX socket stream.
  164. The newly created socket wrapper takes ownership of the socket being passed in.
  165. The existing socket must already be connected.
  166. :param sock_or_fd: a socket object or file descriptor
  167. :return: a UNIX socket stream
  168. """
  169. sock = _validate_socket(
  170. sock_or_fd, socket.SOCK_STREAM, socket.AF_UNIX, require_connected=True
  171. )
  172. return await get_async_backend().wrap_unix_stream_socket(sock)
  173. @abstractmethod
  174. async def send_fds(self, message: bytes, fds: Collection[int | IOBase]) -> None:
  175. """
  176. Send file descriptors along with a message to the peer.
  177. :param message: a non-empty bytestring
  178. :param fds: a collection of files (either numeric file descriptors or open file
  179. or socket objects)
  180. """
  181. @abstractmethod
  182. async def receive_fds(self, msglen: int, maxfds: int) -> tuple[bytes, list[int]]:
  183. """
  184. Receive file descriptors along with a message from the peer.
  185. :param msglen: length of the message to expect from the peer
  186. :param maxfds: maximum number of file descriptors to expect from the peer
  187. :return: a tuple of (message, file descriptors)
  188. """
  189. class SocketListener(Listener[SocketStream], _SocketProvider):
  190. """
  191. Listens to incoming socket connections.
  192. Supports all relevant extra attributes from :class:`~SocketAttribute`.
  193. """
  194. @classmethod
  195. async def from_socket(
  196. cls,
  197. sock_or_fd: socket.socket | int,
  198. ) -> SocketListener:
  199. """
  200. Wrap an existing socket object or file descriptor as a socket listener.
  201. The newly created listener takes ownership of the socket being passed in.
  202. :param sock_or_fd: a socket object or file descriptor
  203. :return: a socket listener
  204. """
  205. sock = _validate_socket(sock_or_fd, socket.SOCK_STREAM, require_bound=True)
  206. return await get_async_backend().wrap_listener_socket(sock)
  207. @abstractmethod
  208. async def accept(self) -> SocketStream:
  209. """Accept an incoming connection."""
  210. async def serve(
  211. self,
  212. handler: Callable[[SocketStream], Any],
  213. task_group: TaskGroup | None = None,
  214. ) -> None:
  215. from .. import create_task_group
  216. async with AsyncExitStack() as stack:
  217. if task_group is None:
  218. task_group = await stack.enter_async_context(create_task_group())
  219. while True:
  220. stream = await self.accept()
  221. task_group.start_soon(handler, stream)
  222. class UDPSocket(UnreliableObjectStream[UDPPacketType], _SocketProvider):
  223. """
  224. Represents an unconnected UDP socket.
  225. Supports all relevant extra attributes from :class:`~SocketAttribute`.
  226. """
  227. @classmethod
  228. async def from_socket(cls, sock_or_fd: socket.socket | int) -> UDPSocket:
  229. """
  230. Wrap an existing socket object or file descriptor as a UDP socket.
  231. The newly created socket wrapper takes ownership of the socket being passed in.
  232. The existing socket must be bound to a local address.
  233. :param sock_or_fd: a socket object or file descriptor
  234. :return: a UDP socket
  235. """
  236. sock = _validate_socket(sock_or_fd, socket.SOCK_DGRAM, require_bound=True)
  237. return await get_async_backend().wrap_udp_socket(sock)
  238. async def sendto(self, data: bytes, host: str, port: int) -> None:
  239. """
  240. Alias for :meth:`~.UnreliableObjectSendStream.send` ((data, (host, port))).
  241. """
  242. return await self.send((data, (host, port)))
  243. class ConnectedUDPSocket(UnreliableObjectStream[bytes], _SocketProvider):
  244. """
  245. Represents an connected UDP socket.
  246. Supports all relevant extra attributes from :class:`~SocketAttribute`.
  247. """
  248. @classmethod
  249. async def from_socket(cls, sock_or_fd: socket.socket | int) -> ConnectedUDPSocket:
  250. """
  251. Wrap an existing socket object or file descriptor as a connected UDP socket.
  252. The newly created socket wrapper takes ownership of the socket being passed in.
  253. The existing socket must already be connected.
  254. :param sock_or_fd: a socket object or file descriptor
  255. :return: a connected UDP socket
  256. """
  257. sock = _validate_socket(
  258. sock_or_fd,
  259. socket.SOCK_DGRAM,
  260. require_connected=True,
  261. )
  262. return await get_async_backend().wrap_connected_udp_socket(sock)
  263. class UNIXDatagramSocket(
  264. UnreliableObjectStream[UNIXDatagramPacketType], _SocketProvider
  265. ):
  266. """
  267. Represents an unconnected Unix datagram socket.
  268. Supports all relevant extra attributes from :class:`~SocketAttribute`.
  269. """
  270. @classmethod
  271. async def from_socket(
  272. cls,
  273. sock_or_fd: socket.socket | int,
  274. ) -> UNIXDatagramSocket:
  275. """
  276. Wrap an existing socket object or file descriptor as a UNIX datagram
  277. socket.
  278. The newly created socket wrapper takes ownership of the socket being passed in.
  279. :param sock_or_fd: a socket object or file descriptor
  280. :return: a UNIX datagram socket
  281. """
  282. sock = _validate_socket(sock_or_fd, socket.SOCK_DGRAM, socket.AF_UNIX)
  283. return await get_async_backend().wrap_unix_datagram_socket(sock)
  284. async def sendto(self, data: bytes, path: str) -> None:
  285. """Alias for :meth:`~.UnreliableObjectSendStream.send` ((data, path))."""
  286. return await self.send((data, path))
  287. class ConnectedUNIXDatagramSocket(UnreliableObjectStream[bytes], _SocketProvider):
  288. """
  289. Represents a connected Unix datagram socket.
  290. Supports all relevant extra attributes from :class:`~SocketAttribute`.
  291. """
  292. @classmethod
  293. async def from_socket(
  294. cls,
  295. sock_or_fd: socket.socket | int,
  296. ) -> ConnectedUNIXDatagramSocket:
  297. """
  298. Wrap an existing socket object or file descriptor as a connected UNIX datagram
  299. socket.
  300. The newly created socket wrapper takes ownership of the socket being passed in.
  301. The existing socket must already be connected.
  302. :param sock_or_fd: a socket object or file descriptor
  303. :return: a connected UNIX datagram socket
  304. """
  305. sock = _validate_socket(
  306. sock_or_fd, socket.SOCK_DGRAM, socket.AF_UNIX, require_connected=True
  307. )
  308. return await get_async_backend().wrap_connected_unix_datagram_socket(sock)