|
- from __future__ import annotations
-
- import http
- import ssl as ssl_module
- import urllib.parse
- from typing import Any, Awaitable, Callable, Literal
-
- from werkzeug.exceptions import NotFound
- from werkzeug.routing import Map, RequestRedirect
-
- from ..http11 import Request, Response
- from .server import Server, ServerConnection, serve
-
-
- __all__ = ["route", "unix_route", "Router"]
-
-
- class Router:
- """WebSocket router supporting :func:`route`."""
-
- def __init__(
- self,
- url_map: Map,
- server_name: str | None = None,
- url_scheme: str = "ws",
- ) -> None:
- self.url_map = url_map
- self.server_name = server_name
- self.url_scheme = url_scheme
- for rule in self.url_map.iter_rules():
- rule.websocket = True
-
- def get_server_name(self, connection: ServerConnection, request: Request) -> str:
- if self.server_name is None:
- return request.headers["Host"]
- else:
- return self.server_name
-
- def redirect(self, connection: ServerConnection, url: str) -> Response:
- response = connection.respond(http.HTTPStatus.FOUND, f"Found at {url}")
- response.headers["Location"] = url
- return response
-
- def not_found(self, connection: ServerConnection) -> Response:
- return connection.respond(http.HTTPStatus.NOT_FOUND, "Not Found")
-
- def route_request(
- self, connection: ServerConnection, request: Request
- ) -> Response | None:
- """Route incoming request."""
- url_map_adapter = self.url_map.bind(
- server_name=self.get_server_name(connection, request),
- url_scheme=self.url_scheme,
- )
- try:
- parsed = urllib.parse.urlparse(request.path)
- handler, kwargs = url_map_adapter.match(
- path_info=parsed.path,
- query_args=parsed.query,
- )
- except RequestRedirect as redirect:
- return self.redirect(connection, redirect.new_url)
- except NotFound:
- return self.not_found(connection)
- connection.handler, connection.handler_kwargs = handler, kwargs
- return None
-
- async def handler(self, connection: ServerConnection) -> None:
- """Handle a connection."""
- return await connection.handler(connection, **connection.handler_kwargs)
-
-
- def route(
- url_map: Map,
- *args: Any,
- server_name: str | None = None,
- ssl: ssl_module.SSLContext | Literal[True] | None = None,
- create_router: type[Router] | None = None,
- **kwargs: Any,
- ) -> Awaitable[Server]:
- """
- Create a WebSocket server dispatching connections to different handlers.
-
- This feature requires the third-party library `werkzeug`_:
-
- .. code-block:: console
-
- $ pip install werkzeug
-
- .. _werkzeug: https://werkzeug.palletsprojects.com/
-
- :func:`route` accepts the same arguments as
- :func:`~websockets.sync.server.serve`, except as described below.
-
- The first argument is a :class:`werkzeug.routing.Map` that maps URL patterns
- to connection handlers. In addition to the connection, handlers receive
- parameters captured in the URL as keyword arguments.
-
- Here's an example::
-
-
- from websockets.asyncio.router import route
- from werkzeug.routing import Map, Rule
-
- async def channel_handler(websocket, channel_id):
- ...
-
- url_map = Map([
- Rule("/channel/<uuid:channel_id>", endpoint=channel_handler),
- ...
- ])
-
- # set this future to exit the server
- stop = asyncio.get_running_loop().create_future()
-
- async with route(url_map, ...) as server:
- await stop
-
-
- Refer to the documentation of :mod:`werkzeug.routing` for details.
-
- If you define redirects with ``Rule(..., redirect_to=...)`` in the URL map,
- when the server runs behind a reverse proxy that modifies the ``Host``
- header or terminates TLS, you need additional configuration:
-
- * Set ``server_name`` to the name of the server as seen by clients. When not
- provided, websockets uses the value of the ``Host`` header.
-
- * Set ``ssl=True`` to generate ``wss://`` URIs without actually enabling
- TLS. Under the hood, this bind the URL map with a ``url_scheme`` of
- ``wss://`` instead of ``ws://``.
-
- There is no need to specify ``websocket=True`` in each rule. It is added
- automatically.
-
- Args:
- url_map: Mapping of URL patterns to connection handlers.
- server_name: Name of the server as seen by clients. If :obj:`None`,
- websockets uses the value of the ``Host`` header.
- ssl: Configuration for enabling TLS on the connection. Set it to
- :obj:`True` if a reverse proxy terminates TLS connections.
- create_router: Factory for the :class:`Router` dispatching requests to
- handlers. Set it to a wrapper or a subclass to customize routing.
-
- """
- url_scheme = "ws" if ssl is None else "wss"
- if ssl is not True and ssl is not None:
- kwargs["ssl"] = ssl
-
- if create_router is None:
- create_router = Router
-
- router = create_router(url_map, server_name, url_scheme)
-
- _process_request: (
- Callable[
- [ServerConnection, Request],
- Awaitable[Response | None] | Response | None,
- ]
- | None
- ) = kwargs.pop("process_request", None)
- if _process_request is None:
- process_request: Callable[
- [ServerConnection, Request],
- Awaitable[Response | None] | Response | None,
- ] = router.route_request
- else:
-
- async def process_request(
- connection: ServerConnection, request: Request
- ) -> Response | None:
- response = _process_request(connection, request)
- if isinstance(response, Awaitable):
- response = await response
- if response is not None:
- return response
- return router.route_request(connection, request)
-
- return serve(router.handler, *args, process_request=process_request, **kwargs)
-
-
- def unix_route(
- url_map: Map,
- path: str | None = None,
- **kwargs: Any,
- ) -> Awaitable[Server]:
- """
- Create a WebSocket Unix server dispatching connections to different handlers.
-
- :func:`unix_route` combines the behaviors of :func:`route` and
- :func:`~websockets.asyncio.server.unix_serve`.
-
- Args:
- url_map: Mapping of URL patterns to connection handlers.
- path: File system path to the Unix socket.
-
- """
- return route(url_map, unix=True, path=path, **kwargs)
|