Add a router based on werkzeug.routing.

Fix #311.

docs/conf.py

@@ -82,7 +82,10 @@
 assert PythonDomain.object_types["data"].roles == ("data", "obj")
 PythonDomain.object_types["data"].roles = ("data", "class", "obj")
 
-intersphinx_mapping = {"python": ("https://docs.python.org/3", None)}
+intersphinx_mapping = {
+    "python": ("https://docs.python.org/3", None),
+    "werkzeug": ("https://werkzeug.palletsprojects.com/en/stable/", None),
+}
 
 spelling_show_suggestions = True
 

docs/project/changelog.rst

@@ -56,6 +56,12 @@ Backwards-incompatible changes
 
     See :doc:`keepalive and latency <../topics/keepalive>` for details.
 
+New features
+............
+
+* Added :func:`~asyncio.router.route` to dispatch connections to different
+  handlers depending on the URL.
+
 Improvements
 ............
 

docs/reference/asyncio/server.rst

@@ -12,6 +12,21 @@ Creating a server
 .. autofunction:: unix_serve
     :async:
 
+Routing connections
+-------------------
+
+.. automodule:: websockets.asyncio.router
+
+.. autofunction:: route
+    :async:
+
+.. autofunction:: unix_route
+    :async:
+
+.. autoclass:: Router
+
+.. currentmodule:: websockets.asyncio.server
+
 Running a server
 ----------------
 
@@ -89,12 +104,12 @@ Using a connection
 Broadcast
 ---------
 
-.. autofunction:: websockets.asyncio.server.broadcast
+.. autofunction:: broadcast
 
 HTTP Basic Authentication
 -------------------------
 
 websockets supports HTTP Basic Authentication according to
 :rfc:`7235` and :rfc:`7617`.
 
-.. autofunction:: websockets.asyncio.server.basic_auth
+.. autofunction:: basic_auth

docs/reference/features.rst

@@ -127,6 +127,8 @@ Server
     +------------------------------------+--------+--------+--------+--------+
     | Perform HTTP Digest Authentication | ❌     | ❌     | ❌     | ❌     |
     +------------------------------------+--------+--------+--------+--------+
+    | Dispatch connections to handlers   | ✅     | ✅     | —      | ❌     |
+    +------------------------------------+--------+--------+--------+--------+
 
 Client
 ------

docs/reference/sync/server.rst

@@ -10,6 +10,33 @@ Creating a server
 
 .. autofunction:: unix_serve
 
+Routing connections
+-------------------
+
+.. automodule:: websockets.sync.router
+
+.. autofunction:: route
+    :async:
+
+.. autofunction:: unix_route
+    :async:
+
+.. autoclass:: Router
+
+.. currentmodule:: websockets.sync.server
+
+Routing connections
+-------------------
+
+.. autofunction:: route
+    :async:
+
+.. autofunction:: unix_route
+    :async:
+
+.. autoclass:: Server
+
+
 Running a server
 ----------------
 
@@ -78,4 +105,4 @@ HTTP Basic Authentication
 websockets supports HTTP Basic Authentication according to
 :rfc:`7235` and :rfc:`7617`.
 
-.. autofunction:: websockets.sync.server.basic_auth
+.. autofunction:: basic_auth

docs/requirements.txt

@@ -6,3 +6,4 @@ sphinx-inline-tabs
 sphinxcontrib-spelling
 sphinxcontrib-trio
 sphinxext-opengraph
+werkzeug

src/websockets/__init__.py

@@ -12,6 +12,10 @@
     "connect",
     "unix_connect",
     "ClientConnection",
+    # .asyncio.router
+    "route",
+    "unix_route",
+    "Router",
     # .asyncio.server
     "basic_auth",
     "broadcast",
@@ -79,6 +83,7 @@
 # When type checking, import non-deprecated aliases eagerly. Else, import on demand.
 if TYPE_CHECKING:
     from .asyncio.client import ClientConnection, connect, unix_connect
+    from .asyncio.router import Router, route, unix_route
     from .asyncio.server import (
         Server,
         ServerConnection,
@@ -138,6 +143,10 @@
             "connect": ".asyncio.client",
             "unix_connect": ".asyncio.client",
             "ClientConnection": ".asyncio.client",
+            # .asyncio.router
+            "route": ".asyncio.router",
+            "unix_route": ".asyncio.router",
+            "Router": ".asyncio.router",
             # .asyncio.server
             "basic_auth": ".asyncio.server",
             "broadcast": ".asyncio.server",

src/websockets/asyncio/router.py

@@ -0,0 +1,186 @@
+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 with several handlers.
+
+    Except for the differences described below, this function accepts the same
+    arguments as :func:`~websockets.sync.server.serve`.
+
+    The first argument is a :class:`werkzeug.routing.Map` mapping URL patterns
+    to connection handlers, instead of a single connection handler::
+
+        from websockets.sync.router import route
+        from werkzeug.routing import Map, Rule
+
+        url_map = Map([
+            Rule("/", endpoint=default_handler),
+            ...
+        ])
+
+        with router(url_map, ...) as server:
+            server.serve_forever()
+
+    Handlers are called with the connection and any keyword arguments captured
+    in the URL.
+
+    There is no need to specify ``websocket=True`` in ``url_map``. It is added
+    to each rule automatically.
+
+    This feature requires the third-party library `werkzeug`_::
+
+        $ pip install werkzeug
+
+    .. _werkzeug: https://werkzeug.palletsprojects.com/
+
+    If you define redirects with ``Rule(..., redirect_to=...)`` in the URL map
+    and the server runs behind a reverse proxy that modifies the ``Host`` header
+    or terminates TLS, you need the following 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://``.
+
+    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 with several handlers.
+
+    This function combines behaviors of :func:`~websockets.sync.router.route`
+    and :func:`~websockets.sync.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)

src/websockets/asyncio/server.py

@@ -9,7 +9,7 @@
 import sys
 from collections.abc import Awaitable, Generator, Iterable, Sequence
 from types import TracebackType
-from typing import Any, Callable, cast
+from typing import Any, Callable, Mapping, cast
 
 from ..exceptions import InvalidHeader
 from ..extensions.base import ServerExtensionFactory
@@ -87,6 +87,8 @@ def __init__(
         self.server = server
         self.request_rcvd: asyncio.Future[None] = self.loop.create_future()
         self.username: str  # see basic_auth()
+        self.handler: Callable[[ServerConnection], Awaitable[None]]  # see route()
+        self.handler_kwargs: Mapping[str, Any]  # see route()
 
     def respond(self, status: StatusLike, text: str) -> Response:
         """