server.py 33 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911
  1. from __future__ import annotations
  2. import asyncio
  3. import hmac
  4. import http
  5. import logging
  6. import re
  7. import socket
  8. from collections.abc import Awaitable, Coroutine, Generator, Iterable, Sequence
  9. from types import TracebackType
  10. from typing import Any, Callable, Mapping, Self, cast
  11. from ..exceptions import InvalidHeader
  12. from ..extensions.base import ServerExtensionFactory
  13. from ..extensions.permessage_deflate import enable_server_permessage_deflate
  14. from ..frames import CloseCode
  15. from ..headers import (
  16. build_www_authenticate_basic,
  17. parse_authorization_basic,
  18. validate_subprotocols,
  19. )
  20. from ..http11 import SERVER, Request, Response
  21. from ..protocol import CONNECTING, OPEN, Event
  22. from ..server import ServerProtocol
  23. from ..typing import LoggerLike, Origin, StatusLike, Subprotocol
  24. from ..utils import get_socket_name
  25. from .connection import Connection, broadcast
  26. __all__ = [
  27. "broadcast",
  28. "serve",
  29. "unix_serve",
  30. "ServerConnection",
  31. "Server",
  32. "basic_auth",
  33. ]
  34. class ServerConnection(Connection):
  35. """
  36. :mod:`asyncio` implementation of a WebSocket server connection.
  37. :class:`ServerConnection` provides :meth:`recv` and :meth:`send` methods for
  38. receiving and sending messages.
  39. It supports asynchronous iteration to receive messages::
  40. async for message in websocket:
  41. await process(message)
  42. The iterator exits normally when the connection is closed with code
  43. 1000 (OK) or 1001 (going away) or without a close code. It raises a
  44. :exc:`~websockets.exceptions.ConnectionClosedError` when the connection is
  45. closed with any other code.
  46. The ``ping_interval``, ``ping_timeout``, ``close_timeout``, ``max_queue``,
  47. and ``write_limit`` arguments have the same meaning as in :func:`serve`.
  48. Args:
  49. protocol: Sans-I/O connection.
  50. server: Server that manages this connection.
  51. """
  52. def __init__(
  53. self,
  54. protocol: ServerProtocol,
  55. server: Server,
  56. *,
  57. ping_interval: float | None = 20,
  58. ping_timeout: float | None = 20,
  59. close_timeout: float | None = 10,
  60. max_queue: int | None | tuple[int | None, int | None] = 16,
  61. write_limit: int | tuple[int, int | None] = 2**15,
  62. ) -> None:
  63. self.protocol: ServerProtocol
  64. super().__init__(
  65. protocol,
  66. ping_interval=ping_interval,
  67. ping_timeout=ping_timeout,
  68. close_timeout=close_timeout,
  69. max_queue=max_queue,
  70. write_limit=write_limit,
  71. )
  72. self.server = server
  73. self.request_rcvd: asyncio.Future[None] = self.loop.create_future()
  74. self.username: str # see basic_auth()
  75. self.handler: Callable[[ServerConnection], Awaitable[None]] # see route()
  76. self.handler_kwargs: Mapping[str, Any] # see route()
  77. def respond(self, status: StatusLike, text: str) -> Response:
  78. """
  79. Create a plain text HTTP response.
  80. ``process_request`` and ``process_response`` may call this method to
  81. return an HTTP response instead of performing the WebSocket opening
  82. handshake.
  83. You can modify the response before returning it, for example by changing
  84. HTTP headers.
  85. Args:
  86. status: HTTP status code.
  87. text: HTTP response body; it will be encoded to UTF-8.
  88. Returns:
  89. HTTP response to send to the client.
  90. """
  91. return self.protocol.reject(status, text)
  92. async def handshake(
  93. self,
  94. process_request: (
  95. Callable[
  96. [ServerConnection, Request],
  97. Awaitable[Response | None] | Response | None,
  98. ]
  99. | None
  100. ) = None,
  101. process_response: (
  102. Callable[
  103. [ServerConnection, Request, Response],
  104. Awaitable[Response | None] | Response | None,
  105. ]
  106. | None
  107. ) = None,
  108. server_header: str | None = SERVER,
  109. ) -> None:
  110. """
  111. Perform the opening handshake.
  112. """
  113. await asyncio.wait(
  114. [self.request_rcvd, self.connection_lost_waiter],
  115. return_when=asyncio.FIRST_COMPLETED,
  116. )
  117. if self.request is not None:
  118. response = None
  119. if process_request is not None:
  120. try:
  121. response = process_request(self, self.request)
  122. if isinstance(response, Awaitable):
  123. response = await response
  124. except Exception as exc:
  125. self.protocol.handshake_exc = exc
  126. self.logger.error("process_request failed", exc_info=True)
  127. response = self.protocol.reject(
  128. http.HTTPStatus.INTERNAL_SERVER_ERROR,
  129. (
  130. "Failed to open a WebSocket connection.\n"
  131. "See server log for more information.\n"
  132. ),
  133. )
  134. if response is None:
  135. self.response = self.protocol.accept(self.request)
  136. else:
  137. assert isinstance(response, Response) # help mypy
  138. self.response = response
  139. if server_header is not None:
  140. self.response.headers["Server"] = server_header
  141. response = None
  142. if process_response is not None:
  143. try:
  144. response = process_response(self, self.request, self.response)
  145. if isinstance(response, Awaitable):
  146. response = await response
  147. except Exception as exc:
  148. self.protocol.handshake_exc = exc
  149. self.logger.error("process_response failed", exc_info=True)
  150. response = self.protocol.reject(
  151. http.HTTPStatus.INTERNAL_SERVER_ERROR,
  152. (
  153. "Failed to open a WebSocket connection.\n"
  154. "See server log for more information.\n"
  155. ),
  156. )
  157. if response is not None:
  158. assert isinstance(response, Response) # help mypy
  159. self.response = response
  160. # Reject the connection if the server started closing during the
  161. # opening handshake. Don't yield before send_response() to avoid
  162. # a race condition after checking if the server is closing.
  163. if (
  164. self.response.status_code == http.HTTPStatus.SWITCHING_PROTOCOLS
  165. and not self.server.is_serving()
  166. ):
  167. self.response = self.protocol.reject(
  168. http.HTTPStatus.SERVICE_UNAVAILABLE,
  169. "Server is shutting down.\n",
  170. )
  171. # Don't respond if the connection was closed during the handshake.
  172. if self.state is CONNECTING:
  173. async with self.send_context(expected_state=CONNECTING):
  174. self.protocol.send_response(self.response)
  175. def process_event(self, event: Event) -> None:
  176. """
  177. Process one incoming event.
  178. """
  179. # First event - handshake request.
  180. if self.request is None:
  181. assert isinstance(event, Request)
  182. self.request = event
  183. self.request_rcvd.set_result(None)
  184. # Later events - frames.
  185. else:
  186. super().process_event(event)
  187. def connection_made(self, transport: asyncio.BaseTransport) -> None:
  188. super().connection_made(transport)
  189. # The handler task must be registered in self.handler_tasks now. If it
  190. # was registered inside the task, a race condition could happen when
  191. # closing the server after scheduling the task but before it executes.
  192. handler_task = self.loop.create_task(self.server.handler(self))
  193. self.server.handler_tasks.add(handler_task)
  194. class Server:
  195. """
  196. WebSocket server returned by :func:`serve`.
  197. This class mirrors most of the API of :class:`asyncio.Server`, with the
  198. following differences:
  199. * You can invoke :func:`serve` as ``async with serve(...) as server: ...``
  200. in addition to ``server = await serve(...)`` to start the server.
  201. * It doesn't provide ``close_clients`` or ``abort_clients``; by default,
  202. :meth:`close` closes existing connections with code 1001 (going away).
  203. Args:
  204. handler: Handler for one connection. It receives an asyncio protocol.
  205. logger: Logger for this server.
  206. It defaults to ``logging.getLogger("websockets.server")``.
  207. See the :doc:`logging guide <../../topics/logging>` for details.
  208. """
  209. def __init__(
  210. self,
  211. create_server: Callable[[], Coroutine[Any, Any, asyncio.Server]],
  212. handler: Callable[[ServerConnection], Coroutine[Any, Any, None]],
  213. logger: LoggerLike | None = None,
  214. ) -> None:
  215. self.create_server = create_server
  216. self.handler = handler
  217. if logger is None:
  218. logger = logging.getLogger("websockets.server")
  219. self.logger = logger
  220. # Keep track of active connections and connection handler tasks.
  221. self.all_connections: set[ServerConnection] = set()
  222. self.handler_tasks: set[asyncio.Task[None]] = set()
  223. # Task responsible for closing the server and terminating connections.
  224. self.close_task: asyncio.Task[None] | None = None
  225. # Completed when the server is closed and connections are terminated.
  226. loop = asyncio.get_running_loop()
  227. self.handlers_waiter: asyncio.Future[None] = loop.create_future()
  228. @property
  229. def connections(self) -> set[ServerConnection]:
  230. """
  231. Set of active connections.
  232. This property contains all connections that completed the opening
  233. handshake successfully and didn't start the closing handshake yet.
  234. It can be useful in combination with :func:`~broadcast`.
  235. """
  236. return {
  237. connection
  238. for connection in self.all_connections
  239. if connection.protocol.state is OPEN
  240. }
  241. def close(
  242. self,
  243. close_connections: bool = True,
  244. code: CloseCode | int = CloseCode.GOING_AWAY,
  245. reason: str = "",
  246. ) -> None:
  247. """
  248. Close the server.
  249. * Close the underlying :class:`asyncio.Server`.
  250. * When ``close_connections`` is :obj:`True`, which is the default, close
  251. existing connections. Specifically:
  252. * Reject opening WebSocket connections with an HTTP 503 (service
  253. unavailable) error. This happens when the server accepted the TCP
  254. connection but didn't complete the opening handshake before closing.
  255. * Close open WebSocket connections with code 1001 (going away).
  256. ``code`` and ``reason`` can be customized, for example to use code
  257. 1012 (service restart).
  258. * Wait until all connection handlers have returned.
  259. :meth:`close` is idempotent.
  260. """
  261. if self.close_task is None:
  262. self.close_task = self.get_loop().create_task(
  263. self._close(close_connections, code, reason)
  264. )
  265. async def _close(
  266. self,
  267. close_connections: bool = True,
  268. code: CloseCode | int = CloseCode.GOING_AWAY,
  269. reason: str = "",
  270. ) -> None:
  271. """
  272. Implementation of :meth:`close`.
  273. This calls :meth:`~asyncio.Server.close` on the underlying
  274. :class:`asyncio.Server` object to stop accepting new connections and
  275. then closes open connections.
  276. """
  277. self.logger.info("server closing")
  278. # Stop accepting new connections.
  279. # Also reject OPENING connections with HTTP 503 — see handshake().
  280. self.server.close()
  281. # Close OPEN connections.
  282. if close_connections:
  283. close_tasks = [
  284. asyncio.create_task(connection.close(code, reason))
  285. for connection in self.all_connections
  286. if connection.protocol.state is OPEN
  287. ]
  288. # asyncio.wait doesn't accept an empty first argument.
  289. if close_tasks:
  290. await asyncio.wait(close_tasks)
  291. # Wait until all TCP connections are closed.
  292. await self.server.wait_closed()
  293. # Wait until all connection handlers have returned.
  294. # asyncio.wait doesn't accept an empty first argument.
  295. if self.handler_tasks:
  296. await asyncio.wait(self.handler_tasks)
  297. # Tell wait_closed() to return.
  298. self.handlers_waiter.set_result(None)
  299. self.logger.info("server closed")
  300. def get_loop(self) -> asyncio.AbstractEventLoop:
  301. """
  302. See :meth:`asyncio.Server.get_loop`.
  303. """
  304. return self.server.get_loop()
  305. async def start_serving(self) -> None:
  306. """
  307. See :meth:`asyncio.Server.start_serving`.
  308. Typical use::
  309. server = await serve(..., start_serving=False)
  310. # perform additional setup here...
  311. # ... then start the server
  312. await server.start_serving()
  313. """
  314. await self.server.start_serving()
  315. async def serve_forever(self) -> None:
  316. """
  317. See :meth:`asyncio.Server.serve_forever`.
  318. Typical use::
  319. server = await serve(...)
  320. # this coroutine doesn't return
  321. # canceling it stops the server
  322. await server.serve_forever()
  323. This is an alternative to using :func:`serve` as an asynchronous context
  324. manager. Shutdown is triggered by canceling :meth:`serve_forever` or by
  325. calling :meth:`~Server.close` from another task.
  326. """
  327. try:
  328. if not self.is_serving():
  329. await self.start_serving()
  330. # If close() is called, wait_closed() will return, and we'll exit.
  331. await self.wait_closed()
  332. except asyncio.CancelledError:
  333. try:
  334. self.close()
  335. await self.wait_closed()
  336. finally:
  337. raise
  338. def is_serving(self) -> bool:
  339. """
  340. See :meth:`asyncio.Server.is_serving`.
  341. """
  342. return self.server.is_serving()
  343. async def wait_closed(self) -> None:
  344. """
  345. Wait until the server is closed.
  346. When :meth:`wait_closed` returns, all TCP connections are closed and
  347. all connection handlers have returned.
  348. To ensure a fast shutdown, a connection handler should always be
  349. awaiting at least one of:
  350. * :meth:`~ServerConnection.recv`: when the connection is closed,
  351. it raises :exc:`~websockets.exceptions.ConnectionClosedOK`;
  352. * :meth:`~ServerConnection.wait_closed`: when the connection is
  353. closed, it returns.
  354. Then the connection handler is immediately notified of the shutdown;
  355. it can clean up and exit.
  356. """
  357. await asyncio.shield(self.handlers_waiter)
  358. @property
  359. def sockets(self) -> tuple[socket.socket, ...]:
  360. """
  361. See :attr:`asyncio.Server.sockets`.
  362. """
  363. return self.server.sockets
  364. async def _await(self) -> Self:
  365. if not hasattr(self, "server"):
  366. self.server = await self.create_server()
  367. if self.server.is_serving():
  368. for sock in self.server.sockets:
  369. self.logger.info("server listening on %s", get_socket_name(sock))
  370. return self
  371. def __await__(self) -> Generator[Any, None, Self]:
  372. # Create a suitable iterator by calling __await__ on a coroutine.
  373. return self._await().__await__()
  374. async def __aenter__(self) -> Self:
  375. return await self
  376. async def __aexit__(
  377. self,
  378. exc_type: type[BaseException] | None,
  379. exc_value: BaseException | None,
  380. traceback: TracebackType | None,
  381. ) -> None:
  382. self.close()
  383. await self.wait_closed()
  384. # serve() is declared as a function rather than a coroutine in order to support
  385. # async with serve(...) as server: ... in addition to server = await serve(...).
  386. def serve(
  387. handler: Callable[[ServerConnection], Awaitable[None]],
  388. host: str | None = None,
  389. port: int | None = None,
  390. *,
  391. # WebSocket
  392. origins: Sequence[Origin | re.Pattern[str] | None] | None = None,
  393. extensions: Sequence[ServerExtensionFactory] | None = None,
  394. subprotocols: Sequence[Subprotocol] | None = None,
  395. select_subprotocol: (
  396. Callable[
  397. [ServerConnection, Sequence[Subprotocol]],
  398. Subprotocol | None,
  399. ]
  400. | None
  401. ) = None,
  402. compression: str | None = "deflate",
  403. # HTTP
  404. process_request: (
  405. Callable[
  406. [ServerConnection, Request],
  407. Awaitable[Response | None] | Response | None,
  408. ]
  409. | None
  410. ) = None,
  411. process_response: (
  412. Callable[
  413. [ServerConnection, Request, Response],
  414. Awaitable[Response | None] | Response | None,
  415. ]
  416. | None
  417. ) = None,
  418. server_header: str | None = SERVER,
  419. # Timeouts
  420. open_timeout: float | None = 10,
  421. ping_interval: float | None = 20,
  422. ping_timeout: float | None = 20,
  423. close_timeout: float | None = 10,
  424. # Limits
  425. max_size: int | None | tuple[int | None, int | None] = 2**20,
  426. max_queue: int | None | tuple[int | None, int | None] = 16,
  427. write_limit: int | tuple[int, int | None] = 2**15,
  428. # Logging
  429. logger: LoggerLike | None = None,
  430. # Escape hatch for advanced customization
  431. create_connection: type[ServerConnection] | None = None,
  432. # Other keyword arguments are passed to loop.create_server
  433. **kwargs: Any,
  434. ) -> Server:
  435. """
  436. Create a WebSocket server listening on ``host`` and ``port``.
  437. Whenever a client connects, the server creates a :class:`ServerConnection`,
  438. performs the opening handshake, and delegates to the ``handler`` coroutine.
  439. The handler receives the :class:`ServerConnection` instance, which you can
  440. use to send and receive messages.
  441. Once the handler completes, either normally or with an exception, the server
  442. performs the closing handshake and closes the connection.
  443. This function returns a :class:`Server` object whose API mirrors
  444. :class:`asyncio.Server`. Treat it as an asynchronous context manager to
  445. serve requests and ensure that the server will be closed gracefully::
  446. from websockets.asyncio.server import serve
  447. async def handler(websocket):
  448. ...
  449. # set this event to exit the server
  450. stop = asyncio.Event()
  451. async with serve(handler, host, port):
  452. await stop.wait()
  453. Alternatively, await it and call :meth:`~Server.serve_forever` to serve
  454. requests, then cancel it or call :meth:`~Server.close` to stop the server::
  455. server = await serve(handler, host, port)
  456. await server.serve_forever()
  457. The following pattern is functional but redundant: by the time the context
  458. manager exits, :meth:`~Server.serve_forever` has already closed the server::
  459. async with serve(handler, host, port) as server:
  460. await server.serve_forever()
  461. Args:
  462. handler: Connection handler. It receives the WebSocket connection,
  463. which is a :class:`ServerConnection`, in argument.
  464. host: Network interfaces the server binds to.
  465. See :meth:`~asyncio.loop.create_server` for details.
  466. port: TCP port the server listens on.
  467. See :meth:`~asyncio.loop.create_server` for details.
  468. origins: Acceptable values of the ``Origin`` header, for defending
  469. against Cross-Site WebSocket Hijacking attacks. Values can be
  470. :class:`str` to test for an exact match or regular expressions
  471. compiled by :func:`re.compile` to test against a pattern. Include
  472. :obj:`None` in the list if the lack of an origin is acceptable.
  473. extensions: List of supported extensions, in order in which they
  474. should be negotiated and run.
  475. subprotocols: List of supported subprotocols, in order of decreasing
  476. preference.
  477. select_subprotocol: Callback for selecting a subprotocol among
  478. those supported by the client and the server. It receives a
  479. :class:`ServerConnection` (not a
  480. :class:`~websockets.server.ServerProtocol`!) instance and a list of
  481. subprotocols offered by the client. Other than the first argument,
  482. it has the same behavior as the
  483. :meth:`ServerProtocol.select_subprotocol
  484. <websockets.server.ServerProtocol.select_subprotocol>` method.
  485. compression: The "permessage-deflate" extension is enabled by default.
  486. Set ``compression`` to :obj:`None` to disable it. See the
  487. :doc:`compression guide <../../topics/compression>` for details.
  488. process_request: Intercept the request during the opening handshake.
  489. Return an HTTP response to force the response or :obj:`None` to
  490. continue normally. When you force an HTTP 101 Continue response, the
  491. handshake is successful. Else, the connection is aborted.
  492. ``process_request`` may be a function or a coroutine.
  493. process_response: Intercept the response during the opening handshake.
  494. Return an HTTP response to force the response or :obj:`None` to
  495. continue normally. When you force an HTTP 101 Continue response, the
  496. handshake is successful. Else, the connection is aborted.
  497. ``process_response`` may be a function or a coroutine.
  498. server_header: Value of the ``Server`` response header.
  499. It defaults to ``"Python/x.y.z websockets/X.Y"``. Setting it to
  500. :obj:`None` removes the header.
  501. open_timeout: Timeout for opening connections in seconds.
  502. :obj:`None` disables the timeout.
  503. ping_interval: Interval between keepalive pings in seconds.
  504. :obj:`None` disables keepalive.
  505. ping_timeout: Timeout for keepalive pings in seconds.
  506. :obj:`None` disables timeouts.
  507. close_timeout: Timeout for closing connections in seconds.
  508. :obj:`None` disables the timeout.
  509. max_size: Maximum size of incoming messages in bytes.
  510. :obj:`None` disables the limit. You may pass a ``(max_message_size,
  511. max_fragment_size)`` tuple to set different limits for messages and
  512. fragments when you expect long messages sent in short fragments.
  513. max_queue: High-water mark of the buffer where frames are received.
  514. It defaults to 16 frames. The low-water mark defaults to ``max_queue
  515. // 4``. You may pass a ``(high, low)`` tuple to set the high-water
  516. and low-water marks. If you want to disable flow control entirely,
  517. you may set it to ``None``, although that's a bad idea.
  518. write_limit: High-water mark of write buffer in bytes. It is passed to
  519. :meth:`~asyncio.WriteTransport.set_write_buffer_limits`. It defaults
  520. to 32 KiB. You may pass a ``(high, low)`` tuple to set the
  521. high-water and low-water marks.
  522. logger: Logger for this server.
  523. It defaults to ``logging.getLogger("websockets.server")``.
  524. See the :doc:`logging guide <../../topics/logging>` for details.
  525. create_connection: Factory for the :class:`ServerConnection` managing
  526. the connection. Set it to a wrapper or a subclass to customize
  527. connection handling.
  528. Any other keyword arguments are passed to the event loop's
  529. :meth:`~asyncio.loop.create_server` method.
  530. For example:
  531. * You can set ``ssl`` to a :class:`~ssl.SSLContext` to enable TLS.
  532. * You can set ``sock`` to provide a preexisting TCP socket. You may call
  533. :func:`socket.create_server` (not to be confused with the event loop's
  534. :meth:`~asyncio.loop.create_server` method) to create a suitable server
  535. socket and customize it.
  536. * You can set ``start_serving`` to ``False`` to start accepting connections
  537. only after you call :meth:`~Server.start_serving()` or
  538. :meth:`~Server.serve_forever()`.
  539. """
  540. if subprotocols is not None:
  541. validate_subprotocols(subprotocols)
  542. if compression == "deflate":
  543. extensions = enable_server_permessage_deflate(extensions)
  544. elif compression is not None:
  545. raise ValueError(f"unsupported compression: {compression}")
  546. if create_connection is None:
  547. create_connection = ServerConnection
  548. if kwargs.get("ssl") is not None:
  549. kwargs.setdefault("ssl_handshake_timeout", open_timeout)
  550. kwargs.setdefault("ssl_shutdown_timeout", close_timeout)
  551. async def create_server() -> asyncio.Server:
  552. loop = asyncio.get_running_loop()
  553. if kwargs.pop("unix", False):
  554. return await loop.create_unix_server(protocol_factory, **kwargs)
  555. else:
  556. # mypy cannot tell that kwargs must provide sock when port is None.
  557. return await loop.create_server(protocol_factory, host, port, **kwargs) # type: ignore[arg-type]
  558. def protocol_factory() -> ServerConnection:
  559. """
  560. Create an asyncio protocol for managing a WebSocket connection.
  561. """
  562. # Create a closure to give select_subprotocol access to connection.
  563. protocol_select_subprotocol: (
  564. Callable[
  565. [ServerProtocol, Sequence[Subprotocol]],
  566. Subprotocol | None,
  567. ]
  568. | None
  569. ) = None
  570. if select_subprotocol is not None:
  571. def protocol_select_subprotocol(
  572. protocol: ServerProtocol,
  573. subprotocols: Sequence[Subprotocol],
  574. ) -> Subprotocol | None:
  575. # mypy doesn't know that select_subprotocol is immutable.
  576. assert select_subprotocol is not None
  577. # Ensure this function is only used in the intended context.
  578. assert protocol is connection.protocol
  579. return select_subprotocol(connection, subprotocols)
  580. # This is a protocol in the Sans-I/O implementation of websockets.
  581. protocol = ServerProtocol(
  582. origins=origins,
  583. extensions=extensions,
  584. subprotocols=subprotocols,
  585. select_subprotocol=protocol_select_subprotocol,
  586. max_size=max_size,
  587. logger=logger,
  588. )
  589. # This is a connection in websockets and a protocol in asyncio.
  590. connection = create_connection(
  591. protocol,
  592. server,
  593. ping_interval=ping_interval,
  594. ping_timeout=ping_timeout,
  595. close_timeout=close_timeout,
  596. max_queue=max_queue,
  597. write_limit=write_limit,
  598. )
  599. return connection
  600. async def protocol_handler(connection: ServerConnection) -> None:
  601. """
  602. Handle the lifecycle of a WebSocket connection.
  603. Since this coroutine doesn't have a caller that can handle
  604. exceptions, it attempts to log relevant ones.
  605. It guarantees that the TCP connection is closed before exiting.
  606. """
  607. try:
  608. # Apply open_timeout to the WebSocket handshake.
  609. # Use ssl_handshake_timeout for the TLS handshake.
  610. async with asyncio.timeout(open_timeout):
  611. await connection.handshake(
  612. process_request,
  613. process_response,
  614. server_header,
  615. )
  616. if connection.protocol.state is not OPEN:
  617. connection.transport.abort()
  618. return
  619. server.all_connections.add(connection)
  620. connection.start_keepalive()
  621. try:
  622. await handler(connection)
  623. except Exception:
  624. connection.logger.error("connection handler failed", exc_info=True)
  625. await connection.close(CloseCode.INTERNAL_ERROR)
  626. else:
  627. await connection.close()
  628. finally:
  629. server.all_connections.discard(connection)
  630. except Exception:
  631. # Don't leak connections when the opening handshake times out or
  632. # an unexpected error occurs.
  633. connection.transport.abort()
  634. finally:
  635. server.handler_tasks.discard(asyncio.current_task())
  636. # The server variable is captured by the closure of conn_handler() and
  637. # protocol_factory().
  638. server = Server(create_server, protocol_handler, logger)
  639. return server
  640. def unix_serve(
  641. handler: Callable[[ServerConnection], Awaitable[None]],
  642. path: str | None = None,
  643. **kwargs: Any,
  644. ) -> Server:
  645. """
  646. Create a WebSocket server listening on a Unix socket.
  647. This function is identical to :func:`serve`, except the ``host`` and
  648. ``port`` arguments are replaced by ``path``. It's only available on Unix.
  649. It's useful for deploying a server behind a reverse proxy such as nginx.
  650. Args:
  651. handler: Connection handler. It receives the WebSocket connection,
  652. which is a :class:`ServerConnection`, in argument.
  653. path: File system path to the Unix socket.
  654. """
  655. return serve(handler, unix=True, path=path, **kwargs)
  656. def is_credentials(credentials: Any) -> bool:
  657. try:
  658. username, password = credentials
  659. except (TypeError, ValueError):
  660. return False
  661. else:
  662. return isinstance(username, str) and isinstance(password, str)
  663. def basic_auth(
  664. realm: str = "",
  665. credentials: tuple[str, str] | Iterable[tuple[str, str]] | None = None,
  666. check_credentials: Callable[[str, str], Awaitable[bool] | bool] | None = None,
  667. ) -> Callable[[ServerConnection, Request], Awaitable[Response | None]]:
  668. """
  669. Factory for ``process_request`` to enforce HTTP Basic Authentication.
  670. :func:`basic_auth` is designed to integrate with :func:`serve` as follows::
  671. from websockets.asyncio.server import basic_auth, serve
  672. async with serve(
  673. ...,
  674. process_request=basic_auth(
  675. realm="my dev server",
  676. credentials=("hello", "iloveyou"),
  677. ),
  678. ):
  679. If authentication succeeds, the connection's ``username`` attribute is set.
  680. If it fails, the server responds with an HTTP 401 Unauthorized status.
  681. One of ``credentials`` or ``check_credentials`` must be provided; not both.
  682. Args:
  683. realm: Scope of protection. It should contain only ASCII characters
  684. because the encoding of non-ASCII characters is undefined. Refer to
  685. section 2.2 of :rfc:`7235` for details.
  686. credentials: Hard coded authorized credentials. It can be a
  687. ``(username, password)`` pair or a list of such pairs.
  688. check_credentials: Function or coroutine that verifies credentials.
  689. It receives ``username`` and ``password`` arguments and returns
  690. whether they're valid.
  691. Raises:
  692. TypeError: If ``credentials`` or ``check_credentials`` is wrong.
  693. ValueError: If ``credentials`` and ``check_credentials`` are both
  694. provided or both not provided.
  695. """
  696. if (credentials is None) == (check_credentials is None):
  697. raise ValueError("provide either credentials or check_credentials")
  698. if credentials is not None:
  699. if is_credentials(credentials):
  700. credentials_list = [cast(tuple[str, str], credentials)]
  701. elif isinstance(credentials, Iterable):
  702. credentials_list = list(cast(Iterable[tuple[str, str]], credentials))
  703. if not all(is_credentials(item) for item in credentials_list):
  704. raise TypeError(f"invalid credentials argument: {credentials}")
  705. else:
  706. raise TypeError(f"invalid credentials argument: {credentials}")
  707. credentials_dict = dict(credentials_list)
  708. def check_credentials(username: str, password: str) -> bool:
  709. try:
  710. expected_password = credentials_dict[username]
  711. except KeyError:
  712. return False
  713. return hmac.compare_digest(expected_password, password)
  714. assert check_credentials is not None # help mypy
  715. async def process_request(
  716. connection: ServerConnection,
  717. request: Request,
  718. ) -> Response | None:
  719. """
  720. Perform HTTP Basic Authentication.
  721. If it succeeds, set the connection's ``username`` attribute and return
  722. :obj:`None`. If it fails, return an HTTP 401 Unauthorized responss.
  723. """
  724. try:
  725. authorization = request.headers["Authorization"]
  726. except KeyError:
  727. response = connection.respond(
  728. http.HTTPStatus.UNAUTHORIZED,
  729. "Missing credentials\n",
  730. )
  731. response.headers["WWW-Authenticate"] = build_www_authenticate_basic(realm)
  732. return response
  733. try:
  734. username, password = parse_authorization_basic(authorization)
  735. except InvalidHeader:
  736. response = connection.respond(
  737. http.HTTPStatus.UNAUTHORIZED,
  738. "Unsupported credentials\n",
  739. )
  740. response.headers["WWW-Authenticate"] = build_www_authenticate_basic(realm)
  741. return response
  742. valid_credentials = check_credentials(username, password)
  743. if isinstance(valid_credentials, Awaitable):
  744. valid_credentials = await valid_credentials
  745. if not valid_credentials:
  746. response = connection.respond(
  747. http.HTTPStatus.UNAUTHORIZED,
  748. "Invalid credentials\n",
  749. )
  750. response.headers["WWW-Authenticate"] = build_www_authenticate_basic(realm)
  751. return response
  752. connection.username = username
  753. return None
  754. return process_request