client.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391
  1. from __future__ import annotations
  2. import os
  3. import random
  4. import warnings
  5. from collections.abc import Generator, Sequence
  6. from typing import Any
  7. from .datastructures import Headers, MultipleValuesError
  8. from .exceptions import (
  9. InvalidHandshake,
  10. InvalidHeader,
  11. InvalidHeaderValue,
  12. InvalidMessage,
  13. InvalidStatus,
  14. InvalidUpgrade,
  15. NegotiationError,
  16. )
  17. from .extensions import ClientExtensionFactory, Extension
  18. from .headers import (
  19. build_authorization_basic,
  20. build_extension,
  21. build_host,
  22. build_subprotocol,
  23. parse_connection,
  24. parse_extension,
  25. parse_subprotocol,
  26. parse_upgrade,
  27. )
  28. from .http11 import Request, Response
  29. from .imports import lazy_import
  30. from .protocol import CLIENT, CONNECTING, OPEN, Protocol, State
  31. from .typing import (
  32. ConnectionOption,
  33. ExtensionHeader,
  34. LoggerLike,
  35. Origin,
  36. Subprotocol,
  37. UpgradeProtocol,
  38. )
  39. from .uri import WebSocketURI
  40. from .utils import accept_key, generate_key
  41. __all__ = ["ClientProtocol"]
  42. class ClientProtocol(Protocol):
  43. """
  44. Sans-I/O implementation of a WebSocket client connection.
  45. Args:
  46. uri: URI of the WebSocket server, parsed
  47. with :func:`~websockets.uri.parse_uri`.
  48. origin: Value of the ``Origin`` header. This is useful when connecting
  49. to a server that validates the ``Origin`` header to defend against
  50. Cross-Site WebSocket Hijacking attacks.
  51. extensions: List of supported extensions, in order in which they
  52. should be tried.
  53. subprotocols: List of supported subprotocols, in order of decreasing
  54. preference.
  55. state: Initial state of the WebSocket connection.
  56. max_size: Maximum size of incoming messages in bytes.
  57. :obj:`None` disables the limit. You may pass a ``(max_message_size,
  58. max_fragment_size)`` tuple to set different limits for messages and
  59. fragments when you expect long messages sent in short fragments.
  60. logger: Logger for this connection;
  61. defaults to ``logging.getLogger("websockets.client")``;
  62. see the :doc:`logging guide <../../topics/logging>` for details.
  63. """
  64. def __init__(
  65. self,
  66. uri: WebSocketURI,
  67. *,
  68. origin: Origin | None = None,
  69. extensions: Sequence[ClientExtensionFactory] | None = None,
  70. subprotocols: Sequence[Subprotocol] | None = None,
  71. state: State = CONNECTING,
  72. max_size: int | None | tuple[int | None, int | None] = 2**20,
  73. logger: LoggerLike | None = None,
  74. ) -> None:
  75. super().__init__(
  76. side=CLIENT,
  77. state=state,
  78. max_size=max_size,
  79. logger=logger,
  80. )
  81. self.uri = uri
  82. self.origin = origin
  83. self.available_extensions = extensions
  84. self.available_subprotocols = subprotocols
  85. self.key = generate_key()
  86. def connect(self) -> Request:
  87. """
  88. Create a handshake request to open a connection.
  89. You must send the handshake request with :meth:`send_request`.
  90. You can modify it before sending it, for example to add HTTP headers.
  91. Returns:
  92. WebSocket handshake request event to send to the server.
  93. """
  94. headers = Headers()
  95. headers["Host"] = build_host(self.uri.host, self.uri.port, self.uri.secure)
  96. if self.uri.user_info:
  97. headers["Authorization"] = build_authorization_basic(*self.uri.user_info)
  98. if self.origin is not None:
  99. headers["Origin"] = self.origin
  100. headers["Upgrade"] = "websocket"
  101. headers["Connection"] = "Upgrade"
  102. headers["Sec-WebSocket-Key"] = self.key
  103. headers["Sec-WebSocket-Version"] = "13"
  104. if self.available_extensions is not None:
  105. headers["Sec-WebSocket-Extensions"] = build_extension(
  106. [
  107. (extension_factory.name, extension_factory.get_request_params())
  108. for extension_factory in self.available_extensions
  109. ]
  110. )
  111. if self.available_subprotocols is not None:
  112. headers["Sec-WebSocket-Protocol"] = build_subprotocol(
  113. self.available_subprotocols
  114. )
  115. return Request(self.uri.resource_name, headers)
  116. def process_response(self, response: Response) -> None:
  117. """
  118. Check a handshake response.
  119. Args:
  120. request: WebSocket handshake response received from the server.
  121. Raises:
  122. InvalidHandshake: If the handshake response is invalid.
  123. """
  124. if response.status_code != 101:
  125. raise InvalidStatus(response)
  126. headers = response.headers
  127. connection: list[ConnectionOption] = sum(
  128. [parse_connection(value) for value in headers.get_all("Connection")], []
  129. )
  130. if not any(value.lower() == "upgrade" for value in connection):
  131. raise InvalidUpgrade(
  132. "Connection", ", ".join(connection) if connection else None
  133. )
  134. upgrade: list[UpgradeProtocol] = sum(
  135. [parse_upgrade(value) for value in headers.get_all("Upgrade")], []
  136. )
  137. # For compatibility with non-strict implementations, ignore case when
  138. # checking the Upgrade header. It's supposed to be 'WebSocket'.
  139. if not (len(upgrade) == 1 and upgrade[0].lower() == "websocket"):
  140. raise InvalidUpgrade("Upgrade", ", ".join(upgrade) if upgrade else None)
  141. try:
  142. s_w_accept = headers["Sec-WebSocket-Accept"]
  143. except KeyError:
  144. raise InvalidHeader("Sec-WebSocket-Accept") from None
  145. except MultipleValuesError:
  146. raise InvalidHeader("Sec-WebSocket-Accept", "multiple values") from None
  147. if s_w_accept != accept_key(self.key):
  148. raise InvalidHeaderValue("Sec-WebSocket-Accept", s_w_accept)
  149. self.extensions = self.process_extensions(headers)
  150. self.subprotocol = self.process_subprotocol(headers)
  151. def process_extensions(self, headers: Headers) -> list[Extension]:
  152. """
  153. Handle the Sec-WebSocket-Extensions HTTP response header.
  154. Check that each extension is supported, as well as its parameters.
  155. :rfc:`6455` leaves the rules up to the specification of each
  156. extension.
  157. To provide this level of flexibility, for each extension accepted by
  158. the server, we check for a match with each extension available in the
  159. client configuration. If no match is found, an exception is raised.
  160. If several variants of the same extension are accepted by the server,
  161. it may be configured several times, which won't make sense in general.
  162. Extensions must implement their own requirements. For this purpose,
  163. the list of previously accepted extensions is provided.
  164. Other requirements, for example related to mandatory extensions or the
  165. order of extensions, may be implemented by overriding this method.
  166. Args:
  167. headers: WebSocket handshake response headers.
  168. Returns:
  169. List of accepted extensions.
  170. Raises:
  171. InvalidHandshake: To abort the handshake.
  172. """
  173. accepted_extensions: list[Extension] = []
  174. extensions = headers.get_all("Sec-WebSocket-Extensions")
  175. if extensions:
  176. if self.available_extensions is None:
  177. raise NegotiationError("no extensions supported")
  178. parsed_extensions: list[ExtensionHeader] = sum(
  179. [parse_extension(header_value) for header_value in extensions], []
  180. )
  181. for name, response_params in parsed_extensions:
  182. for extension_factory in self.available_extensions:
  183. # Skip non-matching extensions based on their name.
  184. if extension_factory.name != name:
  185. continue
  186. # Skip non-matching extensions based on their params.
  187. try:
  188. extension = extension_factory.process_response_params(
  189. response_params, accepted_extensions
  190. )
  191. except NegotiationError:
  192. continue
  193. # Add matching extension to the final list.
  194. accepted_extensions.append(extension)
  195. # Break out of the loop once we have a match.
  196. break
  197. # If we didn't break from the loop, no extension in our list
  198. # matched what the server sent. Fail the connection.
  199. else:
  200. raise NegotiationError(
  201. f"Unsupported extension: "
  202. f"name = {name}, params = {response_params}"
  203. )
  204. return accepted_extensions
  205. def process_subprotocol(self, headers: Headers) -> Subprotocol | None:
  206. """
  207. Handle the Sec-WebSocket-Protocol HTTP response header.
  208. If provided, check that it contains exactly one supported subprotocol.
  209. Args:
  210. headers: WebSocket handshake response headers.
  211. Returns:
  212. Subprotocol, if one was selected.
  213. """
  214. subprotocol: Subprotocol | None = None
  215. subprotocols = headers.get_all("Sec-WebSocket-Protocol")
  216. if subprotocols:
  217. if self.available_subprotocols is None:
  218. raise NegotiationError("no subprotocols supported")
  219. parsed_subprotocols: Sequence[Subprotocol] = sum(
  220. [parse_subprotocol(header_value) for header_value in subprotocols], []
  221. )
  222. if len(parsed_subprotocols) > 1:
  223. raise InvalidHeader(
  224. "Sec-WebSocket-Protocol",
  225. f"multiple values: {', '.join(parsed_subprotocols)}",
  226. )
  227. subprotocol = parsed_subprotocols[0]
  228. if subprotocol not in self.available_subprotocols:
  229. raise NegotiationError(f"unsupported subprotocol: {subprotocol}")
  230. return subprotocol
  231. def send_request(self, request: Request) -> None:
  232. """
  233. Send a handshake request to the server.
  234. Args:
  235. request: WebSocket handshake request event.
  236. """
  237. if self.debug:
  238. self.logger.debug("> GET %s HTTP/1.1", request.path)
  239. for key, value in request.headers.raw_items():
  240. self.logger.debug("> %s: %s", key, value)
  241. self.writes.append(request.serialize())
  242. def parse(self) -> Generator[None]:
  243. if self.state is CONNECTING:
  244. try:
  245. response = yield from Response.parse(
  246. self.reader.read_line,
  247. self.reader.read_exact,
  248. self.reader.read_to_eof,
  249. )
  250. except Exception as exc:
  251. self.handshake_exc = InvalidMessage(
  252. "did not receive a valid HTTP response"
  253. )
  254. self.handshake_exc.__cause__ = exc
  255. self.send_eof()
  256. self.parser = self.discard()
  257. next(self.parser) # start coroutine
  258. yield
  259. if self.debug:
  260. code, phrase = response.status_code, response.reason_phrase
  261. self.logger.debug("< HTTP/1.1 %d %s", code, phrase)
  262. for key, value in response.headers.raw_items():
  263. self.logger.debug("< %s: %s", key, value)
  264. if response.body:
  265. self.logger.debug("< [body] (%d bytes)", len(response.body))
  266. try:
  267. self.process_response(response)
  268. except InvalidHandshake as exc:
  269. response._exception = exc
  270. self.events.append(response)
  271. self.handshake_exc = exc
  272. self.send_eof()
  273. self.parser = self.discard()
  274. next(self.parser) # start coroutine
  275. yield
  276. assert self.state is CONNECTING
  277. self.state = OPEN
  278. self.events.append(response)
  279. yield from super().parse()
  280. class ClientConnection(ClientProtocol):
  281. def __init__(self, *args: Any, **kwargs: Any) -> None:
  282. warnings.warn( # deprecated in 11.0 - 2023-04-02
  283. "ClientConnection was renamed to ClientProtocol",
  284. DeprecationWarning,
  285. )
  286. super().__init__(*args, **kwargs)
  287. BACKOFF_INITIAL_DELAY = float(os.environ.get("WEBSOCKETS_BACKOFF_INITIAL_DELAY", "5"))
  288. BACKOFF_MIN_DELAY = float(os.environ.get("WEBSOCKETS_BACKOFF_MIN_DELAY", "3.1"))
  289. BACKOFF_MAX_DELAY = float(os.environ.get("WEBSOCKETS_BACKOFF_MAX_DELAY", "90.0"))
  290. BACKOFF_FACTOR = float(os.environ.get("WEBSOCKETS_BACKOFF_FACTOR", "1.618"))
  291. def backoff(
  292. initial_delay: float = BACKOFF_INITIAL_DELAY,
  293. min_delay: float = BACKOFF_MIN_DELAY,
  294. max_delay: float = BACKOFF_MAX_DELAY,
  295. factor: float = BACKOFF_FACTOR,
  296. ) -> Generator[float]:
  297. """
  298. Generate a series of backoff delays between reconnection attempts.
  299. Yields:
  300. How many seconds to wait before retrying to connect.
  301. """
  302. # Add a random initial delay between 0 and 5 seconds.
  303. # See 7.2.3. Recovering from Abnormal Closure in RFC 6455.
  304. yield random.random() * initial_delay
  305. delay = min_delay
  306. while delay < max_delay:
  307. yield delay
  308. delay *= factor
  309. while True:
  310. yield max_delay
  311. lazy_import(
  312. globals(),
  313. deprecated_aliases={
  314. # deprecated in 14.0 - 2024-11-09
  315. "WebSocketClientProtocol": ".legacy.client",
  316. "connect": ".legacy.client",
  317. "unix_connect": ".legacy.client",
  318. },
  319. )