您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符
 
 
 
 

199 行
6.4 KiB

  1. from __future__ import annotations
  2. import http
  3. import ssl as ssl_module
  4. import urllib.parse
  5. from typing import Any, Awaitable, Callable, Literal
  6. from werkzeug.exceptions import NotFound
  7. from werkzeug.routing import Map, RequestRedirect
  8. from ..http11 import Request, Response
  9. from .server import Server, ServerConnection, serve
  10. __all__ = ["route", "unix_route", "Router"]
  11. class Router:
  12. """WebSocket router supporting :func:`route`."""
  13. def __init__(
  14. self,
  15. url_map: Map,
  16. server_name: str | None = None,
  17. url_scheme: str = "ws",
  18. ) -> None:
  19. self.url_map = url_map
  20. self.server_name = server_name
  21. self.url_scheme = url_scheme
  22. for rule in self.url_map.iter_rules():
  23. rule.websocket = True
  24. def get_server_name(self, connection: ServerConnection, request: Request) -> str:
  25. if self.server_name is None:
  26. return request.headers["Host"]
  27. else:
  28. return self.server_name
  29. def redirect(self, connection: ServerConnection, url: str) -> Response:
  30. response = connection.respond(http.HTTPStatus.FOUND, f"Found at {url}")
  31. response.headers["Location"] = url
  32. return response
  33. def not_found(self, connection: ServerConnection) -> Response:
  34. return connection.respond(http.HTTPStatus.NOT_FOUND, "Not Found")
  35. def route_request(
  36. self, connection: ServerConnection, request: Request
  37. ) -> Response | None:
  38. """Route incoming request."""
  39. url_map_adapter = self.url_map.bind(
  40. server_name=self.get_server_name(connection, request),
  41. url_scheme=self.url_scheme,
  42. )
  43. try:
  44. parsed = urllib.parse.urlparse(request.path)
  45. handler, kwargs = url_map_adapter.match(
  46. path_info=parsed.path,
  47. query_args=parsed.query,
  48. )
  49. except RequestRedirect as redirect:
  50. return self.redirect(connection, redirect.new_url)
  51. except NotFound:
  52. return self.not_found(connection)
  53. connection.handler, connection.handler_kwargs = handler, kwargs
  54. return None
  55. async def handler(self, connection: ServerConnection) -> None:
  56. """Handle a connection."""
  57. return await connection.handler(connection, **connection.handler_kwargs)
  58. def route(
  59. url_map: Map,
  60. *args: Any,
  61. server_name: str | None = None,
  62. ssl: ssl_module.SSLContext | Literal[True] | None = None,
  63. create_router: type[Router] | None = None,
  64. **kwargs: Any,
  65. ) -> Awaitable[Server]:
  66. """
  67. Create a WebSocket server dispatching connections to different handlers.
  68. This feature requires the third-party library `werkzeug`_:
  69. .. code-block:: console
  70. $ pip install werkzeug
  71. .. _werkzeug: https://werkzeug.palletsprojects.com/
  72. :func:`route` accepts the same arguments as
  73. :func:`~websockets.sync.server.serve`, except as described below.
  74. The first argument is a :class:`werkzeug.routing.Map` that maps URL patterns
  75. to connection handlers. In addition to the connection, handlers receive
  76. parameters captured in the URL as keyword arguments.
  77. Here's an example::
  78. from websockets.asyncio.router import route
  79. from werkzeug.routing import Map, Rule
  80. async def channel_handler(websocket, channel_id):
  81. ...
  82. url_map = Map([
  83. Rule("/channel/<uuid:channel_id>", endpoint=channel_handler),
  84. ...
  85. ])
  86. # set this future to exit the server
  87. stop = asyncio.get_running_loop().create_future()
  88. async with route(url_map, ...) as server:
  89. await stop
  90. Refer to the documentation of :mod:`werkzeug.routing` for details.
  91. If you define redirects with ``Rule(..., redirect_to=...)`` in the URL map,
  92. when the server runs behind a reverse proxy that modifies the ``Host``
  93. header or terminates TLS, you need additional configuration:
  94. * Set ``server_name`` to the name of the server as seen by clients. When not
  95. provided, websockets uses the value of the ``Host`` header.
  96. * Set ``ssl=True`` to generate ``wss://`` URIs without actually enabling
  97. TLS. Under the hood, this bind the URL map with a ``url_scheme`` of
  98. ``wss://`` instead of ``ws://``.
  99. There is no need to specify ``websocket=True`` in each rule. It is added
  100. automatically.
  101. Args:
  102. url_map: Mapping of URL patterns to connection handlers.
  103. server_name: Name of the server as seen by clients. If :obj:`None`,
  104. websockets uses the value of the ``Host`` header.
  105. ssl: Configuration for enabling TLS on the connection. Set it to
  106. :obj:`True` if a reverse proxy terminates TLS connections.
  107. create_router: Factory for the :class:`Router` dispatching requests to
  108. handlers. Set it to a wrapper or a subclass to customize routing.
  109. """
  110. url_scheme = "ws" if ssl is None else "wss"
  111. if ssl is not True and ssl is not None:
  112. kwargs["ssl"] = ssl
  113. if create_router is None:
  114. create_router = Router
  115. router = create_router(url_map, server_name, url_scheme)
  116. _process_request: (
  117. Callable[
  118. [ServerConnection, Request],
  119. Awaitable[Response | None] | Response | None,
  120. ]
  121. | None
  122. ) = kwargs.pop("process_request", None)
  123. if _process_request is None:
  124. process_request: Callable[
  125. [ServerConnection, Request],
  126. Awaitable[Response | None] | Response | None,
  127. ] = router.route_request
  128. else:
  129. async def process_request(
  130. connection: ServerConnection, request: Request
  131. ) -> Response | None:
  132. response = _process_request(connection, request)
  133. if isinstance(response, Awaitable):
  134. response = await response
  135. if response is not None:
  136. return response
  137. return router.route_request(connection, request)
  138. return serve(router.handler, *args, process_request=process_request, **kwargs)
  139. def unix_route(
  140. url_map: Map,
  141. path: str | None = None,
  142. **kwargs: Any,
  143. ) -> Awaitable[Server]:
  144. """
  145. Create a WebSocket Unix server dispatching connections to different handlers.
  146. :func:`unix_route` combines the behaviors of :func:`route` and
  147. :func:`~websockets.asyncio.server.unix_serve`.
  148. Args:
  149. url_map: Mapping of URL patterns to connection handlers.
  150. path: File system path to the Unix socket.
  151. """
  152. return route(url_map, unix=True, path=path, **kwargs)