Funktionierender Prototyp des Serious Games zur Vermittlung von Wissen zu Software-Engineering-Arbeitsmodellen.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

server.py 44KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196
  1. from __future__ import annotations
  2. import asyncio
  3. import email.utils
  4. import functools
  5. import http
  6. import inspect
  7. import logging
  8. import socket
  9. import warnings
  10. from types import TracebackType
  11. from typing import (
  12. Any,
  13. Awaitable,
  14. Callable,
  15. Generator,
  16. Iterable,
  17. List,
  18. Optional,
  19. Sequence,
  20. Set,
  21. Tuple,
  22. Type,
  23. Union,
  24. cast,
  25. )
  26. from ..datastructures import Headers, HeadersLike, MultipleValuesError
  27. from ..exceptions import (
  28. AbortHandshake,
  29. InvalidHandshake,
  30. InvalidHeader,
  31. InvalidMessage,
  32. InvalidOrigin,
  33. InvalidUpgrade,
  34. NegotiationError,
  35. )
  36. from ..extensions import Extension, ServerExtensionFactory
  37. from ..extensions.permessage_deflate import enable_server_permessage_deflate
  38. from ..headers import (
  39. build_extension,
  40. parse_extension,
  41. parse_subprotocol,
  42. validate_subprotocols,
  43. )
  44. from ..http import USER_AGENT
  45. from ..protocol import State
  46. from ..typing import ExtensionHeader, LoggerLike, Origin, Subprotocol
  47. from .compatibility import asyncio_timeout, loop_if_py_lt_38
  48. from .handshake import build_response, check_request
  49. from .http import read_request
  50. from .protocol import WebSocketCommonProtocol
  51. __all__ = ["serve", "unix_serve", "WebSocketServerProtocol", "WebSocketServer"]
  52. HeadersLikeOrCallable = Union[HeadersLike, Callable[[str, Headers], HeadersLike]]
  53. HTTPResponse = Tuple[http.HTTPStatus, HeadersLike, bytes]
  54. class WebSocketServerProtocol(WebSocketCommonProtocol):
  55. """
  56. WebSocket server connection.
  57. :class:`WebSocketServerProtocol` provides :meth:`recv` and :meth:`send`
  58. coroutines for receiving and sending messages.
  59. It supports asynchronous iteration to receive messages::
  60. async for message in websocket:
  61. await process(message)
  62. The iterator exits normally when the connection is closed with close code
  63. 1000 (OK) or 1001 (going away) or without a close code. It raises
  64. a :exc:`~websockets.exceptions.ConnectionClosedError` when the connection
  65. is closed with any other code.
  66. You may customize the opening handshake in a subclass by
  67. overriding :meth:`process_request` or :meth:`select_subprotocol`.
  68. Args:
  69. ws_server: WebSocket server that created this connection.
  70. See :func:`serve` for the documentation of ``ws_handler``, ``logger``, ``origins``,
  71. ``extensions``, ``subprotocols``, ``extra_headers``, and ``server_header``.
  72. See :class:`~websockets.legacy.protocol.WebSocketCommonProtocol` for the
  73. documentation of ``ping_interval``, ``ping_timeout``, ``close_timeout``,
  74. ``max_size``, ``max_queue``, ``read_limit``, and ``write_limit``.
  75. """
  76. is_client = False
  77. side = "server"
  78. def __init__(
  79. self,
  80. ws_handler: Union[
  81. Callable[[WebSocketServerProtocol], Awaitable[Any]],
  82. Callable[[WebSocketServerProtocol, str], Awaitable[Any]], # deprecated
  83. ],
  84. ws_server: WebSocketServer,
  85. *,
  86. logger: Optional[LoggerLike] = None,
  87. origins: Optional[Sequence[Optional[Origin]]] = None,
  88. extensions: Optional[Sequence[ServerExtensionFactory]] = None,
  89. subprotocols: Optional[Sequence[Subprotocol]] = None,
  90. extra_headers: Optional[HeadersLikeOrCallable] = None,
  91. server_header: Optional[str] = USER_AGENT,
  92. process_request: Optional[
  93. Callable[[str, Headers], Awaitable[Optional[HTTPResponse]]]
  94. ] = None,
  95. select_subprotocol: Optional[
  96. Callable[[Sequence[Subprotocol], Sequence[Subprotocol]], Subprotocol]
  97. ] = None,
  98. open_timeout: Optional[float] = 10,
  99. **kwargs: Any,
  100. ) -> None:
  101. if logger is None:
  102. logger = logging.getLogger("websockets.server")
  103. super().__init__(logger=logger, **kwargs)
  104. # For backwards compatibility with 6.0 or earlier.
  105. if origins is not None and "" in origins:
  106. warnings.warn("use None instead of '' in origins", DeprecationWarning)
  107. origins = [None if origin == "" else origin for origin in origins]
  108. # For backwards compatibility with 10.0 or earlier. Done here in
  109. # addition to serve to trigger the deprecation warning on direct
  110. # use of WebSocketServerProtocol.
  111. self.ws_handler = remove_path_argument(ws_handler)
  112. self.ws_server = ws_server
  113. self.origins = origins
  114. self.available_extensions = extensions
  115. self.available_subprotocols = subprotocols
  116. self.extra_headers = extra_headers
  117. self.server_header = server_header
  118. self._process_request = process_request
  119. self._select_subprotocol = select_subprotocol
  120. self.open_timeout = open_timeout
  121. def connection_made(self, transport: asyncio.BaseTransport) -> None:
  122. """
  123. Register connection and initialize a task to handle it.
  124. """
  125. super().connection_made(transport)
  126. # Register the connection with the server before creating the handler
  127. # task. Registering at the beginning of the handler coroutine would
  128. # create a race condition between the creation of the task, which
  129. # schedules its execution, and the moment the handler starts running.
  130. self.ws_server.register(self)
  131. self.handler_task = self.loop.create_task(self.handler())
  132. async def handler(self) -> None:
  133. """
  134. Handle the lifecycle of a WebSocket connection.
  135. Since this method doesn't have a caller able to handle exceptions, it
  136. attempts to log relevant ones and guarantees that the TCP connection is
  137. closed before exiting.
  138. """
  139. try:
  140. try:
  141. async with asyncio_timeout(self.open_timeout):
  142. await self.handshake(
  143. origins=self.origins,
  144. available_extensions=self.available_extensions,
  145. available_subprotocols=self.available_subprotocols,
  146. extra_headers=self.extra_headers,
  147. )
  148. # Remove this branch when dropping support for Python < 3.8
  149. # because CancelledError no longer inherits Exception.
  150. except asyncio.CancelledError: # pragma: no cover
  151. raise
  152. except asyncio.TimeoutError: # pragma: no cover
  153. raise
  154. except ConnectionError:
  155. raise
  156. except Exception as exc:
  157. if isinstance(exc, AbortHandshake):
  158. status, headers, body = exc.status, exc.headers, exc.body
  159. elif isinstance(exc, InvalidOrigin):
  160. if self.debug:
  161. self.logger.debug("! invalid origin", exc_info=True)
  162. status, headers, body = (
  163. http.HTTPStatus.FORBIDDEN,
  164. Headers(),
  165. f"Failed to open a WebSocket connection: {exc}.\n".encode(),
  166. )
  167. elif isinstance(exc, InvalidUpgrade):
  168. if self.debug:
  169. self.logger.debug("! invalid upgrade", exc_info=True)
  170. status, headers, body = (
  171. http.HTTPStatus.UPGRADE_REQUIRED,
  172. Headers([("Upgrade", "websocket")]),
  173. (
  174. f"Failed to open a WebSocket connection: {exc}.\n"
  175. f"\n"
  176. f"You cannot access a WebSocket server directly "
  177. f"with a browser. You need a WebSocket client.\n"
  178. ).encode(),
  179. )
  180. elif isinstance(exc, InvalidHandshake):
  181. if self.debug:
  182. self.logger.debug("! invalid handshake", exc_info=True)
  183. status, headers, body = (
  184. http.HTTPStatus.BAD_REQUEST,
  185. Headers(),
  186. f"Failed to open a WebSocket connection: {exc}.\n".encode(),
  187. )
  188. else:
  189. self.logger.error("opening handshake failed", exc_info=True)
  190. status, headers, body = (
  191. http.HTTPStatus.INTERNAL_SERVER_ERROR,
  192. Headers(),
  193. (
  194. b"Failed to open a WebSocket connection.\n"
  195. b"See server log for more information.\n"
  196. ),
  197. )
  198. headers.setdefault("Date", email.utils.formatdate(usegmt=True))
  199. if self.server_header is not None:
  200. headers.setdefault("Server", self.server_header)
  201. headers.setdefault("Content-Length", str(len(body)))
  202. headers.setdefault("Content-Type", "text/plain")
  203. headers.setdefault("Connection", "close")
  204. self.write_http_response(status, headers, body)
  205. self.logger.info(
  206. "connection failed (%d %s)", status.value, status.phrase
  207. )
  208. await self.close_transport()
  209. return
  210. try:
  211. await self.ws_handler(self)
  212. except Exception:
  213. self.logger.error("connection handler failed", exc_info=True)
  214. if not self.closed:
  215. self.fail_connection(1011)
  216. raise
  217. try:
  218. await self.close()
  219. except ConnectionError:
  220. raise
  221. except Exception:
  222. self.logger.error("closing handshake failed", exc_info=True)
  223. raise
  224. except Exception:
  225. # Last-ditch attempt to avoid leaking connections on errors.
  226. try:
  227. self.transport.close()
  228. except Exception: # pragma: no cover
  229. pass
  230. finally:
  231. # Unregister the connection with the server when the handler task
  232. # terminates. Registration is tied to the lifecycle of the handler
  233. # task because the server waits for tasks attached to registered
  234. # connections before terminating.
  235. self.ws_server.unregister(self)
  236. self.logger.info("connection closed")
  237. async def read_http_request(self) -> Tuple[str, Headers]:
  238. """
  239. Read request line and headers from the HTTP request.
  240. If the request contains a body, it may be read from ``self.reader``
  241. after this coroutine returns.
  242. Raises:
  243. InvalidMessage: if the HTTP message is malformed or isn't an
  244. HTTP/1.1 GET request.
  245. """
  246. try:
  247. path, headers = await read_request(self.reader)
  248. except asyncio.CancelledError: # pragma: no cover
  249. raise
  250. except Exception as exc:
  251. raise InvalidMessage("did not receive a valid HTTP request") from exc
  252. if self.debug:
  253. self.logger.debug("< GET %s HTTP/1.1", path)
  254. for key, value in headers.raw_items():
  255. self.logger.debug("< %s: %s", key, value)
  256. self.path = path
  257. self.request_headers = headers
  258. return path, headers
  259. def write_http_response(
  260. self, status: http.HTTPStatus, headers: Headers, body: Optional[bytes] = None
  261. ) -> None:
  262. """
  263. Write status line and headers to the HTTP response.
  264. This coroutine is also able to write a response body.
  265. """
  266. self.response_headers = headers
  267. if self.debug:
  268. self.logger.debug("> HTTP/1.1 %d %s", status.value, status.phrase)
  269. for key, value in headers.raw_items():
  270. self.logger.debug("> %s: %s", key, value)
  271. if body is not None:
  272. self.logger.debug("> [body] (%d bytes)", len(body))
  273. # Since the status line and headers only contain ASCII characters,
  274. # we can keep this simple.
  275. response = f"HTTP/1.1 {status.value} {status.phrase}\r\n"
  276. response += str(headers)
  277. self.transport.write(response.encode())
  278. if body is not None:
  279. self.transport.write(body)
  280. async def process_request(
  281. self, path: str, request_headers: Headers
  282. ) -> Optional[HTTPResponse]:
  283. """
  284. Intercept the HTTP request and return an HTTP response if appropriate.
  285. You may override this method in a :class:`WebSocketServerProtocol`
  286. subclass, for example:
  287. * to return an HTTP 200 OK response on a given path; then a load
  288. balancer can use this path for a health check;
  289. * to authenticate the request and return an HTTP 401 Unauthorized or an
  290. HTTP 403 Forbidden when authentication fails.
  291. You may also override this method with the ``process_request``
  292. argument of :func:`serve` and :class:`WebSocketServerProtocol`. This
  293. is equivalent, except ``process_request`` won't have access to the
  294. protocol instance, so it can't store information for later use.
  295. :meth:`process_request` is expected to complete quickly. If it may run
  296. for a long time, then it should await :meth:`wait_closed` and exit if
  297. :meth:`wait_closed` completes, or else it could prevent the server
  298. from shutting down.
  299. Args:
  300. path: request path, including optional query string.
  301. request_headers: request headers.
  302. Returns:
  303. Optional[Tuple[http.HTTPStatus, HeadersLike, bytes]]: :obj:`None`
  304. to continue the WebSocket handshake normally.
  305. An HTTP response, represented by a 3-uple of the response status,
  306. headers, and body, to abort the WebSocket handshake and return
  307. that HTTP response instead.
  308. """
  309. if self._process_request is not None:
  310. response = self._process_request(path, request_headers)
  311. if isinstance(response, Awaitable):
  312. return await response
  313. else:
  314. # For backwards compatibility with 7.0.
  315. warnings.warn(
  316. "declare process_request as a coroutine", DeprecationWarning
  317. )
  318. return response
  319. return None
  320. @staticmethod
  321. def process_origin(
  322. headers: Headers, origins: Optional[Sequence[Optional[Origin]]] = None
  323. ) -> Optional[Origin]:
  324. """
  325. Handle the Origin HTTP request header.
  326. Args:
  327. headers: request headers.
  328. origins: optional list of acceptable origins.
  329. Raises:
  330. InvalidOrigin: if the origin isn't acceptable.
  331. """
  332. # "The user agent MUST NOT include more than one Origin header field"
  333. # per https://www.rfc-editor.org/rfc/rfc6454.html#section-7.3.
  334. try:
  335. origin = cast(Optional[Origin], headers.get("Origin"))
  336. except MultipleValuesError as exc:
  337. raise InvalidHeader("Origin", "more than one Origin header found") from exc
  338. if origins is not None:
  339. if origin not in origins:
  340. raise InvalidOrigin(origin)
  341. return origin
  342. @staticmethod
  343. def process_extensions(
  344. headers: Headers,
  345. available_extensions: Optional[Sequence[ServerExtensionFactory]],
  346. ) -> Tuple[Optional[str], List[Extension]]:
  347. """
  348. Handle the Sec-WebSocket-Extensions HTTP request header.
  349. Accept or reject each extension proposed in the client request.
  350. Negotiate parameters for accepted extensions.
  351. Return the Sec-WebSocket-Extensions HTTP response header and the list
  352. of accepted extensions.
  353. :rfc:`6455` leaves the rules up to the specification of each
  354. :extension.
  355. To provide this level of flexibility, for each extension proposed by
  356. the client, we check for a match with each extension available in the
  357. server configuration. If no match is found, the extension is ignored.
  358. If several variants of the same extension are proposed by the client,
  359. it may be accepted several times, which won't make sense in general.
  360. Extensions must implement their own requirements. For this purpose,
  361. the list of previously accepted extensions is provided.
  362. This process doesn't allow the server to reorder extensions. It can
  363. only select a subset of the extensions proposed by the client.
  364. Other requirements, for example related to mandatory extensions or the
  365. order of extensions, may be implemented by overriding this method.
  366. Args:
  367. headers: request headers.
  368. extensions: optional list of supported extensions.
  369. Raises:
  370. InvalidHandshake: to abort the handshake with an HTTP 400 error.
  371. """
  372. response_header_value: Optional[str] = None
  373. extension_headers: List[ExtensionHeader] = []
  374. accepted_extensions: List[Extension] = []
  375. header_values = headers.get_all("Sec-WebSocket-Extensions")
  376. if header_values and available_extensions:
  377. parsed_header_values: List[ExtensionHeader] = sum(
  378. [parse_extension(header_value) for header_value in header_values], []
  379. )
  380. for name, request_params in parsed_header_values:
  381. for ext_factory in available_extensions:
  382. # Skip non-matching extensions based on their name.
  383. if ext_factory.name != name:
  384. continue
  385. # Skip non-matching extensions based on their params.
  386. try:
  387. response_params, extension = ext_factory.process_request_params(
  388. request_params, accepted_extensions
  389. )
  390. except NegotiationError:
  391. continue
  392. # Add matching extension to the final list.
  393. extension_headers.append((name, response_params))
  394. accepted_extensions.append(extension)
  395. # Break out of the loop once we have a match.
  396. break
  397. # If we didn't break from the loop, no extension in our list
  398. # matched what the client sent. The extension is declined.
  399. # Serialize extension header.
  400. if extension_headers:
  401. response_header_value = build_extension(extension_headers)
  402. return response_header_value, accepted_extensions
  403. # Not @staticmethod because it calls self.select_subprotocol()
  404. def process_subprotocol(
  405. self, headers: Headers, available_subprotocols: Optional[Sequence[Subprotocol]]
  406. ) -> Optional[Subprotocol]:
  407. """
  408. Handle the Sec-WebSocket-Protocol HTTP request header.
  409. Return Sec-WebSocket-Protocol HTTP response header, which is the same
  410. as the selected subprotocol.
  411. Args:
  412. headers: request headers.
  413. available_subprotocols: optional list of supported subprotocols.
  414. Raises:
  415. InvalidHandshake: to abort the handshake with an HTTP 400 error.
  416. """
  417. subprotocol: Optional[Subprotocol] = None
  418. header_values = headers.get_all("Sec-WebSocket-Protocol")
  419. if header_values and available_subprotocols:
  420. parsed_header_values: List[Subprotocol] = sum(
  421. [parse_subprotocol(header_value) for header_value in header_values], []
  422. )
  423. subprotocol = self.select_subprotocol(
  424. parsed_header_values, available_subprotocols
  425. )
  426. return subprotocol
  427. def select_subprotocol(
  428. self,
  429. client_subprotocols: Sequence[Subprotocol],
  430. server_subprotocols: Sequence[Subprotocol],
  431. ) -> Optional[Subprotocol]:
  432. """
  433. Pick a subprotocol among those supported by the client and the server.
  434. If several subprotocols are available, select the preferred subprotocol
  435. by giving equal weight to the preferences of the client and the server.
  436. If no subprotocol is available, proceed without a subprotocol.
  437. You may provide a ``select_subprotocol`` argument to :func:`serve` or
  438. :class:`WebSocketServerProtocol` to override this logic. For example,
  439. you could reject the handshake if the client doesn't support a
  440. particular subprotocol, rather than accept the handshake without that
  441. subprotocol.
  442. Args:
  443. client_subprotocols: list of subprotocols offered by the client.
  444. server_subprotocols: list of subprotocols available on the server.
  445. Returns:
  446. Optional[Subprotocol]: Selected subprotocol, if a common subprotocol
  447. was found.
  448. :obj:`None` to continue without a subprotocol.
  449. """
  450. if self._select_subprotocol is not None:
  451. return self._select_subprotocol(client_subprotocols, server_subprotocols)
  452. subprotocols = set(client_subprotocols) & set(server_subprotocols)
  453. if not subprotocols:
  454. return None
  455. return sorted(
  456. subprotocols,
  457. key=lambda p: client_subprotocols.index(p) + server_subprotocols.index(p),
  458. )[0]
  459. async def handshake(
  460. self,
  461. origins: Optional[Sequence[Optional[Origin]]] = None,
  462. available_extensions: Optional[Sequence[ServerExtensionFactory]] = None,
  463. available_subprotocols: Optional[Sequence[Subprotocol]] = None,
  464. extra_headers: Optional[HeadersLikeOrCallable] = None,
  465. ) -> str:
  466. """
  467. Perform the server side of the opening handshake.
  468. Args:
  469. origins: list of acceptable values of the Origin HTTP header;
  470. include :obj:`None` if the lack of an origin is acceptable.
  471. extensions: list of supported extensions, in order in which they
  472. should be tried.
  473. subprotocols: list of supported subprotocols, in order of
  474. decreasing preference.
  475. extra_headers: arbitrary HTTP headers to add to the response when
  476. the handshake succeeds.
  477. Returns:
  478. str: path of the URI of the request.
  479. Raises:
  480. InvalidHandshake: if the handshake fails.
  481. """
  482. path, request_headers = await self.read_http_request()
  483. # Hook for customizing request handling, for example checking
  484. # authentication or treating some paths as plain HTTP endpoints.
  485. early_response_awaitable = self.process_request(path, request_headers)
  486. if isinstance(early_response_awaitable, Awaitable):
  487. early_response = await early_response_awaitable
  488. else:
  489. # For backwards compatibility with 7.0.
  490. warnings.warn("declare process_request as a coroutine", DeprecationWarning)
  491. early_response = early_response_awaitable
  492. # The connection may drop while process_request is running.
  493. if self.state is State.CLOSED:
  494. # This subclass of ConnectionError is silently ignored in handler().
  495. raise BrokenPipeError("connection closed during opening handshake")
  496. # Change the response to a 503 error if the server is shutting down.
  497. if not self.ws_server.is_serving():
  498. early_response = (
  499. http.HTTPStatus.SERVICE_UNAVAILABLE,
  500. [],
  501. b"Server is shutting down.\n",
  502. )
  503. if early_response is not None:
  504. raise AbortHandshake(*early_response)
  505. key = check_request(request_headers)
  506. self.origin = self.process_origin(request_headers, origins)
  507. extensions_header, self.extensions = self.process_extensions(
  508. request_headers, available_extensions
  509. )
  510. protocol_header = self.subprotocol = self.process_subprotocol(
  511. request_headers, available_subprotocols
  512. )
  513. response_headers = Headers()
  514. build_response(response_headers, key)
  515. if extensions_header is not None:
  516. response_headers["Sec-WebSocket-Extensions"] = extensions_header
  517. if protocol_header is not None:
  518. response_headers["Sec-WebSocket-Protocol"] = protocol_header
  519. if callable(extra_headers):
  520. extra_headers = extra_headers(path, self.request_headers)
  521. if extra_headers is not None:
  522. response_headers.update(extra_headers)
  523. response_headers.setdefault("Date", email.utils.formatdate(usegmt=True))
  524. if self.server_header is not None:
  525. response_headers.setdefault("Server", self.server_header)
  526. self.write_http_response(http.HTTPStatus.SWITCHING_PROTOCOLS, response_headers)
  527. self.logger.info("connection open")
  528. self.connection_open()
  529. return path
  530. class WebSocketServer:
  531. """
  532. WebSocket server returned by :func:`serve`.
  533. This class provides the same interface as :class:`~asyncio.Server`,
  534. notably the :meth:`~asyncio.Server.close`
  535. and :meth:`~asyncio.Server.wait_closed` methods.
  536. It keeps track of WebSocket connections in order to close them properly
  537. when shutting down.
  538. Args:
  539. logger: Logger for this server.
  540. It defaults to ``logging.getLogger("websockets.server")``.
  541. See the :doc:`logging guide <../../topics/logging>` for details.
  542. """
  543. def __init__(self, logger: Optional[LoggerLike] = None):
  544. if logger is None:
  545. logger = logging.getLogger("websockets.server")
  546. self.logger = logger
  547. # Keep track of active connections.
  548. self.websockets: Set[WebSocketServerProtocol] = set()
  549. # Task responsible for closing the server and terminating connections.
  550. self.close_task: Optional[asyncio.Task[None]] = None
  551. # Completed when the server is closed and connections are terminated.
  552. self.closed_waiter: asyncio.Future[None]
  553. def wrap(self, server: asyncio.base_events.Server) -> None:
  554. """
  555. Attach to a given :class:`~asyncio.Server`.
  556. Since :meth:`~asyncio.loop.create_server` doesn't support injecting a
  557. custom ``Server`` class, the easiest solution that doesn't rely on
  558. private :mod:`asyncio` APIs is to:
  559. - instantiate a :class:`WebSocketServer`
  560. - give the protocol factory a reference to that instance
  561. - call :meth:`~asyncio.loop.create_server` with the factory
  562. - attach the resulting :class:`~asyncio.Server` with this method
  563. """
  564. self.server = server
  565. for sock in server.sockets:
  566. if sock.family == socket.AF_INET:
  567. name = "%s:%d" % sock.getsockname()
  568. elif sock.family == socket.AF_INET6:
  569. name = "[%s]:%d" % sock.getsockname()[:2]
  570. elif sock.family == socket.AF_UNIX:
  571. name = sock.getsockname()
  572. # In the unlikely event that someone runs websockets over a
  573. # protocol other than IP or Unix sockets, avoid crashing.
  574. else: # pragma: no cover
  575. name = str(sock.getsockname())
  576. self.logger.info("server listening on %s", name)
  577. # Initialized here because we need a reference to the event loop.
  578. # This should be moved back to __init__ when dropping Python < 3.10.
  579. self.closed_waiter = server.get_loop().create_future()
  580. def register(self, protocol: WebSocketServerProtocol) -> None:
  581. """
  582. Register a connection with this server.
  583. """
  584. self.websockets.add(protocol)
  585. def unregister(self, protocol: WebSocketServerProtocol) -> None:
  586. """
  587. Unregister a connection with this server.
  588. """
  589. self.websockets.remove(protocol)
  590. def close(self, close_connections: bool = True) -> None:
  591. """
  592. Close the server.
  593. * Close the underlying :class:`~asyncio.Server`.
  594. * When ``close_connections`` is :obj:`True`, which is the default,
  595. close existing connections. Specifically:
  596. * Reject opening WebSocket connections with an HTTP 503 (service
  597. unavailable) error. This happens when the server accepted the TCP
  598. connection but didn't complete the opening handshake before closing.
  599. * Close open WebSocket connections with close code 1001 (going away).
  600. * Wait until all connection handlers terminate.
  601. :meth:`close` is idempotent.
  602. """
  603. if self.close_task is None:
  604. self.close_task = self.get_loop().create_task(
  605. self._close(close_connections)
  606. )
  607. async def _close(self, close_connections: bool) -> None:
  608. """
  609. Implementation of :meth:`close`.
  610. This calls :meth:`~asyncio.Server.close` on the underlying
  611. :class:`~asyncio.Server` object to stop accepting new connections and
  612. then closes open connections with close code 1001.
  613. """
  614. self.logger.info("server closing")
  615. # Stop accepting new connections.
  616. self.server.close()
  617. # Wait until self.server.close() completes.
  618. await self.server.wait_closed()
  619. # Wait until all accepted connections reach connection_made() and call
  620. # register(). See https://bugs.python.org/issue34852 for details.
  621. await asyncio.sleep(0, **loop_if_py_lt_38(self.get_loop()))
  622. if close_connections:
  623. # Close OPEN connections with status code 1001. Since the server was
  624. # closed, handshake() closes OPENING connections with an HTTP 503
  625. # error. Wait until all connections are closed.
  626. close_tasks = [
  627. asyncio.create_task(websocket.close(1001))
  628. for websocket in self.websockets
  629. if websocket.state is not State.CONNECTING
  630. ]
  631. # asyncio.wait doesn't accept an empty first argument.
  632. if close_tasks:
  633. await asyncio.wait(
  634. close_tasks,
  635. **loop_if_py_lt_38(self.get_loop()),
  636. )
  637. # Wait until all connection handlers are complete.
  638. # asyncio.wait doesn't accept an empty first argument.
  639. if self.websockets:
  640. await asyncio.wait(
  641. [websocket.handler_task for websocket in self.websockets],
  642. **loop_if_py_lt_38(self.get_loop()),
  643. )
  644. # Tell wait_closed() to return.
  645. self.closed_waiter.set_result(None)
  646. self.logger.info("server closed")
  647. async def wait_closed(self) -> None:
  648. """
  649. Wait until the server is closed.
  650. When :meth:`wait_closed` returns, all TCP connections are closed and
  651. all connection handlers have returned.
  652. To ensure a fast shutdown, a connection handler should always be
  653. awaiting at least one of:
  654. * :meth:`~WebSocketServerProtocol.recv`: when the connection is closed,
  655. it raises :exc:`~websockets.exceptions.ConnectionClosedOK`;
  656. * :meth:`~WebSocketServerProtocol.wait_closed`: when the connection is
  657. closed, it returns.
  658. Then the connection handler is immediately notified of the shutdown;
  659. it can clean up and exit.
  660. """
  661. await asyncio.shield(self.closed_waiter)
  662. def get_loop(self) -> asyncio.AbstractEventLoop:
  663. """
  664. See :meth:`asyncio.Server.get_loop`.
  665. """
  666. return self.server.get_loop()
  667. def is_serving(self) -> bool:
  668. """
  669. See :meth:`asyncio.Server.is_serving`.
  670. """
  671. return self.server.is_serving()
  672. async def start_serving(self) -> None: # pragma: no cover
  673. """
  674. See :meth:`asyncio.Server.start_serving`.
  675. Typical use::
  676. server = await serve(..., start_serving=False)
  677. # perform additional setup here...
  678. # ... then start the server
  679. await server.start_serving()
  680. """
  681. await self.server.start_serving()
  682. async def serve_forever(self) -> None: # pragma: no cover
  683. """
  684. See :meth:`asyncio.Server.serve_forever`.
  685. Typical use::
  686. server = await serve(...)
  687. # this coroutine doesn't return
  688. # canceling it stops the server
  689. await server.serve_forever()
  690. This is an alternative to using :func:`serve` as an asynchronous context
  691. manager. Shutdown is triggered by canceling :meth:`serve_forever`
  692. instead of exiting a :func:`serve` context.
  693. """
  694. await self.server.serve_forever()
  695. @property
  696. def sockets(self) -> Iterable[socket.socket]:
  697. """
  698. See :attr:`asyncio.Server.sockets`.
  699. """
  700. return self.server.sockets
  701. async def __aenter__(self) -> WebSocketServer: # pragma: no cover
  702. return self
  703. async def __aexit__(
  704. self,
  705. exc_type: Optional[Type[BaseException]],
  706. exc_value: Optional[BaseException],
  707. traceback: Optional[TracebackType],
  708. ) -> None: # pragma: no cover
  709. self.close()
  710. await self.wait_closed()
  711. class Serve:
  712. """
  713. Start a WebSocket server listening on ``host`` and ``port``.
  714. Whenever a client connects, the server creates a
  715. :class:`WebSocketServerProtocol`, performs the opening handshake, and
  716. delegates to the connection handler, ``ws_handler``.
  717. The handler receives the :class:`WebSocketServerProtocol` and uses it to
  718. send and receive messages.
  719. Once the handler completes, either normally or with an exception, the
  720. server performs the closing handshake and closes the connection.
  721. Awaiting :func:`serve` yields a :class:`WebSocketServer`. This object
  722. provides a :meth:`~WebSocketServer.close` method to shut down the server::
  723. stop = asyncio.Future() # set this future to exit the server
  724. server = await serve(...)
  725. await stop
  726. await server.close()
  727. :func:`serve` can be used as an asynchronous context manager. Then, the
  728. server is shut down automatically when exiting the context::
  729. stop = asyncio.Future() # set this future to exit the server
  730. async with serve(...):
  731. await stop
  732. Args:
  733. ws_handler: Connection handler. It receives the WebSocket connection,
  734. which is a :class:`WebSocketServerProtocol`, in argument.
  735. host: Network interfaces the server binds to.
  736. See :meth:`~asyncio.loop.create_server` for details.
  737. port: TCP port the server listens on.
  738. See :meth:`~asyncio.loop.create_server` for details.
  739. create_protocol: Factory for the :class:`asyncio.Protocol` managing
  740. the connection. It defaults to :class:`WebSocketServerProtocol`.
  741. Set it to a wrapper or a subclass to customize connection handling.
  742. logger: Logger for this server.
  743. It defaults to ``logging.getLogger("websockets.server")``.
  744. See the :doc:`logging guide <../../topics/logging>` for details.
  745. compression: The "permessage-deflate" extension is enabled by default.
  746. Set ``compression`` to :obj:`None` to disable it. See the
  747. :doc:`compression guide <../../topics/compression>` for details.
  748. origins: Acceptable values of the ``Origin`` header, for defending
  749. against Cross-Site WebSocket Hijacking attacks. Include :obj:`None`
  750. in the list if the lack of an origin is acceptable.
  751. extensions: List of supported extensions, in order in which they
  752. should be negotiated and run.
  753. subprotocols: List of supported subprotocols, in order of decreasing
  754. preference.
  755. extra_headers (Union[HeadersLike, Callable[[str, Headers], HeadersLike]]):
  756. Arbitrary HTTP headers to add to the response. This can be
  757. a :data:`~websockets.datastructures.HeadersLike` or a callable
  758. taking the request path and headers in arguments and returning
  759. a :data:`~websockets.datastructures.HeadersLike`.
  760. server_header: Value of the ``Server`` response header.
  761. It defaults to ``"Python/x.y.z websockets/X.Y"``.
  762. Setting it to :obj:`None` removes the header.
  763. process_request (Optional[Callable[[str, Headers], \
  764. Awaitable[Optional[Tuple[http.HTTPStatus, HeadersLike, bytes]]]]]):
  765. Intercept HTTP request before the opening handshake.
  766. See :meth:`~WebSocketServerProtocol.process_request` for details.
  767. select_subprotocol: Select a subprotocol supported by the client.
  768. See :meth:`~WebSocketServerProtocol.select_subprotocol` for details.
  769. open_timeout: Timeout for opening connections in seconds.
  770. :obj:`None` disables the timeout.
  771. See :class:`~websockets.legacy.protocol.WebSocketCommonProtocol` for the
  772. documentation of ``ping_interval``, ``ping_timeout``, ``close_timeout``,
  773. ``max_size``, ``max_queue``, ``read_limit``, and ``write_limit``.
  774. Any other keyword arguments are passed the event loop's
  775. :meth:`~asyncio.loop.create_server` method.
  776. For example:
  777. * You can set ``ssl`` to a :class:`~ssl.SSLContext` to enable TLS.
  778. * You can set ``sock`` to a :obj:`~socket.socket` that you created
  779. outside of websockets.
  780. Returns:
  781. WebSocketServer: WebSocket server.
  782. """
  783. def __init__(
  784. self,
  785. ws_handler: Union[
  786. Callable[[WebSocketServerProtocol], Awaitable[Any]],
  787. Callable[[WebSocketServerProtocol, str], Awaitable[Any]], # deprecated
  788. ],
  789. host: Optional[Union[str, Sequence[str]]] = None,
  790. port: Optional[int] = None,
  791. *,
  792. create_protocol: Optional[Callable[..., WebSocketServerProtocol]] = None,
  793. logger: Optional[LoggerLike] = None,
  794. compression: Optional[str] = "deflate",
  795. origins: Optional[Sequence[Optional[Origin]]] = None,
  796. extensions: Optional[Sequence[ServerExtensionFactory]] = None,
  797. subprotocols: Optional[Sequence[Subprotocol]] = None,
  798. extra_headers: Optional[HeadersLikeOrCallable] = None,
  799. server_header: Optional[str] = USER_AGENT,
  800. process_request: Optional[
  801. Callable[[str, Headers], Awaitable[Optional[HTTPResponse]]]
  802. ] = None,
  803. select_subprotocol: Optional[
  804. Callable[[Sequence[Subprotocol], Sequence[Subprotocol]], Subprotocol]
  805. ] = None,
  806. open_timeout: Optional[float] = 10,
  807. ping_interval: Optional[float] = 20,
  808. ping_timeout: Optional[float] = 20,
  809. close_timeout: Optional[float] = None,
  810. max_size: Optional[int] = 2**20,
  811. max_queue: Optional[int] = 2**5,
  812. read_limit: int = 2**16,
  813. write_limit: int = 2**16,
  814. **kwargs: Any,
  815. ) -> None:
  816. # Backwards compatibility: close_timeout used to be called timeout.
  817. timeout: Optional[float] = kwargs.pop("timeout", None)
  818. if timeout is None:
  819. timeout = 10
  820. else:
  821. warnings.warn("rename timeout to close_timeout", DeprecationWarning)
  822. # If both are specified, timeout is ignored.
  823. if close_timeout is None:
  824. close_timeout = timeout
  825. # Backwards compatibility: create_protocol used to be called klass.
  826. klass: Optional[Type[WebSocketServerProtocol]] = kwargs.pop("klass", None)
  827. if klass is None:
  828. klass = WebSocketServerProtocol
  829. else:
  830. warnings.warn("rename klass to create_protocol", DeprecationWarning)
  831. # If both are specified, klass is ignored.
  832. if create_protocol is None:
  833. create_protocol = klass
  834. # Backwards compatibility: recv() used to return None on closed connections
  835. legacy_recv: bool = kwargs.pop("legacy_recv", False)
  836. # Backwards compatibility: the loop parameter used to be supported.
  837. _loop: Optional[asyncio.AbstractEventLoop] = kwargs.pop("loop", None)
  838. if _loop is None:
  839. loop = asyncio.get_event_loop()
  840. else:
  841. loop = _loop
  842. warnings.warn("remove loop argument", DeprecationWarning)
  843. ws_server = WebSocketServer(logger=logger)
  844. secure = kwargs.get("ssl") is not None
  845. if compression == "deflate":
  846. extensions = enable_server_permessage_deflate(extensions)
  847. elif compression is not None:
  848. raise ValueError(f"unsupported compression: {compression}")
  849. if subprotocols is not None:
  850. validate_subprotocols(subprotocols)
  851. factory = functools.partial(
  852. create_protocol,
  853. # For backwards compatibility with 10.0 or earlier. Done here in
  854. # addition to WebSocketServerProtocol to trigger the deprecation
  855. # warning once per serve() call rather than once per connection.
  856. remove_path_argument(ws_handler),
  857. ws_server,
  858. host=host,
  859. port=port,
  860. secure=secure,
  861. open_timeout=open_timeout,
  862. ping_interval=ping_interval,
  863. ping_timeout=ping_timeout,
  864. close_timeout=close_timeout,
  865. max_size=max_size,
  866. max_queue=max_queue,
  867. read_limit=read_limit,
  868. write_limit=write_limit,
  869. loop=_loop,
  870. legacy_recv=legacy_recv,
  871. origins=origins,
  872. extensions=extensions,
  873. subprotocols=subprotocols,
  874. extra_headers=extra_headers,
  875. server_header=server_header,
  876. process_request=process_request,
  877. select_subprotocol=select_subprotocol,
  878. logger=logger,
  879. )
  880. if kwargs.pop("unix", False):
  881. path: Optional[str] = kwargs.pop("path", None)
  882. # unix_serve(path) must not specify host and port parameters.
  883. assert host is None and port is None
  884. create_server = functools.partial(
  885. loop.create_unix_server, factory, path, **kwargs
  886. )
  887. else:
  888. create_server = functools.partial(
  889. loop.create_server, factory, host, port, **kwargs
  890. )
  891. # This is a coroutine function.
  892. self._create_server = create_server
  893. self.ws_server = ws_server
  894. # async with serve(...)
  895. async def __aenter__(self) -> WebSocketServer:
  896. return await self
  897. async def __aexit__(
  898. self,
  899. exc_type: Optional[Type[BaseException]],
  900. exc_value: Optional[BaseException],
  901. traceback: Optional[TracebackType],
  902. ) -> None:
  903. self.ws_server.close()
  904. await self.ws_server.wait_closed()
  905. # await serve(...)
  906. def __await__(self) -> Generator[Any, None, WebSocketServer]:
  907. # Create a suitable iterator by calling __await__ on a coroutine.
  908. return self.__await_impl__().__await__()
  909. async def __await_impl__(self) -> WebSocketServer:
  910. server = await self._create_server()
  911. self.ws_server.wrap(server)
  912. return self.ws_server
  913. # yield from serve(...) - remove when dropping Python < 3.10
  914. __iter__ = __await__
  915. serve = Serve
  916. def unix_serve(
  917. ws_handler: Union[
  918. Callable[[WebSocketServerProtocol], Awaitable[Any]],
  919. Callable[[WebSocketServerProtocol, str], Awaitable[Any]], # deprecated
  920. ],
  921. path: Optional[str] = None,
  922. **kwargs: Any,
  923. ) -> Serve:
  924. """
  925. Start a WebSocket server listening on a Unix socket.
  926. This function is identical to :func:`serve`, except the ``host`` and
  927. ``port`` arguments are replaced by ``path``. It is only available on Unix.
  928. Unrecognized keyword arguments are passed the event loop's
  929. :meth:`~asyncio.loop.create_unix_server` method.
  930. It's useful for deploying a server behind a reverse proxy such as nginx.
  931. Args:
  932. path: File system path to the Unix socket.
  933. """
  934. return serve(ws_handler, path=path, unix=True, **kwargs)
  935. def remove_path_argument(
  936. ws_handler: Union[
  937. Callable[[WebSocketServerProtocol], Awaitable[Any]],
  938. Callable[[WebSocketServerProtocol, str], Awaitable[Any]],
  939. ]
  940. ) -> Callable[[WebSocketServerProtocol], Awaitable[Any]]:
  941. try:
  942. inspect.signature(ws_handler).bind(None)
  943. except TypeError:
  944. try:
  945. inspect.signature(ws_handler).bind(None, "")
  946. except TypeError: # pragma: no cover
  947. # ws_handler accepts neither one nor two arguments; leave it alone.
  948. pass
  949. else:
  950. # ws_handler accepts two arguments; activate backwards compatibility.
  951. # Enable deprecation warning and announce deprecation in 11.0.
  952. # warnings.warn("remove second argument of ws_handler", DeprecationWarning)
  953. async def _ws_handler(websocket: WebSocketServerProtocol) -> Any:
  954. return await cast(
  955. Callable[[WebSocketServerProtocol, str], Awaitable[Any]],
  956. ws_handler,
  957. )(websocket, websocket.path)
  958. return _ws_handler
  959. return cast(
  960. Callable[[WebSocketServerProtocol], Awaitable[Any]],
  961. ws_handler,
  962. )