server.py 34 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942
  1. from __future__ import annotations
  2. import concurrent.futures
  3. import hmac
  4. import http
  5. import logging
  6. import os
  7. import re
  8. import selectors
  9. import socket
  10. import ssl as ssl_module
  11. import sys
  12. import threading
  13. import time
  14. import warnings
  15. from collections.abc import Iterable, Sequence
  16. from types import TracebackType
  17. from typing import Any, Callable, Mapping, Self, cast
  18. from ..exceptions import InvalidHeader
  19. from ..extensions.base import ServerExtensionFactory
  20. from ..extensions.permessage_deflate import enable_server_permessage_deflate
  21. from ..frames import CloseCode
  22. from ..headers import (
  23. build_www_authenticate_basic,
  24. parse_authorization_basic,
  25. validate_subprotocols,
  26. )
  27. from ..http11 import SERVER, Request, Response
  28. from ..protocol import CONNECTING, OPEN, Event
  29. from ..server import ServerProtocol
  30. from ..typing import LoggerLike, Origin, PathLike, StatusLike, Subprotocol
  31. from ..utils import get_socket_name
  32. from .connection import Connection, broadcast
  33. from .utils import Deadline
  34. __all__ = [
  35. "broadcast",
  36. "serve",
  37. "unix_serve",
  38. "ServerConnection",
  39. "Server",
  40. "basic_auth",
  41. ]
  42. class ServerConnection(Connection):
  43. """
  44. :mod:`threading` implementation of a WebSocket server connection.
  45. :class:`ServerConnection` provides :meth:`recv` and :meth:`send` methods for
  46. receiving and sending messages.
  47. It supports iteration to receive messages::
  48. for message in websocket:
  49. process(message)
  50. The iterator exits normally when the connection is closed with code
  51. 1000 (OK) or 1001 (going away) or without a close code. It raises a
  52. :exc:`~websockets.exceptions.ConnectionClosedError` when the connection is
  53. closed with any other code.
  54. The ``ping_interval``, ``ping_timeout``, ``close_timeout``, and
  55. ``max_queue`` arguments have the same meaning as in :func:`serve`.
  56. Args:
  57. socket: Socket connected to a WebSocket client.
  58. protocol: Sans-I/O connection.
  59. server: Server that manages this connection.
  60. """
  61. def __init__(
  62. self,
  63. sock: socket.socket,
  64. protocol: ServerProtocol,
  65. server: Server,
  66. *,
  67. ping_interval: float | None = 20,
  68. ping_timeout: float | None = 20,
  69. close_timeout: float | None = 10,
  70. max_queue: int | None | tuple[int | None, int | None] = 16,
  71. ) -> None:
  72. self.protocol: ServerProtocol
  73. self.request_rcvd = threading.Event()
  74. super().__init__(
  75. sock,
  76. protocol,
  77. ping_interval=ping_interval,
  78. ping_timeout=ping_timeout,
  79. close_timeout=close_timeout,
  80. max_queue=max_queue,
  81. )
  82. self.server = server
  83. self.username: str # see basic_auth()
  84. self.handler: Callable[[ServerConnection], None] # see route()
  85. self.handler_kwargs: Mapping[str, Any] # see route()
  86. def respond(self, status: StatusLike, text: str) -> Response:
  87. """
  88. Create a plain text HTTP response.
  89. ``process_request`` and ``process_response`` may call this method to
  90. return an HTTP response instead of performing the WebSocket opening
  91. handshake.
  92. You can modify the response before returning it, for example by changing
  93. HTTP headers.
  94. Args:
  95. status: HTTP status code.
  96. text: HTTP response body; it will be encoded to UTF-8.
  97. Returns:
  98. HTTP response to send to the client.
  99. """
  100. return self.protocol.reject(status, text)
  101. def handshake(
  102. self,
  103. process_request: (
  104. Callable[
  105. [ServerConnection, Request],
  106. Response | None,
  107. ]
  108. | None
  109. ) = None,
  110. process_response: (
  111. Callable[
  112. [ServerConnection, Request, Response],
  113. Response | None,
  114. ]
  115. | None
  116. ) = None,
  117. server_header: str | None = SERVER,
  118. timeout: float | None = None,
  119. ) -> None:
  120. """
  121. Perform the opening handshake.
  122. """
  123. if not self.request_rcvd.wait(timeout):
  124. raise TimeoutError("timed out while waiting for handshake request")
  125. if self.request is not None:
  126. response = None
  127. if process_request is not None:
  128. try:
  129. response = process_request(self, self.request)
  130. except Exception as exc:
  131. self.protocol.handshake_exc = exc
  132. self.logger.error("process_request failed", exc_info=True)
  133. response = self.protocol.reject(
  134. http.HTTPStatus.INTERNAL_SERVER_ERROR,
  135. (
  136. "Failed to open a WebSocket connection.\n"
  137. "See server log for more information.\n"
  138. ),
  139. )
  140. if response is None:
  141. self.response = self.protocol.accept(self.request)
  142. else:
  143. self.response = response
  144. if server_header is not None:
  145. self.response.headers["Server"] = server_header
  146. response = None
  147. if process_response is not None:
  148. try:
  149. response = process_response(self, self.request, self.response)
  150. except Exception as exc:
  151. self.protocol.handshake_exc = exc
  152. self.logger.error("process_response failed", exc_info=True)
  153. response = self.protocol.reject(
  154. http.HTTPStatus.INTERNAL_SERVER_ERROR,
  155. (
  156. "Failed to open a WebSocket connection.\n"
  157. "See server log for more information.\n"
  158. ),
  159. )
  160. if response is not None:
  161. self.response = response
  162. # Reject the connection if the server started closing during the
  163. # opening handshake. shutdown() runs a loop to catch cases where
  164. # the server shuts down between this check and send_response().
  165. if (
  166. self.response.status_code == http.HTTPStatus.SWITCHING_PROTOCOLS
  167. and self.server.socket_closed.is_set()
  168. ):
  169. self.response = self.protocol.reject(
  170. http.HTTPStatus.SERVICE_UNAVAILABLE,
  171. "Server is shutting down.\n",
  172. )
  173. # Don't respond if the connection was closed during the handshake.
  174. if self.state is CONNECTING:
  175. with self.send_context(expected_state=CONNECTING):
  176. self.protocol.send_response(self.response)
  177. def process_event(self, event: Event) -> None:
  178. """
  179. Process one incoming event.
  180. """
  181. # First event - handshake request.
  182. if self.request is None:
  183. assert isinstance(event, Request)
  184. self.request = event
  185. self.request_rcvd.set()
  186. # Later events - frames.
  187. else:
  188. super().process_event(event)
  189. def recv_events(self) -> None:
  190. """
  191. Read incoming data from the socket and process events.
  192. """
  193. try:
  194. super().recv_events()
  195. finally:
  196. # If the connection is closed during the handshake, unblock it.
  197. self.request_rcvd.set()
  198. class Server:
  199. """
  200. WebSocket server returned by :func:`serve`.
  201. This class mirrors partially the API of :class:`~socketserver.BaseServer`.
  202. Args:
  203. socket: Server socket accepting new connections.
  204. handler: Handler for one connection. It receives the socket and address
  205. returned by :meth:`~socket.socket.accept`.
  206. logger: Logger for this server.
  207. It defaults to ``logging.getLogger("websockets.server")``.
  208. See the :doc:`logging guide <../../topics/logging>` for details.
  209. """
  210. SHUTDOWN_POLLING_INTERVAL = 0.1 # seconds
  211. def __init__(
  212. self,
  213. sock: socket.socket,
  214. handler: Callable[[socket.socket, Any], None],
  215. logger: LoggerLike | None = None,
  216. ) -> None:
  217. self.socket = sock
  218. self.handler = handler
  219. if logger is None:
  220. logger = logging.getLogger("websockets.server")
  221. self.logger = logger
  222. # Synchronize access to all_connections and handler_threads.
  223. self.lock = threading.Lock()
  224. # Keep track of active connections and connection handler threads.
  225. self.all_connections: set[ServerConnection] = set()
  226. self.handler_threads: set[threading.Thread] = set()
  227. # On Windows, closing the socket wakes up the poller in serve_forever(),
  228. # making the notification mechanism unnecessary.
  229. if sys.platform != "win32":
  230. self.shutdown_watcher, self.shutdown_notifier = socket.socketpair()
  231. # Set when serve_forever() no longer accepts new connections and starts
  232. # threads to handle them.
  233. self.socket_closed = threading.Event()
  234. @property
  235. def connections(self) -> set[ServerConnection]:
  236. """
  237. Set of active connections.
  238. This property contains all connections that completed the opening
  239. handshake successfully and didn't start the closing handshake yet.
  240. It can be useful in combination with :func:`~broadcast`.
  241. """
  242. with self.lock:
  243. return {
  244. connection
  245. for connection in self.all_connections
  246. if connection.protocol.state is OPEN
  247. }
  248. def serve_forever(self) -> None:
  249. """
  250. See :meth:`socketserver.BaseServer.serve_forever`.
  251. This method doesn't return. Calling :meth:`shutdown` from another thread
  252. stops the server.
  253. Typical use::
  254. with serve(...) as server:
  255. server.serve_forever()
  256. """
  257. poller = selectors.DefaultSelector()
  258. if sys.platform != "win32":
  259. poller.register(self.shutdown_watcher, selectors.EVENT_READ)
  260. try:
  261. try:
  262. # Capture the socket name to prevent a race condition in case
  263. # the socket is closed between registering it with the poller
  264. # and logging the "server listening on ..." message.
  265. socket_name = get_socket_name(self.socket)
  266. poller.register(self.socket, selectors.EVENT_READ)
  267. except (OSError, ValueError): # pragma: no cover
  268. # shutdown() ran before get_socket_name() or poller.register().
  269. # This may result in:
  270. # * OSError: [Errno 9] Bad file descriptor
  271. # (only seen on free-threaded Python for poller.register())
  272. # * ValueError: Invalid file descriptor: -1
  273. return
  274. self.logger.info("server listening on %s", socket_name)
  275. while True:
  276. poller.select()
  277. try:
  278. # If the socket is closed, this raises an exception and
  279. # exits the loop; no need to check what select() returned.
  280. sock, addr = self.socket.accept()
  281. except OSError:
  282. break
  283. # shutdown() can let existing connections terminate on their own
  284. # or close them. Either way, it waits for connection handlers to
  285. # terminate, so there's no point using daemon threads.
  286. thread = threading.Thread(target=self.handler, args=(sock, addr))
  287. # The thread must be registered in self.handler_threads now,
  288. # before it's started. If it was registered in sock_handler(),
  289. # a race condition could happen when closing the server after
  290. # starting the thread but before it executes.
  291. with self.lock:
  292. self.handler_threads.add(thread)
  293. thread.start()
  294. finally:
  295. self.socket_closed.set()
  296. if sys.platform != "win32":
  297. self.shutdown_watcher.close()
  298. def shutdown(
  299. self,
  300. close_connections: bool = True,
  301. code: CloseCode | int = CloseCode.GOING_AWAY,
  302. reason: str = "",
  303. ) -> None:
  304. """
  305. Close the server.
  306. * Close the listening socket to stop accepting new connections.
  307. * When ``close_connections`` is :obj:`True`, which is the default, close
  308. existing connections. Specifically:
  309. * Reject opening WebSocket connections with an HTTP 503 (service
  310. unavailable) error. This happens when the server accepted the TCP
  311. connection but didn't complete the opening handshake before closing.
  312. * Close open WebSocket connections with code 1001 (going away).
  313. ``code`` and ``reason`` can be customized, for example to use code
  314. 1012 (service restart).
  315. * Wait until all connection handlers terminate.
  316. :meth:`shutdown` is idempotent.
  317. """
  318. self.logger.info("server closing")
  319. # Stop accepting new connections.
  320. self.socket.close()
  321. if sys.platform != "win32":
  322. try:
  323. self.shutdown_notifier.send(b"x")
  324. except OSError:
  325. pass # shutdown() was already called
  326. finally:
  327. self.shutdown_notifier.close()
  328. # Wait until serve_forever() no longer accepts new connections nor
  329. # starts threads to handle them, meaning that self.handler_threads
  330. # won't get new entries.
  331. # Also reject OPENING connections with HTTP 503 — see handshake().
  332. self.socket_closed.wait()
  333. # Close OPEN connections.
  334. if close_connections:
  335. # At this point, all threads are started, but some may still be in
  336. # the opening handshake. Close open connections until no thread is
  337. # executing anymore. Some threads may be cleaning up; in that case
  338. # they're expected to terminate quickly, so waiting is fine.
  339. while True:
  340. with self.lock:
  341. # Inline self.connections because it acquires self.lock,
  342. # which isn't reentrant.
  343. connections = [
  344. connection
  345. for connection in self.all_connections
  346. if connection.protocol.state is OPEN
  347. ]
  348. threads = list(self.handler_threads)
  349. # No threads are executing anymore. Server is fully closed.
  350. if not threads:
  351. break
  352. # Some threads are still executing, but no connections are OPEN.
  353. # Wait for connections to complete the opening handshake, or for
  354. # handler threads to terminate.
  355. if not connections:
  356. time.sleep(self.SHUTDOWN_POLLING_INTERVAL)
  357. continue
  358. # Close open connections and wait until they're closed.
  359. with concurrent.futures.ThreadPoolExecutor() as executor:
  360. for connection in connections:
  361. executor.submit(connection.close, code, reason)
  362. else:
  363. # At this point, all threads are started.
  364. with self.lock:
  365. threads = list(self.handler_threads)
  366. # Wait until all connection handlers terminate.
  367. for thread in threads:
  368. # This raises RuntimeError if shutdown() is called from a
  369. # connection handler. It's documented to return after all
  370. # connection handlers terminate, which is impossible when
  371. # it's called from a connection handler.
  372. thread.join()
  373. self.logger.info("server closed")
  374. def fileno(self) -> int:
  375. """
  376. See :meth:`socketserver.BaseServer.fileno`.
  377. """
  378. return self.socket.fileno()
  379. def __enter__(self) -> Self:
  380. return self
  381. def __exit__(
  382. self,
  383. exc_type: type[BaseException] | None,
  384. exc_value: BaseException | None,
  385. traceback: TracebackType | None,
  386. ) -> None:
  387. self.shutdown()
  388. def __getattr__(name: str) -> Any:
  389. if name == "WebSocketServer":
  390. warnings.warn( # deprecated in 13.0 - 2024-08-20
  391. "WebSocketServer was renamed to Server",
  392. DeprecationWarning,
  393. )
  394. return Server
  395. raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
  396. def serve(
  397. handler: Callable[[ServerConnection], None],
  398. host: str | None = None,
  399. port: int | None = None,
  400. *,
  401. # TCP/TLS
  402. sock: socket.socket | None = None,
  403. ssl: ssl_module.SSLContext | None = None,
  404. # WebSocket
  405. origins: Sequence[Origin | re.Pattern[str] | None] | None = None,
  406. extensions: Sequence[ServerExtensionFactory] | None = None,
  407. subprotocols: Sequence[Subprotocol] | None = None,
  408. select_subprotocol: (
  409. Callable[
  410. [ServerConnection, Sequence[Subprotocol]],
  411. Subprotocol | None,
  412. ]
  413. | None
  414. ) = None,
  415. compression: str | None = "deflate",
  416. # HTTP
  417. process_request: (
  418. Callable[
  419. [ServerConnection, Request],
  420. Response | None,
  421. ]
  422. | None
  423. ) = None,
  424. process_response: (
  425. Callable[
  426. [ServerConnection, Request, Response],
  427. Response | None,
  428. ]
  429. | None
  430. ) = None,
  431. server_header: str | None = SERVER,
  432. # Timeouts
  433. open_timeout: float | None = 10,
  434. ping_interval: float | None = 20,
  435. ping_timeout: float | None = 20,
  436. close_timeout: float | None = 10,
  437. # Limits
  438. max_size: int | None | tuple[int | None, int | None] = 2**20,
  439. max_queue: int | None | tuple[int | None, int | None] = 16,
  440. # Logging
  441. logger: LoggerLike | None = None,
  442. # Escape hatch for advanced customization
  443. create_connection: type[ServerConnection] | None = None,
  444. **kwargs: Any,
  445. ) -> Server:
  446. """
  447. Create a WebSocket server listening on ``host`` and ``port``.
  448. Whenever a client connects, the server creates a :class:`ServerConnection`,
  449. performs the opening handshake, and delegates to the ``handler`` function.
  450. The handler receives the :class:`ServerConnection` instance, which you can
  451. use to send and receive messages.
  452. Once the handler completes, either normally or with an exception, the server
  453. performs the closing handshake and closes the connection.
  454. This function returns a :class:`Server` object whose API mirrors
  455. :class:`~socketserver.BaseServer`. Treat it as a context manager to ensure
  456. that it will be closed gracefully and call :meth:`~Server.serve_forever` to
  457. serve requests::
  458. from websockets.sync.server import serve
  459. def handler(websocket):
  460. ...
  461. with serve(handler, ...) as server:
  462. server.serve_forever()
  463. You can stop the server gracefully by calling its :meth:`~Server.shutdown`
  464. method from another thread::
  465. server.shutdown()
  466. By default, closing the server closes connections with code 1001 (going
  467. away). You can keep connections alive and wait for clients to disconnect::
  468. server.shutdown(close_connections=False)
  469. Args:
  470. handler: Connection handler. It receives the WebSocket connection,
  471. which is a :class:`ServerConnection`, in argument.
  472. host: Network interfaces the server binds to.
  473. See :func:`~socket.create_server` for details.
  474. port: TCP port the server listens on.
  475. See :func:`~socket.create_server` for details.
  476. sock: Preexisting TCP socket. ``sock`` replaces ``host`` and ``port``.
  477. You may call :func:`socket.create_server` to create a suitable TCP
  478. socket.
  479. ssl: Configuration for enabling TLS on the connection.
  480. origins: Acceptable values of the ``Origin`` header, for defending
  481. against Cross-Site WebSocket Hijacking attacks. Values can be
  482. :class:`str` to test for an exact match or regular expressions
  483. compiled by :func:`re.compile` to test against a pattern. Include
  484. :obj:`None` in the list if the lack of an origin is acceptable.
  485. extensions: List of supported extensions, in order in which they
  486. should be negotiated and run.
  487. subprotocols: List of supported subprotocols, in order of decreasing
  488. preference.
  489. select_subprotocol: Callback for selecting a subprotocol among
  490. those supported by the client and the server. It receives a
  491. :class:`ServerConnection` (not a
  492. :class:`~websockets.server.ServerProtocol`!) instance and a list of
  493. subprotocols offered by the client. Other than the first argument,
  494. it has the same behavior as the
  495. :meth:`ServerProtocol.select_subprotocol
  496. <websockets.server.ServerProtocol.select_subprotocol>` method.
  497. compression: The "permessage-deflate" extension is enabled by default.
  498. Set ``compression`` to :obj:`None` to disable it. See the
  499. :doc:`compression guide <../../topics/compression>` for details.
  500. process_request: Intercept the request during the opening handshake.
  501. Return an HTTP response to force the response. Return :obj:`None` to
  502. continue normally. When you force an HTTP 101 Continue response, the
  503. handshake is successful. Else, the connection is aborted.
  504. process_response: Intercept the response during the opening handshake.
  505. Modify the response or return a new HTTP response to force the
  506. response. Return :obj:`None` to continue normally. When you force an
  507. HTTP 101 Continue response, the handshake is successful. Else, the
  508. connection is aborted.
  509. server_header: Value of the ``Server`` response header.
  510. It defaults to ``"Python/x.y.z websockets/X.Y"``. Setting it to
  511. :obj:`None` removes the header.
  512. open_timeout: Timeout for opening connections in seconds.
  513. :obj:`None` disables the timeout.
  514. ping_interval: Interval between keepalive pings in seconds.
  515. :obj:`None` disables keepalive.
  516. ping_timeout: Timeout for keepalive pings in seconds.
  517. :obj:`None` disables timeouts.
  518. close_timeout: Timeout for closing connections in seconds.
  519. :obj:`None` disables the timeout.
  520. max_size: Maximum size of incoming messages in bytes.
  521. :obj:`None` disables the limit. You may pass a ``(max_message_size,
  522. max_fragment_size)`` tuple to set different limits for messages and
  523. fragments when you expect long messages sent in short fragments.
  524. max_queue: High-water mark of the buffer where frames are received.
  525. It defaults to 16 frames. The low-water mark defaults to ``max_queue
  526. // 4``. You may pass a ``(high, low)`` tuple to set the high-water
  527. and low-water marks. If you want to disable flow control entirely,
  528. you may set it to ``None``, although that's a bad idea.
  529. logger: Logger for this server.
  530. It defaults to ``logging.getLogger("websockets.server")``.
  531. See the :doc:`logging guide <../../topics/logging>` for details.
  532. create_connection: Factory for the :class:`ServerConnection` managing
  533. the connection. Set it to a wrapper or a subclass to customize
  534. connection handling.
  535. Any other keyword arguments are passed to :func:`~socket.create_server`.
  536. """
  537. # Process parameters
  538. # Backwards compatibility: ssl used to be called ssl_context.
  539. if ssl is None and "ssl_context" in kwargs:
  540. ssl = kwargs.pop("ssl_context")
  541. warnings.warn( # deprecated in 13.0 - 2024-08-20
  542. "ssl_context was renamed to ssl",
  543. DeprecationWarning,
  544. )
  545. if subprotocols is not None:
  546. validate_subprotocols(subprotocols)
  547. if compression == "deflate":
  548. extensions = enable_server_permessage_deflate(extensions)
  549. elif compression is not None:
  550. raise ValueError(f"unsupported compression: {compression}")
  551. if create_connection is None:
  552. create_connection = ServerConnection
  553. # Bind socket and listen
  554. # Private APIs for unix_connect()
  555. unix: bool = kwargs.pop("unix", False)
  556. path: str | None = kwargs.pop("path", None)
  557. if sock is None:
  558. if unix:
  559. if path is None:
  560. raise ValueError("missing path argument")
  561. kwargs.setdefault("family", socket.AF_UNIX)
  562. sock = socket.create_server(os.fspath(path), **kwargs)
  563. else:
  564. sock = socket.create_server((host, port), **kwargs)
  565. else:
  566. if host is not None:
  567. raise ValueError("host is incompatible with sock")
  568. if port is not None:
  569. raise ValueError("port is incompatible with sock")
  570. if path is not None:
  571. raise ValueError("path is incompatible with sock")
  572. # Initialize TLS wrapper
  573. if ssl is not None:
  574. sock = ssl.wrap_socket(
  575. sock,
  576. server_side=True,
  577. # Delay TLS handshake until after we set a timeout on the socket.
  578. do_handshake_on_connect=False,
  579. )
  580. # Define request handler
  581. def sock_handler(sock: socket.socket, addr: Any) -> None:
  582. """
  583. Handle the lifecycle of a WebSocket connection.
  584. Since this function doesn't have a caller that can handle exceptions,
  585. it attempts to log relevant ones.
  586. It guarantees that the TCP connection is closed before exiting.
  587. """
  588. # Calculate timeouts on the TLS and WebSocket handshakes.
  589. # The TLS timeout must be set on the socket, then removed
  590. # to avoid conflicting with the WebSocket timeout in handshake().
  591. deadline = Deadline(open_timeout)
  592. try:
  593. # Disable Nagle algorithm
  594. if sock.family in {socket.AF_INET, socket.AF_INET6}:
  595. sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True)
  596. # Perform TLS handshake
  597. if ssl is not None:
  598. sock.settimeout(deadline.timeout())
  599. # mypy cannot figure this out
  600. assert isinstance(sock, ssl_module.SSLSocket)
  601. sock.do_handshake()
  602. sock.settimeout(None)
  603. # Create a closure to give select_subprotocol access to connection.
  604. protocol_select_subprotocol: (
  605. Callable[
  606. [ServerProtocol, Sequence[Subprotocol]],
  607. Subprotocol | None,
  608. ]
  609. | None
  610. ) = None
  611. if select_subprotocol is not None:
  612. def protocol_select_subprotocol(
  613. protocol: ServerProtocol,
  614. subprotocols: Sequence[Subprotocol],
  615. ) -> Subprotocol | None:
  616. # mypy doesn't know that select_subprotocol is immutable.
  617. assert select_subprotocol is not None
  618. # Ensure this function is only used in the intended context.
  619. assert protocol is connection.protocol
  620. return select_subprotocol(connection, subprotocols)
  621. # Initialize WebSocket protocol
  622. protocol = ServerProtocol(
  623. origins=origins,
  624. extensions=extensions,
  625. subprotocols=subprotocols,
  626. select_subprotocol=protocol_select_subprotocol,
  627. max_size=max_size,
  628. logger=logger,
  629. )
  630. # Initialize WebSocket connection
  631. assert create_connection is not None # help mypy
  632. connection = create_connection(
  633. sock,
  634. protocol,
  635. server,
  636. ping_interval=ping_interval,
  637. ping_timeout=ping_timeout,
  638. close_timeout=close_timeout,
  639. max_queue=max_queue,
  640. )
  641. except Exception:
  642. try:
  643. sock.close()
  644. return
  645. finally:
  646. with server.lock:
  647. server.handler_threads.discard(threading.current_thread())
  648. try:
  649. connection.handshake(
  650. process_request,
  651. process_response,
  652. server_header,
  653. deadline.timeout(),
  654. )
  655. if connection.protocol.state is not OPEN:
  656. connection.close_socket()
  657. return
  658. with server.lock:
  659. server.all_connections.add(connection)
  660. connection.start_keepalive()
  661. try:
  662. with connection:
  663. handler(connection)
  664. except Exception:
  665. connection.logger.error("connection handler failed", exc_info=True)
  666. finally:
  667. with server.lock:
  668. server.all_connections.discard(connection)
  669. except Exception:
  670. # Don't leak sockets when the opening handshake times out or an
  671. # unexpected error occurs.
  672. connection.close_socket()
  673. finally:
  674. with server.lock:
  675. server.handler_threads.discard(threading.current_thread())
  676. # Initialize server
  677. # The server variable is captured by the closure of sock_handler().
  678. server = Server(sock, sock_handler, logger)
  679. return server
  680. def unix_serve(
  681. handler: Callable[[ServerConnection], None],
  682. path: PathLike | None = None,
  683. **kwargs: Any,
  684. ) -> Server:
  685. """
  686. Create a WebSocket server listening on a Unix socket.
  687. This function accepts the same keyword arguments as :func:`serve`.
  688. It's only available on Unix.
  689. It's useful for deploying a server behind a reverse proxy such as nginx.
  690. Args:
  691. handler: Connection handler. It receives the WebSocket connection,
  692. which is a :class:`ServerConnection`, in argument.
  693. path: File system path to the Unix socket.
  694. """
  695. return serve(handler, unix=True, path=path, **kwargs)
  696. def is_credentials(credentials: Any) -> bool:
  697. try:
  698. username, password = credentials
  699. except (TypeError, ValueError):
  700. return False
  701. else:
  702. return isinstance(username, str) and isinstance(password, str)
  703. def basic_auth(
  704. realm: str = "",
  705. credentials: tuple[str, str] | Iterable[tuple[str, str]] | None = None,
  706. check_credentials: Callable[[str, str], bool] | None = None,
  707. ) -> Callable[[ServerConnection, Request], Response | None]:
  708. """
  709. Factory for ``process_request`` to enforce HTTP Basic Authentication.
  710. :func:`basic_auth` is designed to integrate with :func:`serve` as follows::
  711. from websockets.sync.server import basic_auth, serve
  712. with serve(
  713. ...,
  714. process_request=basic_auth(
  715. realm="my dev server",
  716. credentials=("hello", "iloveyou"),
  717. ),
  718. ):
  719. If authentication succeeds, the connection's ``username`` attribute is set.
  720. If it fails, the server responds with an HTTP 401 Unauthorized status.
  721. One of ``credentials`` or ``check_credentials`` must be provided; not both.
  722. Args:
  723. realm: Scope of protection. It should contain only ASCII characters
  724. because the encoding of non-ASCII characters is undefined. Refer to
  725. section 2.2 of :rfc:`7235` for details.
  726. credentials: Hard coded authorized credentials. It can be a
  727. ``(username, password)`` pair or a list of such pairs.
  728. check_credentials: Function that verifies credentials.
  729. It receives ``username`` and ``password`` arguments and returns
  730. whether they're valid.
  731. Raises:
  732. TypeError: If ``credentials`` or ``check_credentials`` is wrong.
  733. ValueError: If ``credentials`` and ``check_credentials`` are both
  734. provided or both not provided.
  735. """
  736. if (credentials is None) == (check_credentials is None):
  737. raise ValueError("provide either credentials or check_credentials")
  738. if credentials is not None:
  739. if is_credentials(credentials):
  740. credentials_list = [cast(tuple[str, str], credentials)]
  741. elif isinstance(credentials, Iterable):
  742. credentials_list = list(cast(Iterable[tuple[str, str]], credentials))
  743. if not all(is_credentials(item) for item in credentials_list):
  744. raise TypeError(f"invalid credentials argument: {credentials}")
  745. else:
  746. raise TypeError(f"invalid credentials argument: {credentials}")
  747. credentials_dict = dict(credentials_list)
  748. def check_credentials(username: str, password: str) -> bool:
  749. try:
  750. expected_password = credentials_dict[username]
  751. except KeyError:
  752. return False
  753. return hmac.compare_digest(expected_password, password)
  754. assert check_credentials is not None # help mypy
  755. def process_request(
  756. connection: ServerConnection,
  757. request: Request,
  758. ) -> Response | None:
  759. """
  760. Perform HTTP Basic Authentication.
  761. If it succeeds, set the connection's ``username`` attribute and return
  762. :obj:`None`. If it fails, return an HTTP 401 Unauthorized responss.
  763. """
  764. try:
  765. authorization = request.headers["Authorization"]
  766. except KeyError:
  767. response = connection.respond(
  768. http.HTTPStatus.UNAUTHORIZED,
  769. "Missing credentials\n",
  770. )
  771. response.headers["WWW-Authenticate"] = build_www_authenticate_basic(realm)
  772. return response
  773. try:
  774. username, password = parse_authorization_basic(authorization)
  775. except InvalidHeader:
  776. response = connection.respond(
  777. http.HTTPStatus.UNAUTHORIZED,
  778. "Unsupported credentials\n",
  779. )
  780. response.headers["WWW-Authenticate"] = build_www_authenticate_basic(realm)
  781. return response
  782. if not check_credentials(username, password):
  783. response = connection.respond(
  784. http.HTTPStatus.UNAUTHORIZED,
  785. "Invalid credentials\n",
  786. )
  787. response.headers["WWW-Authenticate"] = build_www_authenticate_basic(realm)
  788. return response
  789. connection.username = username
  790. return None
  791. return process_request