router.py 7.3 KB

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