tls.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421
  1. from __future__ import annotations
  2. __all__ = (
  3. "TLSAttribute",
  4. "TLSConnectable",
  5. "TLSListener",
  6. "TLSStream",
  7. )
  8. import logging
  9. import re
  10. import ssl
  11. import sys
  12. from collections.abc import Callable, Mapping
  13. from dataclasses import dataclass
  14. from functools import wraps
  15. from ssl import SSLContext
  16. from typing import Any, TypeAlias, TypeVar
  17. from .. import (
  18. BrokenResourceError,
  19. EndOfStream,
  20. aclose_forcefully,
  21. get_cancelled_exc_class,
  22. to_thread,
  23. )
  24. from .._core._typedattr import TypedAttributeSet, typed_attribute
  25. from ..abc import (
  26. AnyByteStream,
  27. AnyByteStreamConnectable,
  28. ByteStream,
  29. ByteStreamConnectable,
  30. Listener,
  31. TaskGroup,
  32. )
  33. if sys.version_info >= (3, 11):
  34. from typing import TypeVarTuple, Unpack
  35. else:
  36. from typing_extensions import TypeVarTuple, Unpack
  37. if sys.version_info >= (3, 12):
  38. from typing import override
  39. else:
  40. from typing_extensions import override
  41. T_Retval = TypeVar("T_Retval")
  42. PosArgsT = TypeVarTuple("PosArgsT")
  43. _PCTRTT: TypeAlias = tuple[tuple[str, str], ...]
  44. _PCTRTTT: TypeAlias = tuple[_PCTRTT, ...]
  45. class TLSAttribute(TypedAttributeSet):
  46. """Contains Transport Layer Security related attributes."""
  47. #: the selected ALPN protocol
  48. alpn_protocol: str | None = typed_attribute()
  49. #: the channel binding for type ``tls-unique``
  50. channel_binding_tls_unique: bytes = typed_attribute()
  51. #: the selected cipher
  52. cipher: tuple[str, str, int] = typed_attribute()
  53. #: the peer certificate in dictionary form (see :meth:`ssl.SSLSocket.getpeercert`
  54. # for more information)
  55. peer_certificate: None | (dict[str, str | _PCTRTTT | _PCTRTT]) = typed_attribute()
  56. #: the peer certificate in binary form
  57. peer_certificate_binary: bytes | None = typed_attribute()
  58. #: ``True`` if this is the server side of the connection
  59. server_side: bool = typed_attribute()
  60. #: ciphers shared by the client during the TLS handshake (``None`` if this is the
  61. #: client side)
  62. shared_ciphers: list[tuple[str, str, int]] | None = typed_attribute()
  63. #: the :class:`~ssl.SSLObject` used for encryption
  64. ssl_object: ssl.SSLObject = typed_attribute()
  65. #: ``True`` if this stream does (and expects) a closing TLS handshake when the
  66. #: stream is being closed
  67. standard_compatible: bool = typed_attribute()
  68. #: the TLS protocol version (e.g. ``TLSv1.2``)
  69. tls_version: str = typed_attribute()
  70. @dataclass(eq=False)
  71. class TLSStream(ByteStream):
  72. """
  73. A stream wrapper that encrypts all sent data and decrypts received data.
  74. This class has no public initializer; use :meth:`wrap` instead.
  75. All extra attributes from :class:`~TLSAttribute` are supported.
  76. :var AnyByteStream transport_stream: the wrapped stream
  77. """
  78. transport_stream: AnyByteStream
  79. standard_compatible: bool
  80. _ssl_object: ssl.SSLObject
  81. _read_bio: ssl.MemoryBIO
  82. _write_bio: ssl.MemoryBIO
  83. @classmethod
  84. async def wrap(
  85. cls,
  86. transport_stream: AnyByteStream,
  87. *,
  88. server_side: bool | None = None,
  89. hostname: str | None = None,
  90. ssl_context: ssl.SSLContext | None = None,
  91. standard_compatible: bool = True,
  92. ) -> TLSStream:
  93. """
  94. Wrap an existing stream with Transport Layer Security.
  95. This performs a TLS handshake with the peer.
  96. :param transport_stream: a bytes-transporting stream to wrap
  97. :param server_side: ``True`` if this is the server side of the connection,
  98. ``False`` if this is the client side (if omitted, will be set to ``False``
  99. if ``hostname`` has been provided, ``False`` otherwise). Used only to create
  100. a default context when an explicit context has not been provided.
  101. :param hostname: host name of the peer (if host name checking is desired)
  102. :param ssl_context: the SSLContext object to use (if not provided, a secure
  103. default will be created)
  104. :param standard_compatible: if ``False``, skip the closing handshake when
  105. closing the connection, and don't raise an exception if the peer does the
  106. same
  107. :raises ~ssl.SSLError: if the TLS handshake fails
  108. """
  109. if server_side is None:
  110. server_side = not hostname
  111. if not ssl_context:
  112. purpose = (
  113. ssl.Purpose.CLIENT_AUTH if server_side else ssl.Purpose.SERVER_AUTH
  114. )
  115. ssl_context = ssl.create_default_context(purpose)
  116. # Re-enable detection of unexpected EOFs if it was disabled by Python
  117. if hasattr(ssl, "OP_IGNORE_UNEXPECTED_EOF"):
  118. ssl_context.options &= ~ssl.OP_IGNORE_UNEXPECTED_EOF
  119. bio_in = ssl.MemoryBIO()
  120. bio_out = ssl.MemoryBIO()
  121. # External SSLContext implementations may do blocking I/O in wrap_bio(),
  122. # but the standard library implementation won't
  123. if type(ssl_context) is ssl.SSLContext:
  124. ssl_object = ssl_context.wrap_bio(
  125. bio_in, bio_out, server_side=server_side, server_hostname=hostname
  126. )
  127. else:
  128. ssl_object = await to_thread.run_sync(
  129. ssl_context.wrap_bio,
  130. bio_in,
  131. bio_out,
  132. server_side,
  133. hostname,
  134. None,
  135. )
  136. wrapper = cls(
  137. transport_stream=transport_stream,
  138. standard_compatible=standard_compatible,
  139. _ssl_object=ssl_object,
  140. _read_bio=bio_in,
  141. _write_bio=bio_out,
  142. )
  143. await wrapper._call_sslobject_method(ssl_object.do_handshake)
  144. return wrapper
  145. async def _call_sslobject_method(
  146. self, func: Callable[[Unpack[PosArgsT]], T_Retval], *args: Unpack[PosArgsT]
  147. ) -> T_Retval:
  148. while True:
  149. try:
  150. result = func(*args)
  151. except ssl.SSLWantReadError:
  152. try:
  153. # Flush any pending writes first
  154. if self._write_bio.pending:
  155. await self.transport_stream.send(self._write_bio.read())
  156. data = await self.transport_stream.receive()
  157. except EndOfStream:
  158. self._read_bio.write_eof()
  159. except OSError as exc:
  160. self._read_bio.write_eof()
  161. self._write_bio.write_eof()
  162. raise BrokenResourceError from exc
  163. else:
  164. self._read_bio.write(data)
  165. except ssl.SSLWantWriteError:
  166. await self.transport_stream.send(self._write_bio.read())
  167. except ssl.SSLSyscallError as exc:
  168. self._read_bio.write_eof()
  169. self._write_bio.write_eof()
  170. raise BrokenResourceError from exc
  171. except ssl.SSLError as exc:
  172. self._read_bio.write_eof()
  173. self._write_bio.write_eof()
  174. if isinstance(exc, ssl.SSLEOFError) or (
  175. exc.strerror and "UNEXPECTED_EOF_WHILE_READING" in exc.strerror
  176. ):
  177. if self.standard_compatible:
  178. raise BrokenResourceError from exc
  179. else:
  180. raise EndOfStream from None
  181. raise
  182. else:
  183. # Flush any pending writes first
  184. if self._write_bio.pending:
  185. await self.transport_stream.send(self._write_bio.read())
  186. return result
  187. async def unwrap(self) -> tuple[AnyByteStream, bytes]:
  188. """
  189. Does the TLS closing handshake.
  190. :return: a tuple of (wrapped byte stream, bytes left in the read buffer)
  191. """
  192. await self._call_sslobject_method(self._ssl_object.unwrap)
  193. self._read_bio.write_eof()
  194. self._write_bio.write_eof()
  195. return self.transport_stream, self._read_bio.read()
  196. async def aclose(self) -> None:
  197. if self.standard_compatible:
  198. try:
  199. await self.unwrap()
  200. except BaseException:
  201. await aclose_forcefully(self.transport_stream)
  202. raise
  203. await self.transport_stream.aclose()
  204. async def receive(self, max_bytes: int = 65536) -> bytes:
  205. data = await self._call_sslobject_method(self._ssl_object.read, max_bytes)
  206. if not data:
  207. raise EndOfStream
  208. return data
  209. async def send(self, item: bytes) -> None:
  210. await self._call_sslobject_method(self._ssl_object.write, item)
  211. async def send_eof(self) -> None:
  212. tls_version = self.extra(TLSAttribute.tls_version)
  213. match = re.match(r"TLSv(\d+)(?:\.(\d+))?", tls_version)
  214. if match:
  215. major, minor = int(match.group(1)), int(match.group(2) or 0)
  216. if (major, minor) < (1, 3):
  217. raise NotImplementedError(
  218. f"send_eof() requires at least TLSv1.3; current "
  219. f"session uses {tls_version}"
  220. )
  221. raise NotImplementedError(
  222. "send_eof() has not yet been implemented for TLS streams"
  223. )
  224. @property
  225. def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
  226. return {
  227. **self.transport_stream.extra_attributes,
  228. TLSAttribute.alpn_protocol: self._ssl_object.selected_alpn_protocol,
  229. TLSAttribute.channel_binding_tls_unique: (
  230. self._ssl_object.get_channel_binding
  231. ),
  232. TLSAttribute.cipher: self._ssl_object.cipher,
  233. TLSAttribute.peer_certificate: lambda: self._ssl_object.getpeercert(False),
  234. TLSAttribute.peer_certificate_binary: lambda: self._ssl_object.getpeercert(
  235. True
  236. ),
  237. TLSAttribute.server_side: lambda: self._ssl_object.server_side,
  238. TLSAttribute.shared_ciphers: lambda: (
  239. self._ssl_object.shared_ciphers()
  240. if self._ssl_object.server_side
  241. else None
  242. ),
  243. TLSAttribute.standard_compatible: lambda: self.standard_compatible,
  244. TLSAttribute.ssl_object: lambda: self._ssl_object,
  245. TLSAttribute.tls_version: self._ssl_object.version,
  246. }
  247. @dataclass(eq=False)
  248. class TLSListener(Listener[TLSStream]):
  249. """
  250. A convenience listener that wraps another listener and auto-negotiates a TLS session
  251. on every accepted connection.
  252. If the TLS handshake times out or raises an exception,
  253. :meth:`handle_handshake_error` is called to do whatever post-mortem processing is
  254. deemed necessary.
  255. Supports only the :attr:`~TLSAttribute.standard_compatible` extra attribute.
  256. :param Listener listener: the listener to wrap
  257. :param ssl_context: the SSL context object
  258. :param standard_compatible: a flag passed through to :meth:`TLSStream.wrap`
  259. :param handshake_timeout: time limit for the TLS handshake
  260. (passed to :func:`~anyio.fail_after`)
  261. """
  262. listener: Listener[Any]
  263. ssl_context: ssl.SSLContext
  264. standard_compatible: bool = True
  265. handshake_timeout: float = 30
  266. @staticmethod
  267. async def handle_handshake_error(exc: BaseException, stream: AnyByteStream) -> None:
  268. """
  269. Handle an exception raised during the TLS handshake.
  270. This method does 3 things:
  271. #. Forcefully closes the original stream
  272. #. Logs the exception (unless it was a cancellation exception) using the
  273. ``anyio.streams.tls`` logger
  274. #. Reraises the exception if it was a base exception or a cancellation exception
  275. :param exc: the exception
  276. :param stream: the original stream
  277. """
  278. await aclose_forcefully(stream)
  279. # Log all except cancellation exceptions
  280. if not isinstance(exc, get_cancelled_exc_class()):
  281. # CPython (as of 3.11.5) returns incorrect `sys.exc_info()` here when using
  282. # any asyncio implementation, so we explicitly pass the exception to log
  283. # (https://github.com/python/cpython/issues/108668). Trio does not have this
  284. # issue because it works around the CPython bug.
  285. logging.getLogger(__name__).exception(
  286. "Error during TLS handshake", exc_info=exc
  287. )
  288. # Only reraise base exceptions and cancellation exceptions
  289. if not isinstance(exc, Exception) or isinstance(exc, get_cancelled_exc_class()):
  290. raise
  291. async def serve(
  292. self,
  293. handler: Callable[[TLSStream], Any],
  294. task_group: TaskGroup | None = None,
  295. ) -> None:
  296. @wraps(handler)
  297. async def handler_wrapper(stream: AnyByteStream) -> None:
  298. from .. import fail_after
  299. try:
  300. with fail_after(self.handshake_timeout):
  301. wrapped_stream = await TLSStream.wrap(
  302. stream,
  303. ssl_context=self.ssl_context,
  304. standard_compatible=self.standard_compatible,
  305. )
  306. except BaseException as exc:
  307. await self.handle_handshake_error(exc, stream)
  308. else:
  309. await handler(wrapped_stream)
  310. await self.listener.serve(handler_wrapper, task_group)
  311. async def aclose(self) -> None:
  312. await self.listener.aclose()
  313. @property
  314. def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
  315. return {
  316. TLSAttribute.standard_compatible: lambda: self.standard_compatible,
  317. }
  318. class TLSConnectable(ByteStreamConnectable):
  319. """
  320. Wraps another connectable and does TLS negotiation after a successful connection.
  321. :param connectable: the connectable to wrap
  322. :param hostname: host name of the server (if host name checking is desired)
  323. :param ssl_context: the SSLContext object to use (if not provided, a secure default
  324. will be created)
  325. :param standard_compatible: if ``False``, skip the closing handshake when closing
  326. the connection, and don't raise an exception if the server does the same
  327. """
  328. def __init__(
  329. self,
  330. connectable: AnyByteStreamConnectable,
  331. *,
  332. hostname: str | None = None,
  333. ssl_context: ssl.SSLContext | None = None,
  334. standard_compatible: bool = True,
  335. ) -> None:
  336. self.connectable = connectable
  337. self.ssl_context: SSLContext = ssl_context or ssl.create_default_context(
  338. ssl.Purpose.SERVER_AUTH
  339. )
  340. if not isinstance(self.ssl_context, ssl.SSLContext):
  341. raise TypeError(
  342. "ssl_context must be an instance of ssl.SSLContext, not "
  343. f"{type(self.ssl_context).__name__}"
  344. )
  345. self.hostname = hostname
  346. self.standard_compatible = standard_compatible
  347. @override
  348. async def connect(self) -> TLSStream:
  349. stream = await self.connectable.connect()
  350. try:
  351. return await TLSStream.wrap(
  352. stream,
  353. hostname=self.hostname,
  354. ssl_context=self.ssl_context,
  355. standard_compatible=self.standard_compatible,
  356. )
  357. except BaseException:
  358. await aclose_forcefully(stream)
  359. raise