Nevar pievienot vairāk kā 25 tēmas Tēmai ir jāsākas ar burtu vai ciparu, tā var saturēt domu zīmes ('-') un var būt līdz 35 simboliem gara.
 
 
 
 

997 rindas
37 KiB

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