server.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429
  1. import base64
  2. import binascii
  3. import collections
  4. import email.utils
  5. import http
  6. import logging
  7. from typing import Callable, Generator, List, Optional, Sequence, Tuple, Union, cast
  8. from .connection import CONNECTING, OPEN, SERVER, Connection
  9. from .datastructures import Headers, HeadersLike, MultipleValuesError
  10. from .exceptions import (
  11. InvalidHandshake,
  12. InvalidHeader,
  13. InvalidHeaderValue,
  14. InvalidOrigin,
  15. InvalidUpgrade,
  16. NegotiationError,
  17. )
  18. from .extensions.base import Extension, ServerExtensionFactory
  19. from .headers import (
  20. build_extension,
  21. parse_connection,
  22. parse_extension,
  23. parse_subprotocol,
  24. parse_upgrade,
  25. )
  26. from .http import USER_AGENT
  27. from .http11 import Request, Response
  28. from .typing import (
  29. ConnectionOption,
  30. ExtensionHeader,
  31. Origin,
  32. Subprotocol,
  33. UpgradeProtocol,
  34. )
  35. from .utils import accept_key
  36. # See #940 for why lazy_import isn't used here for backwards compatibility.
  37. from .legacy.server import * # isort:skip # noqa
  38. __all__ = ["ServerConnection"]
  39. logger = logging.getLogger(__name__)
  40. HeadersLikeOrCallable = Union[HeadersLike, Callable[[str, Headers], HeadersLike]]
  41. class ServerConnection(Connection):
  42. side = SERVER
  43. def __init__(
  44. self,
  45. origins: Optional[Sequence[Optional[Origin]]] = None,
  46. extensions: Optional[Sequence[ServerExtensionFactory]] = None,
  47. subprotocols: Optional[Sequence[Subprotocol]] = None,
  48. extra_headers: Optional[HeadersLikeOrCallable] = None,
  49. max_size: Optional[int] = 2 ** 20,
  50. ):
  51. super().__init__(side=SERVER, state=CONNECTING, max_size=max_size)
  52. self.origins = origins
  53. self.available_extensions = extensions
  54. self.available_subprotocols = subprotocols
  55. self.extra_headers = extra_headers
  56. def accept(self, request: Request) -> Response:
  57. """
  58. Create a WebSocket handshake response event to send to the client.
  59. If the connection cannot be established, the response rejects the
  60. connection, which may be unexpected.
  61. """
  62. # TODO: when changing Request to a dataclass, set the exception
  63. # attribute on the request rather than the Response, which will
  64. # be semantically more correct.
  65. try:
  66. key, extensions_header, protocol_header = self.process_request(request)
  67. except InvalidOrigin as exc:
  68. logger.debug("Invalid origin", exc_info=True)
  69. return self.reject(
  70. http.HTTPStatus.FORBIDDEN,
  71. f"Failed to open a WebSocket connection: {exc}.\n",
  72. )._replace(exception=exc)
  73. except InvalidUpgrade as exc:
  74. logger.debug("Invalid upgrade", exc_info=True)
  75. return self.reject(
  76. http.HTTPStatus.UPGRADE_REQUIRED,
  77. (
  78. f"Failed to open a WebSocket connection: {exc}.\n"
  79. f"\n"
  80. f"You cannot access a WebSocket server directly "
  81. f"with a browser. You need a WebSocket client.\n"
  82. ),
  83. headers=Headers([("Upgrade", "websocket")]),
  84. )._replace(exception=exc)
  85. except InvalidHandshake as exc:
  86. logger.debug("Invalid handshake", exc_info=True)
  87. return self.reject(
  88. http.HTTPStatus.BAD_REQUEST,
  89. f"Failed to open a WebSocket connection: {exc}.\n",
  90. )._replace(exception=exc)
  91. except Exception as exc:
  92. logger.warning("Error in opening handshake", exc_info=True)
  93. return self.reject(
  94. http.HTTPStatus.INTERNAL_SERVER_ERROR,
  95. (
  96. "Failed to open a WebSocket connection.\n"
  97. "See server log for more information.\n"
  98. ),
  99. )._replace(exception=exc)
  100. headers = Headers()
  101. headers["Upgrade"] = "websocket"
  102. headers["Connection"] = "Upgrade"
  103. headers["Sec-WebSocket-Accept"] = accept_key(key)
  104. if extensions_header is not None:
  105. headers["Sec-WebSocket-Extensions"] = extensions_header
  106. if protocol_header is not None:
  107. headers["Sec-WebSocket-Protocol"] = protocol_header
  108. extra_headers: Optional[HeadersLike]
  109. if callable(self.extra_headers):
  110. extra_headers = self.extra_headers(request.path, request.headers)
  111. else:
  112. extra_headers = self.extra_headers
  113. if extra_headers is not None:
  114. if isinstance(extra_headers, Headers):
  115. extra_headers = extra_headers.raw_items()
  116. elif isinstance(extra_headers, collections.abc.Mapping):
  117. extra_headers = extra_headers.items()
  118. for name, value in extra_headers:
  119. headers[name] = value
  120. headers.setdefault("Date", email.utils.formatdate(usegmt=True))
  121. headers.setdefault("Server", USER_AGENT)
  122. return Response(101, "Switching Protocols", headers)
  123. def process_request(
  124. self, request: Request
  125. ) -> Tuple[str, Optional[str], Optional[str]]:
  126. """
  127. Check a handshake request received from the client.
  128. This function doesn't verify that the request is an HTTP/1.1 or higher GET
  129. request and doesn't perform ``Host`` and ``Origin`` checks. These controls
  130. are usually performed earlier in the HTTP request handling code. They're
  131. the responsibility of the caller.
  132. :param request: request
  133. :returns: ``key`` which must be passed to :func:`build_response`
  134. :raises ~websockets.exceptions.InvalidHandshake: if the handshake request
  135. is invalid; then the server must return 400 Bad Request error
  136. """
  137. headers = request.headers
  138. connection: List[ConnectionOption] = sum(
  139. [parse_connection(value) for value in headers.get_all("Connection")], []
  140. )
  141. if not any(value.lower() == "upgrade" for value in connection):
  142. raise InvalidUpgrade(
  143. "Connection", ", ".join(connection) if connection else None
  144. )
  145. upgrade: List[UpgradeProtocol] = sum(
  146. [parse_upgrade(value) for value in headers.get_all("Upgrade")], []
  147. )
  148. # For compatibility with non-strict implementations, ignore case when
  149. # checking the Upgrade header. The RFC always uses "websocket", except
  150. # in section 11.2. (IANA registration) where it uses "WebSocket".
  151. if not (len(upgrade) == 1 and upgrade[0].lower() == "websocket"):
  152. raise InvalidUpgrade("Upgrade", ", ".join(upgrade) if upgrade else None)
  153. try:
  154. key = headers["Sec-WebSocket-Key"]
  155. except KeyError as exc:
  156. raise InvalidHeader("Sec-WebSocket-Key") from exc
  157. except MultipleValuesError as exc:
  158. raise InvalidHeader(
  159. "Sec-WebSocket-Key", "more than one Sec-WebSocket-Key header found"
  160. ) from exc
  161. try:
  162. raw_key = base64.b64decode(key.encode(), validate=True)
  163. except binascii.Error as exc:
  164. raise InvalidHeaderValue("Sec-WebSocket-Key", key) from exc
  165. if len(raw_key) != 16:
  166. raise InvalidHeaderValue("Sec-WebSocket-Key", key)
  167. try:
  168. version = headers["Sec-WebSocket-Version"]
  169. except KeyError as exc:
  170. raise InvalidHeader("Sec-WebSocket-Version") from exc
  171. except MultipleValuesError as exc:
  172. raise InvalidHeader(
  173. "Sec-WebSocket-Version",
  174. "more than one Sec-WebSocket-Version header found",
  175. ) from exc
  176. if version != "13":
  177. raise InvalidHeaderValue("Sec-WebSocket-Version", version)
  178. self.origin = self.process_origin(headers)
  179. extensions_header, self.extensions = self.process_extensions(headers)
  180. protocol_header = self.subprotocol = self.process_subprotocol(headers)
  181. return key, extensions_header, protocol_header
  182. def process_origin(self, headers: Headers) -> Optional[Origin]:
  183. """
  184. Handle the Origin HTTP request header.
  185. :param headers: request headers
  186. :raises ~websockets.exceptions.InvalidOrigin: if the origin isn't
  187. acceptable
  188. """
  189. # "The user agent MUST NOT include more than one Origin header field"
  190. # per https://tools.ietf.org/html/rfc6454#section-7.3.
  191. try:
  192. origin = cast(Optional[Origin], headers.get("Origin"))
  193. except MultipleValuesError as exc:
  194. raise InvalidHeader("Origin", "more than one Origin header found") from exc
  195. if self.origins is not None:
  196. if origin not in self.origins:
  197. raise InvalidOrigin(origin)
  198. return origin
  199. def process_extensions(
  200. self,
  201. headers: Headers,
  202. ) -> Tuple[Optional[str], List[Extension]]:
  203. """
  204. Handle the Sec-WebSocket-Extensions HTTP request header.
  205. Accept or reject each extension proposed in the client request.
  206. Negotiate parameters for accepted extensions.
  207. Return the Sec-WebSocket-Extensions HTTP response header and the list
  208. of accepted extensions.
  209. :rfc:`6455` leaves the rules up to the specification of each
  210. :extension.
  211. To provide this level of flexibility, for each extension proposed by
  212. the client, we check for a match with each extension available in the
  213. server configuration. If no match is found, the extension is ignored.
  214. If several variants of the same extension are proposed by the client,
  215. it may be accepted several times, which won't make sense in general.
  216. Extensions must implement their own requirements. For this purpose,
  217. the list of previously accepted extensions is provided.
  218. This process doesn't allow the server to reorder extensions. It can
  219. only select a subset of the extensions proposed by the client.
  220. Other requirements, for example related to mandatory extensions or the
  221. order of extensions, may be implemented by overriding this method.
  222. :param headers: request headers
  223. :raises ~websockets.exceptions.InvalidHandshake: to abort the
  224. handshake with an HTTP 400 error code
  225. """
  226. response_header_value: Optional[str] = None
  227. extension_headers: List[ExtensionHeader] = []
  228. accepted_extensions: List[Extension] = []
  229. header_values = headers.get_all("Sec-WebSocket-Extensions")
  230. if header_values and self.available_extensions:
  231. parsed_header_values: List[ExtensionHeader] = sum(
  232. [parse_extension(header_value) for header_value in header_values], []
  233. )
  234. for name, request_params in parsed_header_values:
  235. for ext_factory in self.available_extensions:
  236. # Skip non-matching extensions based on their name.
  237. if ext_factory.name != name:
  238. continue
  239. # Skip non-matching extensions based on their params.
  240. try:
  241. response_params, extension = ext_factory.process_request_params(
  242. request_params, accepted_extensions
  243. )
  244. except NegotiationError:
  245. continue
  246. # Add matching extension to the final list.
  247. extension_headers.append((name, response_params))
  248. accepted_extensions.append(extension)
  249. # Break out of the loop once we have a match.
  250. break
  251. # If we didn't break from the loop, no extension in our list
  252. # matched what the client sent. The extension is declined.
  253. # Serialize extension header.
  254. if extension_headers:
  255. response_header_value = build_extension(extension_headers)
  256. return response_header_value, accepted_extensions
  257. def process_subprotocol(self, headers: Headers) -> Optional[Subprotocol]:
  258. """
  259. Handle the Sec-WebSocket-Protocol HTTP request header.
  260. Return Sec-WebSocket-Protocol HTTP response header, which is the same
  261. as the selected subprotocol.
  262. :param headers: request headers
  263. :raises ~websockets.exceptions.InvalidHandshake: to abort the
  264. handshake with an HTTP 400 error code
  265. """
  266. subprotocol: Optional[Subprotocol] = None
  267. header_values = headers.get_all("Sec-WebSocket-Protocol")
  268. if header_values and self.available_subprotocols:
  269. parsed_header_values: List[Subprotocol] = sum(
  270. [parse_subprotocol(header_value) for header_value in header_values], []
  271. )
  272. subprotocol = self.select_subprotocol(
  273. parsed_header_values, self.available_subprotocols
  274. )
  275. return subprotocol
  276. def select_subprotocol(
  277. self,
  278. client_subprotocols: Sequence[Subprotocol],
  279. server_subprotocols: Sequence[Subprotocol],
  280. ) -> Optional[Subprotocol]:
  281. """
  282. Pick a subprotocol among those offered by the client.
  283. If several subprotocols are supported by the client and the server,
  284. the default implementation selects the preferred subprotocols by
  285. giving equal value to the priorities of the client and the server.
  286. If no common subprotocol is supported by the client and the server, it
  287. proceeds without a subprotocol.
  288. This is unlikely to be the most useful implementation in practice, as
  289. many servers providing a subprotocol will require that the client uses
  290. that subprotocol.
  291. :param client_subprotocols: list of subprotocols offered by the client
  292. :param server_subprotocols: list of subprotocols available on the server
  293. """
  294. subprotocols = set(client_subprotocols) & set(server_subprotocols)
  295. if not subprotocols:
  296. return None
  297. priority = lambda p: (
  298. client_subprotocols.index(p) + server_subprotocols.index(p)
  299. )
  300. return sorted(subprotocols, key=priority)[0]
  301. def reject(
  302. self,
  303. status: http.HTTPStatus,
  304. text: str,
  305. headers: Optional[Headers] = None,
  306. exception: Optional[Exception] = None,
  307. ) -> Response:
  308. """
  309. Create a HTTP response event to send to the client.
  310. A short plain text response is the best fallback when failing to
  311. establish a WebSocket connection.
  312. """
  313. body = text.encode()
  314. if headers is None:
  315. headers = Headers()
  316. headers.setdefault("Date", email.utils.formatdate(usegmt=True))
  317. headers.setdefault("Server", USER_AGENT)
  318. headers.setdefault("Content-Length", str(len(body)))
  319. headers.setdefault("Content-Type", "text/plain; charset=utf-8")
  320. headers.setdefault("Connection", "close")
  321. return Response(status.value, status.phrase, headers, body)
  322. def send_response(self, response: Response) -> None:
  323. """
  324. Send a WebSocket handshake response to the client.
  325. """
  326. if response.status_code == 101:
  327. self.set_state(OPEN)
  328. logger.debug(
  329. "%s > HTTP/1.1 %d %s",
  330. self.side,
  331. response.status_code,
  332. response.reason_phrase,
  333. )
  334. logger.debug("%s > %r", self.side, response.headers)
  335. if response.body is not None:
  336. logger.debug("%s > body (%d bytes)", self.side, len(response.body))
  337. self.writes.append(response.serialize())
  338. def parse(self) -> Generator[None, None, None]:
  339. request = yield from Request.parse(self.reader.read_line)
  340. assert self.state == CONNECTING
  341. self.events.append(request)
  342. yield from super().parse()