client.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293
  1. import collections
  2. import logging
  3. from typing import Generator, List, Optional, Sequence
  4. from .connection import CLIENT, CONNECTING, OPEN, Connection
  5. from .datastructures import Headers, HeadersLike, MultipleValuesError
  6. from .exceptions import (
  7. InvalidHandshake,
  8. InvalidHeader,
  9. InvalidHeaderValue,
  10. InvalidStatusCode,
  11. InvalidUpgrade,
  12. NegotiationError,
  13. )
  14. from .extensions.base import ClientExtensionFactory, Extension
  15. from .headers import (
  16. build_authorization_basic,
  17. build_extension,
  18. build_subprotocol,
  19. parse_connection,
  20. parse_extension,
  21. parse_subprotocol,
  22. parse_upgrade,
  23. )
  24. from .http import USER_AGENT, build_host
  25. from .http11 import Request, Response
  26. from .typing import (
  27. ConnectionOption,
  28. ExtensionHeader,
  29. Origin,
  30. Subprotocol,
  31. UpgradeProtocol,
  32. )
  33. from .uri import parse_uri
  34. from .utils import accept_key, generate_key
  35. # See #940 for why lazy_import isn't used here for backwards compatibility.
  36. from .legacy.client import * # isort:skip # noqa
  37. __all__ = ["ClientConnection"]
  38. logger = logging.getLogger(__name__)
  39. class ClientConnection(Connection):
  40. def __init__(
  41. self,
  42. uri: str,
  43. origin: Optional[Origin] = None,
  44. extensions: Optional[Sequence[ClientExtensionFactory]] = None,
  45. subprotocols: Optional[Sequence[Subprotocol]] = None,
  46. extra_headers: Optional[HeadersLike] = None,
  47. max_size: Optional[int] = 2 ** 20,
  48. ):
  49. super().__init__(side=CLIENT, state=CONNECTING, max_size=max_size)
  50. self.wsuri = parse_uri(uri)
  51. self.origin = origin
  52. self.available_extensions = extensions
  53. self.available_subprotocols = subprotocols
  54. self.extra_headers = extra_headers
  55. self.key = generate_key()
  56. def connect(self) -> Request: # noqa: F811
  57. """
  58. Create a WebSocket handshake request event to send to the server.
  59. """
  60. headers = Headers()
  61. headers["Host"] = build_host(
  62. self.wsuri.host, self.wsuri.port, self.wsuri.secure
  63. )
  64. if self.wsuri.user_info:
  65. headers["Authorization"] = build_authorization_basic(*self.wsuri.user_info)
  66. if self.origin is not None:
  67. headers["Origin"] = self.origin
  68. headers["Upgrade"] = "websocket"
  69. headers["Connection"] = "Upgrade"
  70. headers["Sec-WebSocket-Key"] = self.key
  71. headers["Sec-WebSocket-Version"] = "13"
  72. if self.available_extensions is not None:
  73. extensions_header = build_extension(
  74. [
  75. (extension_factory.name, extension_factory.get_request_params())
  76. for extension_factory in self.available_extensions
  77. ]
  78. )
  79. headers["Sec-WebSocket-Extensions"] = extensions_header
  80. if self.available_subprotocols is not None:
  81. protocol_header = build_subprotocol(self.available_subprotocols)
  82. headers["Sec-WebSocket-Protocol"] = protocol_header
  83. if self.extra_headers is not None:
  84. extra_headers = self.extra_headers
  85. if isinstance(extra_headers, Headers):
  86. extra_headers = extra_headers.raw_items()
  87. elif isinstance(extra_headers, collections.abc.Mapping):
  88. extra_headers = extra_headers.items()
  89. for name, value in extra_headers:
  90. headers[name] = value
  91. headers.setdefault("User-Agent", USER_AGENT)
  92. return Request(self.wsuri.resource_name, headers)
  93. def process_response(self, response: Response) -> None:
  94. """
  95. Check a handshake response received from the server.
  96. :param response: response
  97. :param key: comes from :func:`build_request`
  98. :raises ~websockets.exceptions.InvalidHandshake: if the handshake response
  99. is invalid
  100. """
  101. if response.status_code != 101:
  102. raise InvalidStatusCode(response.status_code)
  103. headers = response.headers
  104. connection: List[ConnectionOption] = sum(
  105. [parse_connection(value) for value in headers.get_all("Connection")], []
  106. )
  107. if not any(value.lower() == "upgrade" for value in connection):
  108. raise InvalidUpgrade(
  109. "Connection", ", ".join(connection) if connection else None
  110. )
  111. upgrade: List[UpgradeProtocol] = sum(
  112. [parse_upgrade(value) for value in headers.get_all("Upgrade")], []
  113. )
  114. # For compatibility with non-strict implementations, ignore case when
  115. # checking the Upgrade header. It's supposed to be 'WebSocket'.
  116. if not (len(upgrade) == 1 and upgrade[0].lower() == "websocket"):
  117. raise InvalidUpgrade("Upgrade", ", ".join(upgrade) if upgrade else None)
  118. try:
  119. s_w_accept = headers["Sec-WebSocket-Accept"]
  120. except KeyError as exc:
  121. raise InvalidHeader("Sec-WebSocket-Accept") from exc
  122. except MultipleValuesError as exc:
  123. raise InvalidHeader(
  124. "Sec-WebSocket-Accept",
  125. "more than one Sec-WebSocket-Accept header found",
  126. ) from exc
  127. if s_w_accept != accept_key(self.key):
  128. raise InvalidHeaderValue("Sec-WebSocket-Accept", s_w_accept)
  129. self.extensions = self.process_extensions(headers)
  130. self.subprotocol = self.process_subprotocol(headers)
  131. def process_extensions(self, headers: Headers) -> List[Extension]:
  132. """
  133. Handle the Sec-WebSocket-Extensions HTTP response header.
  134. Check that each extension is supported, as well as its parameters.
  135. Return the list of accepted extensions.
  136. Raise :exc:`~websockets.exceptions.InvalidHandshake` to abort the
  137. connection.
  138. :rfc:`6455` leaves the rules up to the specification of each
  139. extension.
  140. To provide this level of flexibility, for each extension accepted by
  141. the server, we check for a match with each extension available in the
  142. client configuration. If no match is found, an exception is raised.
  143. If several variants of the same extension are accepted by the server,
  144. it may be configured severel times, which won't make sense in general.
  145. Extensions must implement their own requirements. For this purpose,
  146. the list of previously accepted extensions is provided.
  147. Other requirements, for example related to mandatory extensions or the
  148. order of extensions, may be implemented by overriding this method.
  149. """
  150. accepted_extensions: List[Extension] = []
  151. extensions = headers.get_all("Sec-WebSocket-Extensions")
  152. if extensions:
  153. if self.available_extensions is None:
  154. raise InvalidHandshake("no extensions supported")
  155. parsed_extensions: List[ExtensionHeader] = sum(
  156. [parse_extension(header_value) for header_value in extensions], []
  157. )
  158. for name, response_params in parsed_extensions:
  159. for extension_factory in self.available_extensions:
  160. # Skip non-matching extensions based on their name.
  161. if extension_factory.name != name:
  162. continue
  163. # Skip non-matching extensions based on their params.
  164. try:
  165. extension = extension_factory.process_response_params(
  166. response_params, accepted_extensions
  167. )
  168. except NegotiationError:
  169. continue
  170. # Add matching extension to the final list.
  171. accepted_extensions.append(extension)
  172. # Break out of the loop once we have a match.
  173. break
  174. # If we didn't break from the loop, no extension in our list
  175. # matched what the server sent. Fail the connection.
  176. else:
  177. raise NegotiationError(
  178. f"Unsupported extension: "
  179. f"name = {name}, params = {response_params}"
  180. )
  181. return accepted_extensions
  182. def process_subprotocol(self, headers: Headers) -> Optional[Subprotocol]:
  183. """
  184. Handle the Sec-WebSocket-Protocol HTTP response header.
  185. Check that it contains exactly one supported subprotocol.
  186. Return the selected subprotocol.
  187. """
  188. subprotocol: Optional[Subprotocol] = None
  189. subprotocols = headers.get_all("Sec-WebSocket-Protocol")
  190. if subprotocols:
  191. if self.available_subprotocols is None:
  192. raise InvalidHandshake("no subprotocols supported")
  193. parsed_subprotocols: Sequence[Subprotocol] = sum(
  194. [parse_subprotocol(header_value) for header_value in subprotocols], []
  195. )
  196. if len(parsed_subprotocols) > 1:
  197. subprotocols_display = ", ".join(parsed_subprotocols)
  198. raise InvalidHandshake(f"multiple subprotocols: {subprotocols_display}")
  199. subprotocol = parsed_subprotocols[0]
  200. if subprotocol not in self.available_subprotocols:
  201. raise NegotiationError(f"unsupported subprotocol: {subprotocol}")
  202. return subprotocol
  203. def send_request(self, request: Request) -> None:
  204. """
  205. Send a WebSocket handshake request to the server.
  206. """
  207. logger.debug("%s > GET %s HTTP/1.1", self.side, request.path)
  208. logger.debug("%s > %r", self.side, request.headers)
  209. self.writes.append(request.serialize())
  210. def parse(self) -> Generator[None, None, None]:
  211. response = yield from Response.parse(
  212. self.reader.read_line, self.reader.read_exact, self.reader.read_to_eof
  213. )
  214. assert self.state == CONNECTING
  215. try:
  216. self.process_response(response)
  217. except InvalidHandshake as exc:
  218. response = response._replace(exception=exc)
  219. logger.debug("Invalid handshake", exc_info=True)
  220. else:
  221. self.set_state(OPEN)
  222. finally:
  223. self.events.append(response)
  224. yield from super().parse()