server.py 37 KB

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