server.py 42 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101
  1. """
  2. :mod:`websockets.legacy.server` defines the WebSocket server APIs.
  3. """
  4. import asyncio
  5. import collections.abc
  6. import email.utils
  7. import functools
  8. import http
  9. import logging
  10. import socket
  11. import sys
  12. import warnings
  13. from types import TracebackType
  14. from typing import (
  15. Any,
  16. Awaitable,
  17. Callable,
  18. Generator,
  19. List,
  20. Optional,
  21. Sequence,
  22. Set,
  23. Tuple,
  24. Type,
  25. Union,
  26. cast,
  27. )
  28. from ..datastructures import Headers, HeadersLike, MultipleValuesError
  29. from ..exceptions import (
  30. AbortHandshake,
  31. InvalidHandshake,
  32. InvalidHeader,
  33. InvalidMessage,
  34. InvalidOrigin,
  35. InvalidUpgrade,
  36. NegotiationError,
  37. )
  38. from ..extensions.base import Extension, ServerExtensionFactory
  39. from ..extensions.permessage_deflate import enable_server_permessage_deflate
  40. from ..headers import build_extension, parse_extension, parse_subprotocol
  41. from ..http import USER_AGENT
  42. from ..typing import ExtensionHeader, Origin, Subprotocol
  43. from .handshake import build_response, check_request
  44. from .http import read_request
  45. from .protocol import WebSocketCommonProtocol
  46. __all__ = ["serve", "unix_serve", "WebSocketServerProtocol", "WebSocketServer"]
  47. logger = logging.getLogger("websockets.server")
  48. HeadersLikeOrCallable = Union[HeadersLike, Callable[[str, Headers], HeadersLike]]
  49. HTTPResponse = Tuple[http.HTTPStatus, HeadersLike, bytes]
  50. class WebSocketServerProtocol(WebSocketCommonProtocol):
  51. """
  52. :class:`~asyncio.Protocol` subclass implementing a WebSocket server.
  53. :class:`WebSocketServerProtocol`:
  54. * performs the opening handshake to establish the connection;
  55. * provides :meth:`recv` and :meth:`send` coroutines for receiving and
  56. sending messages;
  57. * deals with control frames automatically;
  58. * performs the closing handshake to terminate the connection.
  59. You may customize the opening handshake by subclassing
  60. :class:`WebSocketServer` and overriding:
  61. * :meth:`process_request` to intercept the client request before any
  62. processing and, if appropriate, to abort the WebSocket request and
  63. return a HTTP response instead;
  64. * :meth:`select_subprotocol` to select a subprotocol, if the client and
  65. the server have multiple subprotocols in common and the default logic
  66. for choosing one isn't suitable (this is rarely needed).
  67. :class:`WebSocketServerProtocol` supports asynchronous iteration::
  68. async for message in websocket:
  69. await process(message)
  70. The iterator yields incoming messages. It exits normally when the
  71. connection is closed with the close code 1000 (OK) or 1001 (going away).
  72. It raises a :exc:`~websockets.exceptions.ConnectionClosedError` exception
  73. when the connection is closed with any other code.
  74. Once the connection is open, a `Ping frame`_ is sent every
  75. ``ping_interval`` seconds. This serves as a keepalive. It helps keeping
  76. the connection open, especially in the presence of proxies with short
  77. timeouts on inactive connections. Set ``ping_interval`` to ``None`` to
  78. disable this behavior.
  79. .. _Ping frame: https://tools.ietf.org/html/rfc6455#section-5.5.2
  80. If the corresponding `Pong frame`_ isn't received within ``ping_timeout``
  81. seconds, the connection is considered unusable and is closed with
  82. code 1011. This ensures that the remote endpoint remains responsive. Set
  83. ``ping_timeout`` to ``None`` to disable this behavior.
  84. .. _Pong frame: https://tools.ietf.org/html/rfc6455#section-5.5.3
  85. The ``close_timeout`` parameter defines a maximum wait time for completing
  86. the closing handshake and terminating the TCP connection. For legacy
  87. reasons, :meth:`close` completes in at most ``4 * close_timeout`` seconds.
  88. ``close_timeout`` needs to be a parameter of the protocol because
  89. websockets usually calls :meth:`close` implicitly when the connection
  90. handler terminates.
  91. To apply a timeout to any other API, wrap it in :func:`~asyncio.wait_for`.
  92. The ``max_size`` parameter enforces the maximum size for incoming messages
  93. in bytes. The default value is 1 MiB. ``None`` disables the limit. If a
  94. message larger than the maximum size is received, :meth:`recv` will
  95. raise :exc:`~websockets.exceptions.ConnectionClosedError` and the
  96. connection will be closed with code 1009.
  97. The ``max_queue`` parameter sets the maximum length of the queue that
  98. holds incoming messages. The default value is ``32``. ``None`` disables
  99. the limit. Messages are added to an in-memory queue when they're received;
  100. then :meth:`recv` pops from that queue. In order to prevent excessive
  101. memory consumption when messages are received faster than they can be
  102. processed, the queue must be bounded. If the queue fills up, the protocol
  103. stops processing incoming data until :meth:`recv` is called. In this
  104. situation, various receive buffers (at least in :mod:`asyncio` and in the
  105. OS) will fill up, then the TCP receive window will shrink, slowing down
  106. transmission to avoid packet loss.
  107. Since Python can use up to 4 bytes of memory to represent a single
  108. character, each connection may use up to ``4 * max_size * max_queue``
  109. bytes of memory to store incoming messages. By default, this is 128 MiB.
  110. You may want to lower the limits, depending on your application's
  111. requirements.
  112. The ``read_limit`` argument sets the high-water limit of the buffer for
  113. incoming bytes. The low-water limit is half the high-water limit. The
  114. default value is 64 KiB, half of asyncio's default (based on the current
  115. implementation of :class:`~asyncio.StreamReader`).
  116. The ``write_limit`` argument sets the high-water limit of the buffer for
  117. outgoing bytes. The low-water limit is a quarter of the high-water limit.
  118. The default value is 64 KiB, equal to asyncio's default (based on the
  119. current implementation of ``FlowControlMixin``).
  120. As soon as the HTTP request and response in the opening handshake are
  121. processed:
  122. * the request path is available in the :attr:`path` attribute;
  123. * the request and response HTTP headers are available in the
  124. :attr:`request_headers` and :attr:`response_headers` attributes,
  125. which are :class:`~websockets.http.Headers` instances.
  126. If a subprotocol was negotiated, it's available in the :attr:`subprotocol`
  127. attribute.
  128. Once the connection is closed, the code is available in the
  129. :attr:`close_code` attribute and the reason in :attr:`close_reason`.
  130. All attributes must be treated as read-only.
  131. """
  132. is_client = False
  133. side = "server"
  134. def __init__(
  135. self,
  136. ws_handler: Callable[["WebSocketServerProtocol", str], Awaitable[Any]],
  137. ws_server: "WebSocketServer",
  138. *,
  139. origins: Optional[Sequence[Optional[Origin]]] = None,
  140. extensions: Optional[Sequence[ServerExtensionFactory]] = None,
  141. subprotocols: Optional[Sequence[Subprotocol]] = None,
  142. extra_headers: Optional[HeadersLikeOrCallable] = None,
  143. process_request: Optional[
  144. Callable[[str, Headers], Awaitable[Optional[HTTPResponse]]]
  145. ] = None,
  146. select_subprotocol: Optional[
  147. Callable[[Sequence[Subprotocol], Sequence[Subprotocol]], Subprotocol]
  148. ] = None,
  149. **kwargs: Any,
  150. ) -> None:
  151. # For backwards compatibility with 6.0 or earlier.
  152. if origins is not None and "" in origins:
  153. warnings.warn("use None instead of '' in origins", DeprecationWarning)
  154. origins = [None if origin == "" else origin for origin in origins]
  155. self.ws_handler = ws_handler
  156. self.ws_server = ws_server
  157. self.origins = origins
  158. self.available_extensions = extensions
  159. self.available_subprotocols = subprotocols
  160. self.extra_headers = extra_headers
  161. self._process_request = process_request
  162. self._select_subprotocol = select_subprotocol
  163. super().__init__(**kwargs)
  164. def connection_made(self, transport: asyncio.BaseTransport) -> None:
  165. """
  166. Register connection and initialize a task to handle it.
  167. """
  168. super().connection_made(transport)
  169. # Register the connection with the server before creating the handler
  170. # task. Registering at the beginning of the handler coroutine would
  171. # create a race condition between the creation of the task, which
  172. # schedules its execution, and the moment the handler starts running.
  173. self.ws_server.register(self)
  174. self.handler_task = self.loop.create_task(self.handler())
  175. async def handler(self) -> None:
  176. """
  177. Handle the lifecycle of a WebSocket connection.
  178. Since this method doesn't have a caller able to handle exceptions, it
  179. attemps to log relevant ones and guarantees that the TCP connection is
  180. closed before exiting.
  181. """
  182. try:
  183. try:
  184. path = await self.handshake(
  185. origins=self.origins,
  186. available_extensions=self.available_extensions,
  187. available_subprotocols=self.available_subprotocols,
  188. extra_headers=self.extra_headers,
  189. )
  190. # Remove this branch when dropping support for Python < 3.8
  191. # because CancelledError no longer inherits Exception.
  192. except asyncio.CancelledError: # pragma: no cover
  193. raise
  194. except ConnectionError:
  195. logger.debug("Connection error in opening handshake", exc_info=True)
  196. raise
  197. except Exception as exc:
  198. if isinstance(exc, AbortHandshake):
  199. status, headers, body = exc.status, exc.headers, exc.body
  200. elif isinstance(exc, InvalidOrigin):
  201. logger.debug("Invalid origin", exc_info=True)
  202. status, headers, body = (
  203. http.HTTPStatus.FORBIDDEN,
  204. Headers(),
  205. f"Failed to open a WebSocket connection: {exc}.\n".encode(),
  206. )
  207. elif isinstance(exc, InvalidUpgrade):
  208. logger.debug("Invalid upgrade", exc_info=True)
  209. status, headers, body = (
  210. http.HTTPStatus.UPGRADE_REQUIRED,
  211. Headers([("Upgrade", "websocket")]),
  212. (
  213. f"Failed to open a WebSocket connection: {exc}.\n"
  214. f"\n"
  215. f"You cannot access a WebSocket server directly "
  216. f"with a browser. You need a WebSocket client.\n"
  217. ).encode(),
  218. )
  219. elif isinstance(exc, InvalidHandshake):
  220. logger.debug("Invalid handshake", exc_info=True)
  221. status, headers, body = (
  222. http.HTTPStatus.BAD_REQUEST,
  223. Headers(),
  224. f"Failed to open a WebSocket connection: {exc}.\n".encode(),
  225. )
  226. else:
  227. logger.warning("Error in opening handshake", exc_info=True)
  228. status, headers, body = (
  229. http.HTTPStatus.INTERNAL_SERVER_ERROR,
  230. Headers(),
  231. (
  232. b"Failed to open a WebSocket connection.\n"
  233. b"See server log for more information.\n"
  234. ),
  235. )
  236. headers.setdefault("Date", email.utils.formatdate(usegmt=True))
  237. headers.setdefault("Server", USER_AGENT)
  238. headers.setdefault("Content-Length", str(len(body)))
  239. headers.setdefault("Content-Type", "text/plain")
  240. headers.setdefault("Connection", "close")
  241. self.write_http_response(status, headers, body)
  242. self.fail_connection()
  243. await self.wait_closed()
  244. return
  245. try:
  246. await self.ws_handler(self, path)
  247. except Exception:
  248. logger.error("Error in connection handler", exc_info=True)
  249. if not self.closed:
  250. self.fail_connection(1011)
  251. raise
  252. try:
  253. await self.close()
  254. except ConnectionError:
  255. logger.debug("Connection error in closing handshake", exc_info=True)
  256. raise
  257. except Exception:
  258. logger.warning("Error in closing handshake", exc_info=True)
  259. raise
  260. except Exception:
  261. # Last-ditch attempt to avoid leaking connections on errors.
  262. try:
  263. self.transport.close()
  264. except Exception: # pragma: no cover
  265. pass
  266. finally:
  267. # Unregister the connection with the server when the handler task
  268. # terminates. Registration is tied to the lifecycle of the handler
  269. # task because the server waits for tasks attached to registered
  270. # connections before terminating.
  271. self.ws_server.unregister(self)
  272. async def read_http_request(self) -> Tuple[str, Headers]:
  273. """
  274. Read request line and headers from the HTTP request.
  275. If the request contains a body, it may be read from ``self.reader``
  276. after this coroutine returns.
  277. :raises ~websockets.exceptions.InvalidMessage: if the HTTP message is
  278. malformed or isn't an HTTP/1.1 GET request
  279. """
  280. try:
  281. path, headers = await read_request(self.reader)
  282. except asyncio.CancelledError: # pragma: no cover
  283. raise
  284. except Exception as exc:
  285. raise InvalidMessage("did not receive a valid HTTP request") from exc
  286. logger.debug("%s < GET %s HTTP/1.1", self.side, path)
  287. logger.debug("%s < %r", self.side, headers)
  288. self.path = path
  289. self.request_headers = headers
  290. return path, headers
  291. def write_http_response(
  292. self, status: http.HTTPStatus, headers: Headers, body: Optional[bytes] = None
  293. ) -> None:
  294. """
  295. Write status line and headers to the HTTP response.
  296. This coroutine is also able to write a response body.
  297. """
  298. self.response_headers = headers
  299. logger.debug("%s > HTTP/1.1 %d %s", self.side, status.value, status.phrase)
  300. logger.debug("%s > %r", self.side, headers)
  301. # Since the status line and headers only contain ASCII characters,
  302. # we can keep this simple.
  303. response = f"HTTP/1.1 {status.value} {status.phrase}\r\n"
  304. response += str(headers)
  305. self.transport.write(response.encode())
  306. if body is not None:
  307. logger.debug("%s > body (%d bytes)", self.side, len(body))
  308. self.transport.write(body)
  309. async def process_request(
  310. self, path: str, request_headers: Headers
  311. ) -> Optional[HTTPResponse]:
  312. """
  313. Intercept the HTTP request and return an HTTP response if appropriate.
  314. If ``process_request`` returns ``None``, the WebSocket handshake
  315. continues. If it returns 3-uple containing a status code, response
  316. headers and a response body, that HTTP response is sent and the
  317. connection is closed. In that case:
  318. * The HTTP status must be a :class:`~http.HTTPStatus`.
  319. * HTTP headers must be a :class:`~websockets.http.Headers` instance, a
  320. :class:`~collections.abc.Mapping`, or an iterable of ``(name,
  321. value)`` pairs.
  322. * The HTTP response body must be :class:`bytes`. It may be empty.
  323. This coroutine may be overridden in a :class:`WebSocketServerProtocol`
  324. subclass, for example:
  325. * to return a HTTP 200 OK response on a given path; then a load
  326. balancer can use this path for a health check;
  327. * to authenticate the request and return a HTTP 401 Unauthorized or a
  328. HTTP 403 Forbidden when authentication fails.
  329. Instead of subclassing, it is possible to override this method by
  330. passing a ``process_request`` argument to the :func:`serve` function
  331. or the :class:`WebSocketServerProtocol` constructor. This is
  332. equivalent, except ``process_request`` won't have access to the
  333. protocol instance, so it can't store information for later use.
  334. ``process_request`` is expected to complete quickly. If it may run for
  335. a long time, then it should await :meth:`wait_closed` and exit if
  336. :meth:`wait_closed` completes, or else it could prevent the server
  337. from shutting down.
  338. :param path: request path, including optional query string
  339. :param request_headers: request headers
  340. """
  341. if self._process_request is not None:
  342. response = self._process_request(path, request_headers)
  343. if isinstance(response, Awaitable):
  344. return await response
  345. else:
  346. # For backwards compatibility with 7.0.
  347. warnings.warn(
  348. "declare process_request as a coroutine", DeprecationWarning
  349. )
  350. return response # type: ignore
  351. return None
  352. @staticmethod
  353. def process_origin(
  354. headers: Headers, origins: Optional[Sequence[Optional[Origin]]] = None
  355. ) -> Optional[Origin]:
  356. """
  357. Handle the Origin HTTP request header.
  358. :param headers: request headers
  359. :param origins: optional list of acceptable origins
  360. :raises ~websockets.exceptions.InvalidOrigin: if the origin isn't
  361. acceptable
  362. """
  363. # "The user agent MUST NOT include more than one Origin header field"
  364. # per https://tools.ietf.org/html/rfc6454#section-7.3.
  365. try:
  366. origin = cast(Optional[Origin], headers.get("Origin"))
  367. except MultipleValuesError as exc:
  368. raise InvalidHeader("Origin", "more than one Origin header found") from exc
  369. if origins is not None:
  370. if origin not in origins:
  371. raise InvalidOrigin(origin)
  372. return origin
  373. @staticmethod
  374. def process_extensions(
  375. headers: Headers,
  376. available_extensions: Optional[Sequence[ServerExtensionFactory]],
  377. ) -> Tuple[Optional[str], List[Extension]]:
  378. """
  379. Handle the Sec-WebSocket-Extensions HTTP request header.
  380. Accept or reject each extension proposed in the client request.
  381. Negotiate parameters for accepted extensions.
  382. Return the Sec-WebSocket-Extensions HTTP response header and the list
  383. of accepted extensions.
  384. :rfc:`6455` leaves the rules up to the specification of each
  385. :extension.
  386. To provide this level of flexibility, for each extension proposed by
  387. the client, we check for a match with each extension available in the
  388. server configuration. If no match is found, the extension is ignored.
  389. If several variants of the same extension are proposed by the client,
  390. it may be accepted several times, which won't make sense in general.
  391. Extensions must implement their own requirements. For this purpose,
  392. the list of previously accepted extensions is provided.
  393. This process doesn't allow the server to reorder extensions. It can
  394. only select a subset of the extensions proposed by the client.
  395. Other requirements, for example related to mandatory extensions or the
  396. order of extensions, may be implemented by overriding this method.
  397. :param headers: request headers
  398. :param extensions: optional list of supported extensions
  399. :raises ~websockets.exceptions.InvalidHandshake: to abort the
  400. handshake with an HTTP 400 error code
  401. """
  402. response_header_value: Optional[str] = None
  403. extension_headers: List[ExtensionHeader] = []
  404. accepted_extensions: List[Extension] = []
  405. header_values = headers.get_all("Sec-WebSocket-Extensions")
  406. if header_values and available_extensions:
  407. parsed_header_values: List[ExtensionHeader] = sum(
  408. [parse_extension(header_value) for header_value in header_values], []
  409. )
  410. for name, request_params in parsed_header_values:
  411. for ext_factory in available_extensions:
  412. # Skip non-matching extensions based on their name.
  413. if ext_factory.name != name:
  414. continue
  415. # Skip non-matching extensions based on their params.
  416. try:
  417. response_params, extension = ext_factory.process_request_params(
  418. request_params, accepted_extensions
  419. )
  420. except NegotiationError:
  421. continue
  422. # Add matching extension to the final list.
  423. extension_headers.append((name, response_params))
  424. accepted_extensions.append(extension)
  425. # Break out of the loop once we have a match.
  426. break
  427. # If we didn't break from the loop, no extension in our list
  428. # matched what the client sent. The extension is declined.
  429. # Serialize extension header.
  430. if extension_headers:
  431. response_header_value = build_extension(extension_headers)
  432. return response_header_value, accepted_extensions
  433. # Not @staticmethod because it calls self.select_subprotocol()
  434. def process_subprotocol(
  435. self, headers: Headers, available_subprotocols: Optional[Sequence[Subprotocol]]
  436. ) -> Optional[Subprotocol]:
  437. """
  438. Handle the Sec-WebSocket-Protocol HTTP request header.
  439. Return Sec-WebSocket-Protocol HTTP response header, which is the same
  440. as the selected subprotocol.
  441. :param headers: request headers
  442. :param available_subprotocols: optional list of supported subprotocols
  443. :raises ~websockets.exceptions.InvalidHandshake: to abort the
  444. handshake with an HTTP 400 error code
  445. """
  446. subprotocol: Optional[Subprotocol] = None
  447. header_values = headers.get_all("Sec-WebSocket-Protocol")
  448. if header_values and available_subprotocols:
  449. parsed_header_values: List[Subprotocol] = sum(
  450. [parse_subprotocol(header_value) for header_value in header_values], []
  451. )
  452. subprotocol = self.select_subprotocol(
  453. parsed_header_values, available_subprotocols
  454. )
  455. return subprotocol
  456. def select_subprotocol(
  457. self,
  458. client_subprotocols: Sequence[Subprotocol],
  459. server_subprotocols: Sequence[Subprotocol],
  460. ) -> Optional[Subprotocol]:
  461. """
  462. Pick a subprotocol among those offered by the client.
  463. If several subprotocols are supported by the client and the server,
  464. the default implementation selects the preferred subprotocols by
  465. giving equal value to the priorities of the client and the server.
  466. If no subprotocol is supported by the client and the server, it
  467. proceeds without a subprotocol.
  468. This is unlikely to be the most useful implementation in practice, as
  469. many servers providing a subprotocol will require that the client uses
  470. that subprotocol. Such rules can be implemented in a subclass.
  471. Instead of subclassing, it is possible to override this method by
  472. passing a ``select_subprotocol`` argument to the :func:`serve`
  473. function or the :class:`WebSocketServerProtocol` constructor.
  474. :param client_subprotocols: list of subprotocols offered by the client
  475. :param server_subprotocols: list of subprotocols available on the server
  476. """
  477. if self._select_subprotocol is not None:
  478. return self._select_subprotocol(client_subprotocols, server_subprotocols)
  479. subprotocols = set(client_subprotocols) & set(server_subprotocols)
  480. if not subprotocols:
  481. return None
  482. priority = lambda p: (
  483. client_subprotocols.index(p) + server_subprotocols.index(p)
  484. )
  485. return sorted(subprotocols, key=priority)[0]
  486. async def handshake(
  487. self,
  488. origins: Optional[Sequence[Optional[Origin]]] = None,
  489. available_extensions: Optional[Sequence[ServerExtensionFactory]] = None,
  490. available_subprotocols: Optional[Sequence[Subprotocol]] = None,
  491. extra_headers: Optional[HeadersLikeOrCallable] = None,
  492. ) -> str:
  493. """
  494. Perform the server side of the opening handshake.
  495. Return the path of the URI of the request.
  496. :param origins: list of acceptable values of the Origin HTTP header;
  497. include ``None`` if the lack of an origin is acceptable
  498. :param available_extensions: list of supported extensions in the order
  499. in which they should be used
  500. :param available_subprotocols: list of supported subprotocols in order
  501. of decreasing preference
  502. :param extra_headers: sets additional HTTP response headers when the
  503. handshake succeeds; it can be a :class:`~websockets.http.Headers`
  504. instance, a :class:`~collections.abc.Mapping`, an iterable of
  505. ``(name, value)`` pairs, or a callable taking the request path and
  506. headers in arguments and returning one of the above.
  507. :raises ~websockets.exceptions.InvalidHandshake: if the handshake
  508. fails
  509. """
  510. path, request_headers = await self.read_http_request()
  511. # Hook for customizing request handling, for example checking
  512. # authentication or treating some paths as plain HTTP endpoints.
  513. early_response_awaitable = self.process_request(path, request_headers)
  514. if isinstance(early_response_awaitable, Awaitable):
  515. early_response = await early_response_awaitable
  516. else:
  517. # For backwards compatibility with 7.0.
  518. warnings.warn("declare process_request as a coroutine", DeprecationWarning)
  519. early_response = early_response_awaitable # type: ignore
  520. # Change the response to a 503 error if the server is shutting down.
  521. if not self.ws_server.is_serving():
  522. early_response = (
  523. http.HTTPStatus.SERVICE_UNAVAILABLE,
  524. [],
  525. b"Server is shutting down.\n",
  526. )
  527. if early_response is not None:
  528. raise AbortHandshake(*early_response)
  529. key = check_request(request_headers)
  530. self.origin = self.process_origin(request_headers, origins)
  531. extensions_header, self.extensions = self.process_extensions(
  532. request_headers, available_extensions
  533. )
  534. protocol_header = self.subprotocol = self.process_subprotocol(
  535. request_headers, available_subprotocols
  536. )
  537. response_headers = Headers()
  538. build_response(response_headers, key)
  539. if extensions_header is not None:
  540. response_headers["Sec-WebSocket-Extensions"] = extensions_header
  541. if protocol_header is not None:
  542. response_headers["Sec-WebSocket-Protocol"] = protocol_header
  543. if callable(extra_headers):
  544. extra_headers = extra_headers(path, self.request_headers)
  545. if extra_headers is not None:
  546. if isinstance(extra_headers, Headers):
  547. extra_headers = extra_headers.raw_items()
  548. elif isinstance(extra_headers, collections.abc.Mapping):
  549. extra_headers = extra_headers.items()
  550. for name, value in extra_headers:
  551. response_headers[name] = value
  552. response_headers.setdefault("Date", email.utils.formatdate(usegmt=True))
  553. response_headers.setdefault("Server", USER_AGENT)
  554. self.write_http_response(http.HTTPStatus.SWITCHING_PROTOCOLS, response_headers)
  555. self.connection_open()
  556. return path
  557. class WebSocketServer:
  558. """
  559. WebSocket server returned by :func:`serve`.
  560. This class provides the same interface as
  561. :class:`~asyncio.AbstractServer`, namely the
  562. :meth:`~asyncio.AbstractServer.close` and
  563. :meth:`~asyncio.AbstractServer.wait_closed` methods.
  564. It keeps track of WebSocket connections in order to close them properly
  565. when shutting down.
  566. Instances of this class store a reference to the :class:`~asyncio.Server`
  567. object returned by :meth:`~asyncio.loop.create_server` rather than inherit
  568. from :class:`~asyncio.Server` in part because
  569. :meth:`~asyncio.loop.create_server` doesn't support passing a custom
  570. :class:`~asyncio.Server` class.
  571. """
  572. def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
  573. # Store a reference to loop to avoid relying on self.server._loop.
  574. self.loop = loop
  575. # Keep track of active connections.
  576. self.websockets: Set[WebSocketServerProtocol] = set()
  577. # Task responsible for closing the server and terminating connections.
  578. self.close_task: Optional[asyncio.Task[None]] = None
  579. # Completed when the server is closed and connections are terminated.
  580. self.closed_waiter: asyncio.Future[None] = loop.create_future()
  581. def wrap(self, server: asyncio.AbstractServer) -> None:
  582. """
  583. Attach to a given :class:`~asyncio.Server`.
  584. Since :meth:`~asyncio.loop.create_server` doesn't support injecting a
  585. custom ``Server`` class, the easiest solution that doesn't rely on
  586. private :mod:`asyncio` APIs is to:
  587. - instantiate a :class:`WebSocketServer`
  588. - give the protocol factory a reference to that instance
  589. - call :meth:`~asyncio.loop.create_server` with the factory
  590. - attach the resulting :class:`~asyncio.Server` with this method
  591. """
  592. self.server = server
  593. def register(self, protocol: WebSocketServerProtocol) -> None:
  594. """
  595. Register a connection with this server.
  596. """
  597. self.websockets.add(protocol)
  598. def unregister(self, protocol: WebSocketServerProtocol) -> None:
  599. """
  600. Unregister a connection with this server.
  601. """
  602. self.websockets.remove(protocol)
  603. def is_serving(self) -> bool:
  604. """
  605. Tell whether the server is accepting new connections or shutting down.
  606. """
  607. try:
  608. # Python ≥ 3.7
  609. return self.server.is_serving()
  610. except AttributeError: # pragma: no cover
  611. # Python < 3.7
  612. return self.server.sockets is not None
  613. def close(self) -> None:
  614. """
  615. Close the server.
  616. This method:
  617. * closes the underlying :class:`~asyncio.Server`;
  618. * rejects new WebSocket connections with an HTTP 503 (service
  619. unavailable) error; this happens when the server accepted the TCP
  620. connection but didn't complete the WebSocket opening handshake prior
  621. to closing;
  622. * closes open WebSocket connections with close code 1001 (going away).
  623. :meth:`close` is idempotent.
  624. """
  625. if self.close_task is None:
  626. self.close_task = self.loop.create_task(self._close())
  627. async def _close(self) -> None:
  628. """
  629. Implementation of :meth:`close`.
  630. This calls :meth:`~asyncio.Server.close` on the underlying
  631. :class:`~asyncio.Server` object to stop accepting new connections and
  632. then closes open connections with close code 1001.
  633. """
  634. # Stop accepting new connections.
  635. self.server.close()
  636. # Wait until self.server.close() completes.
  637. await self.server.wait_closed()
  638. # Wait until all accepted connections reach connection_made() and call
  639. # register(). See https://bugs.python.org/issue34852 for details.
  640. await asyncio.sleep(
  641. 0, loop=self.loop if sys.version_info[:2] < (3, 8) else None
  642. )
  643. # Close OPEN connections with status code 1001. Since the server was
  644. # closed, handshake() closes OPENING conections with a HTTP 503 error.
  645. # Wait until all connections are closed.
  646. # asyncio.wait doesn't accept an empty first argument
  647. if self.websockets:
  648. await asyncio.wait(
  649. [
  650. asyncio.ensure_future(websocket.close(1001))
  651. for websocket in self.websockets
  652. ],
  653. loop=self.loop if sys.version_info[:2] < (3, 8) else None,
  654. )
  655. # Wait until all connection handlers are complete.
  656. # asyncio.wait doesn't accept an empty first argument.
  657. if self.websockets:
  658. await asyncio.wait(
  659. [websocket.handler_task for websocket in self.websockets],
  660. loop=self.loop if sys.version_info[:2] < (3, 8) else None,
  661. )
  662. # Tell wait_closed() to return.
  663. self.closed_waiter.set_result(None)
  664. async def wait_closed(self) -> None:
  665. """
  666. Wait until the server is closed.
  667. When :meth:`wait_closed` returns, all TCP connections are closed and
  668. all connection handlers have returned.
  669. """
  670. await asyncio.shield(self.closed_waiter)
  671. @property
  672. def sockets(self) -> Optional[List[socket.socket]]:
  673. """
  674. List of :class:`~socket.socket` objects the server is listening to.
  675. ``None`` if the server is closed.
  676. """
  677. return self.server.sockets
  678. class Serve:
  679. """
  680. Create, start, and return a WebSocket server on ``host`` and ``port``.
  681. Whenever a client connects, the server accepts the connection, creates a
  682. :class:`WebSocketServerProtocol`, performs the opening handshake, and
  683. delegates to the connection handler defined by ``ws_handler``. Once the
  684. handler completes, either normally or with an exception, the server
  685. performs the closing handshake and closes the connection.
  686. Awaiting :func:`serve` yields a :class:`WebSocketServer`. This instance
  687. provides :meth:`~WebSocketServer.close` and
  688. :meth:`~WebSocketServer.wait_closed` methods for terminating the server
  689. and cleaning up its resources.
  690. When a server is closed with :meth:`~WebSocketServer.close`, it closes all
  691. connections with close code 1001 (going away). Connections handlers, which
  692. are running the ``ws_handler`` coroutine, will receive a
  693. :exc:`~websockets.exceptions.ConnectionClosedOK` exception on their
  694. current or next interaction with the WebSocket connection.
  695. :func:`serve` can also be used as an asynchronous context manager::
  696. stop = asyncio.Future() # set this future to exit the server
  697. async with serve(...):
  698. await stop
  699. In this case, the server is shut down when exiting the context.
  700. :func:`serve` is a wrapper around the event loop's
  701. :meth:`~asyncio.loop.create_server` method. It creates and starts a
  702. :class:`asyncio.Server` with :meth:`~asyncio.loop.create_server`. Then it
  703. wraps the :class:`asyncio.Server` in a :class:`WebSocketServer` and
  704. returns the :class:`WebSocketServer`.
  705. ``ws_handler`` is the WebSocket handler. It must be a coroutine accepting
  706. two arguments: the WebSocket connection, which is an instance of
  707. :class:`WebSocketServerProtocol`, and the path of the request.
  708. The ``host`` and ``port`` arguments, as well as unrecognized keyword
  709. arguments, are passed to :meth:`~asyncio.loop.create_server`.
  710. For example, you can set the ``ssl`` keyword argument to a
  711. :class:`~ssl.SSLContext` to enable TLS.
  712. ``create_protocol`` defaults to :class:`WebSocketServerProtocol`. It may
  713. be replaced by a wrapper or a subclass to customize the protocol that
  714. manages the connection.
  715. The behavior of ``ping_interval``, ``ping_timeout``, ``close_timeout``,
  716. ``max_size``, ``max_queue``, ``read_limit``, and ``write_limit`` is
  717. described in :class:`WebSocketServerProtocol`.
  718. :func:`serve` also accepts the following optional arguments:
  719. * ``compression`` is a shortcut to configure compression extensions;
  720. by default it enables the "permessage-deflate" extension; set it to
  721. ``None`` to disable compression.
  722. * ``origins`` defines acceptable Origin HTTP headers; include ``None`` in
  723. the list if the lack of an origin is acceptable.
  724. * ``extensions`` is a list of supported extensions in order of
  725. decreasing preference.
  726. * ``subprotocols`` is a list of supported subprotocols in order of
  727. decreasing preference.
  728. * ``extra_headers`` sets additional HTTP response headers when the
  729. handshake succeeds; it can be a :class:`~websockets.http.Headers`
  730. instance, a :class:`~collections.abc.Mapping`, an iterable of ``(name,
  731. value)`` pairs, or a callable taking the request path and headers in
  732. arguments and returning one of the above.
  733. * ``process_request`` allows intercepting the HTTP request; it must be a
  734. coroutine taking the request path and headers in argument; see
  735. :meth:`~WebSocketServerProtocol.process_request` for details.
  736. * ``select_subprotocol`` allows customizing the logic for selecting a
  737. subprotocol; it must be a callable taking the subprotocols offered by
  738. the client and available on the server in argument; see
  739. :meth:`~WebSocketServerProtocol.select_subprotocol` for details.
  740. Since there's no useful way to propagate exceptions triggered in handlers,
  741. they're sent to the ``"websockets.server"`` logger instead.
  742. Debugging is much easier if you configure logging to print them::
  743. import logging
  744. logger = logging.getLogger("websockets.server")
  745. logger.setLevel(logging.ERROR)
  746. logger.addHandler(logging.StreamHandler())
  747. """
  748. def __init__(
  749. self,
  750. ws_handler: Callable[[WebSocketServerProtocol, str], Awaitable[Any]],
  751. host: Optional[Union[str, Sequence[str]]] = None,
  752. port: Optional[int] = None,
  753. *,
  754. create_protocol: Optional[Callable[[Any], WebSocketServerProtocol]] = None,
  755. ping_interval: Optional[float] = 20,
  756. ping_timeout: Optional[float] = 20,
  757. close_timeout: Optional[float] = None,
  758. max_size: Optional[int] = 2 ** 20,
  759. max_queue: Optional[int] = 2 ** 5,
  760. read_limit: int = 2 ** 16,
  761. write_limit: int = 2 ** 16,
  762. loop: Optional[asyncio.AbstractEventLoop] = None,
  763. compression: Optional[str] = "deflate",
  764. origins: Optional[Sequence[Optional[Origin]]] = None,
  765. extensions: Optional[Sequence[ServerExtensionFactory]] = None,
  766. subprotocols: Optional[Sequence[Subprotocol]] = None,
  767. extra_headers: Optional[HeadersLikeOrCallable] = None,
  768. process_request: Optional[
  769. Callable[[str, Headers], Awaitable[Optional[HTTPResponse]]]
  770. ] = None,
  771. select_subprotocol: Optional[
  772. Callable[[Sequence[Subprotocol], Sequence[Subprotocol]], Subprotocol]
  773. ] = None,
  774. **kwargs: Any,
  775. ) -> None:
  776. # Backwards compatibility: close_timeout used to be called timeout.
  777. timeout: Optional[float] = kwargs.pop("timeout", None)
  778. if timeout is None:
  779. timeout = 10
  780. else:
  781. warnings.warn("rename timeout to close_timeout", DeprecationWarning)
  782. # If both are specified, timeout is ignored.
  783. if close_timeout is None:
  784. close_timeout = timeout
  785. # Backwards compatibility: create_protocol used to be called klass.
  786. klass: Optional[Type[WebSocketServerProtocol]] = kwargs.pop("klass", None)
  787. if klass is None:
  788. klass = WebSocketServerProtocol
  789. else:
  790. warnings.warn("rename klass to create_protocol", DeprecationWarning)
  791. # If both are specified, klass is ignored.
  792. if create_protocol is None:
  793. create_protocol = klass
  794. # Backwards compatibility: recv() used to return None on closed connections
  795. legacy_recv: bool = kwargs.pop("legacy_recv", False)
  796. if loop is None:
  797. loop = asyncio.get_event_loop()
  798. ws_server = WebSocketServer(loop)
  799. secure = kwargs.get("ssl") is not None
  800. if compression == "deflate":
  801. extensions = enable_server_permessage_deflate(extensions)
  802. elif compression is not None:
  803. raise ValueError(f"unsupported compression: {compression}")
  804. factory = functools.partial(
  805. create_protocol,
  806. ws_handler,
  807. ws_server,
  808. host=host,
  809. port=port,
  810. secure=secure,
  811. ping_interval=ping_interval,
  812. ping_timeout=ping_timeout,
  813. close_timeout=close_timeout,
  814. max_size=max_size,
  815. max_queue=max_queue,
  816. read_limit=read_limit,
  817. write_limit=write_limit,
  818. loop=loop,
  819. legacy_recv=legacy_recv,
  820. origins=origins,
  821. extensions=extensions,
  822. subprotocols=subprotocols,
  823. extra_headers=extra_headers,
  824. process_request=process_request,
  825. select_subprotocol=select_subprotocol,
  826. )
  827. if kwargs.pop("unix", False):
  828. path: Optional[str] = kwargs.pop("path", None)
  829. # unix_serve(path) must not specify host and port parameters.
  830. assert host is None and port is None
  831. create_server = functools.partial(
  832. loop.create_unix_server, factory, path, **kwargs
  833. )
  834. else:
  835. create_server = functools.partial(
  836. loop.create_server, factory, host, port, **kwargs
  837. )
  838. # This is a coroutine function.
  839. self._create_server = create_server
  840. self.ws_server = ws_server
  841. # async with serve(...)
  842. async def __aenter__(self) -> WebSocketServer:
  843. return await self
  844. async def __aexit__(
  845. self,
  846. exc_type: Optional[Type[BaseException]],
  847. exc_value: Optional[BaseException],
  848. traceback: Optional[TracebackType],
  849. ) -> None:
  850. self.ws_server.close()
  851. await self.ws_server.wait_closed()
  852. # await serve(...)
  853. def __await__(self) -> Generator[Any, None, WebSocketServer]:
  854. # Create a suitable iterator by calling __await__ on a coroutine.
  855. return self.__await_impl__().__await__()
  856. async def __await_impl__(self) -> WebSocketServer:
  857. server = await self._create_server()
  858. self.ws_server.wrap(server)
  859. return self.ws_server
  860. # yield from serve(...)
  861. __iter__ = __await__
  862. serve = Serve
  863. def unix_serve(
  864. ws_handler: Callable[[WebSocketServerProtocol, str], Awaitable[Any]],
  865. path: Optional[str] = None,
  866. **kwargs: Any,
  867. ) -> Serve:
  868. """
  869. Similar to :func:`serve`, but for listening on Unix sockets.
  870. This function calls the event loop's
  871. :meth:`~asyncio.loop.create_unix_server` method.
  872. It is only available on Unix.
  873. It's useful for deploying a server behind a reverse proxy such as nginx.
  874. :param path: file system path to the Unix socket
  875. """
  876. return serve(ws_handler, path=path, unix=True, **kwargs)