Repository: falconry/falcon

Issue:
WebSocket sinks
It seems that using sinks to match a WebSocket request is currently undocumented, and it leads to sloppy behaviour/mismatched typing, although it surprisingly works, in a sense. But the sink then receives a `ws` object instead of `resp`, so a sink method that is unaware of the possibility will likely just error out when trying to set unexpected attributes on `ws`.

A decision is needed how we want this to work, and what can be considered a fix, and not a breaking change.
OTOH, we have clearly documented the sink callable in `add_sink()` as taking `req` and `resp` parameters, so I think we can afford some redesign here. Alternatively, we could add a flag governing this behaviour.

One way to architect this is to add a `ws` parameter along the lines of [`App.add_error_handler()`](https://falcon.readthedocs.io/en/stable/api/app.html#falcon.asgi.App.add_error_handler). However, sinks can accept arbitrary `**kwargs`...


Discussion:

Patch:
diff --git a/docs/_newsfragments/2414.bugfix.rst b/docs/_newsfragments/2414.bugfix.rst
new file mode 100644
index 000000000..757042ef4
--- /dev/null
+++ b/docs/_newsfragments/2414.bugfix.rst
@@ -0,0 +1,15 @@
+Previously, Falcon's :ref:`WebSocket implementation <ws>` was not documented to
+route the request to any :meth:`sink <falcon.asgi.App.add_sink>`. However, in
+the case of a missing route, a matching sink was actually invoked, passing
+:class:`ws <falcon.asgi.WebSocket>` in place of the incompatible
+:class:`resp <falcon.asgi.Response>`.
+
+This mismatch has been addressed by introducing a `ws` keyword argument
+(similar to ASGI :meth:`error handlers <falcon.asgi.App.add_error_handler>`)
+for sink functions meant to accept WebSocket connections.
+
+For backwards-compatibility, when `ws` is absent from the sink's signature, the
+:class:`~falcon.asgi.WebSocket` object is still passed in place of the
+incompatible `resp`.
+This behavior will change in Falcon 5.0: when draining a WebSocket connection,
+`resp` will always be set to ``None``.
diff --git a/docs/api/websocket.rst b/docs/api/websocket.rst
index 6e632c521..2fbc9cb3e 100644
--- a/docs/api/websocket.rst
+++ b/docs/api/websocket.rst
@@ -90,11 +90,31 @@ framework would expect an ``on_websocket()`` responder similar to this:
 
 .. code:: python
 
-    async def on_websocket(self, req: Request, ws: WebSocket, account_id: str):
+    async def on_websocket(self, req: Request, ws: WebSocket, account_id: str) -> None:
         pass
 
-If no route matches the path requested in the WebSocket handshake, control then
-passes to a default responder that simply raises an instance of
+Just like other HTTP requests, WebSocket connections can also be handled
+dynamically by :meth:`sinks <falcon.asgi.App.add_sink>`. In order to receive
+a :class:`~falcon.asgi.WebSocket` connection object, the sink method ought to
+define a `ws` keyword argument:
+
+.. code:: python
+
+    async def sink(
+        req: Request,
+        resp: Response | None,
+        ws: WebSocket | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if ws is not None:
+            # WebSocket connection (resp is None)
+            ...
+        else:
+            # Ordinary HTTP request (ws is None)
+            ...
+
+If no route or sink matches the path requested in the WebSocket handshake,
+control then passes to a default responder that simply raises an instance of
 :class:`~.HTTPRouteNotFound`. By default, this error will be rendered as a 403
 response with a 3404 close code. This behavior can be modified by adding a
 custom error handler (see also: :meth:`~falcon.asgi.App.add_error_handler`).
diff --git a/falcon/_typing.py b/falcon/_typing.py
index 6145b3c82..2adb466be 100644
--- a/falcon/_typing.py
+++ b/falcon/_typing.py
@@ -99,7 +99,7 @@ def __call__(
 
 class AsgiSinkCallable(Protocol):
     async def __call__(
-        self, req: AsgiRequest, resp: AsgiResponse, **kwargs: Optional[str]
+        self, req: AsgiRequest, resp: Optional[AsgiResponse], **kwargs: Optional[str]
     ) -> None: ...
 
 
diff --git a/falcon/app.py b/falcon/app.py
index 3e89765f5..c01913d98 100644
--- a/falcon/app.py
+++ b/falcon/app.py
@@ -764,7 +764,30 @@ def add_sink(self, sink: SinkCallable, prefix: SinkPrefix = r'/') -> None:
 
                 Note:
                     When using an async version of the ``App``, this must be a
-                    coroutine.
+                    coroutine function taking the form
+                    ``func(req, resp, ws=None, **kwargs)``.
+
+                    Similar to
+                    :meth:`error handlers <falcon.asgi.App.add_error_handler>`,
+                    in the case of a WebSocket connection, the
+                    :class:`resp <falcon.asgi.Response>` argument will be
+                    ``None``, whereas the `ws` keyword argument will receive
+                    the :class:`~falcon.asgi.WebSocket` connection object.
+
+                    For backwards-compatibility, when `ws` is absent from the
+                    sink's signature, or a regex match (see **prefix** below)
+                    contains a group named 'ws', the
+                    :class:`~falcon.asgi.WebSocket` object is passed in place
+                    of the incompatible `resp`.
+
+                    This behavior will change in Falcon 5.0: when draining a
+                    WebSocket connection, `resp` will always be set to ``None``
+                    regardless of the sink's signature.
+
+                .. versionadded:: 4.1
+                    If an asynchronous sink callable explicitly defines a `ws`
+                    argument, it is used to pass the
+                    :class:`~falcon.asgi.WebSocket` connection object.
 
             prefix (str): A regex string, typically starting with '/', which
                 will trigger the sink if it matches the path portion of the
diff --git a/falcon/asgi/app.py b/falcon/asgi/app.py
index 651e50723..ad5ddf4d6 100644
--- a/falcon/asgi/app.py
+++ b/falcon/asgi/app.py
@@ -36,6 +36,7 @@
     TypeVar,
     Union,
 )
+import warnings
 
 from falcon import constants
 from falcon import responders
@@ -64,7 +65,8 @@
 from falcon.http_error import HTTPError
 from falcon.http_status import HTTPStatus
 from falcon.media.multipart import MultipartFormHandler
-from falcon.util import get_argnames
+from falcon.util.deprecation import DeprecatedWarning
+from falcon.util.misc import _has_arg_name
 from falcon.util.misc import is_python_func
 from falcon.util.sync import _should_wrap_non_coroutines
 from falcon.util.sync import _wrap_non_coroutine_unsafe
@@ -1160,7 +1162,41 @@ async def _handle_websocket(
                 for process_resource_ws in resource_mw:
                     await process_resource_ws(req, web_socket, resource, params)
 
-            await on_websocket(req, web_socket, **params)
+                await on_websocket(req, web_socket, **params)
+
+            else:
+                # NOTE(vytas): The request did not match any route;
+                #   on_websocket is either a sink or a default responder.
+                #   Check whether it has a ws argument, and we are not
+                #   shadowing a sink parameter from regex; otherwise pass ws
+                #   instead of resp for backwards compatibility.
+                # TODO(vytas): Always pass resp=None in Falcon 5.0
+                #   (breaking change).
+                if _has_arg_name(on_websocket, 'ws') and 'ws' not in params:
+                    params['ws'] = web_socket
+                    await on_websocket(req, None, **params)  # type: ignore[arg-type]
+                else:
+                    cls = type(self)
+                    # NOTE(vytas): We could add the ws=None parameter to the
+                    #   default responders, but let's keep it simple for now,
+                    #   and skip the warning part.
+                    if on_websocket not in (
+                        cls._default_responder_bad_request,
+                        cls._default_responder_path_not_found,
+                    ):
+                        warnings.warn(
+                            f'Please define the sink coroutine function '
+                            f'{on_websocket.__qualname__} as '  # type: ignore[attr-defined]
+                            f'{on_websocket.__name__}'  # type: ignore[attr-defined]
+                            f'(req, resp, ws=None, **kwargs) if '
+                            f'it is intended to handle WebSocket requests; receiving '
+                            f'ws in place of resp is deprecated, and in Falcon 5.0+, '
+                            f'resp will be set to None when invoking WebSocket sinks.',
+                            category=DeprecatedWarning,
+                        )
+
+                    await on_websocket(req, web_socket, **params)
+
             await web_socket.close()
 
         except Exception as ex:
@@ -1298,7 +1334,10 @@ async def _handle_exception(  # type: ignore[override]
             try:
                 kwargs = {}
 
-                if ws and 'ws' in get_argnames(err_handler):
+                # PERF(vytas): Using the LRU-cache backed misc._has_arg_name
+                #   here as inspect.signature (and by proxy misc.get_argnames)
+                #   can be very slow (on the order of magnitude of 10 µs).
+                if ws is not None and _has_arg_name(err_handler, 'ws'):
                     kwargs['ws'] = ws
 
                 await err_handler(req, resp, ex, params, **kwargs)
diff --git a/falcon/routing/static.py b/falcon/routing/static.py
index bc7096e96..edb733075 100644
--- a/falcon/routing/static.py
+++ b/falcon/routing/static.py
@@ -307,7 +307,16 @@ def __call__(self, req: Request, resp: Response, **kw: Any) -> None:
 class StaticRouteAsync(StaticRoute):
     """Subclass of StaticRoute with modifications to support ASGI apps."""
 
-    async def __call__(self, req: asgi.Request, resp: asgi.Response, **kw: Any) -> None:  # type: ignore[override]
+    async def __call__(  # type: ignore[override]
+        self,
+        req: asgi.Request,
+        resp: asgi.Response,
+        ws: Optional[asgi.WebSocket] = None,
+        **kw: Any,
+    ) -> None:
+        if ws is not None:
+            raise falcon.HTTPBadRequest()
+
         super().__call__(req, resp, **kw)
         if resp.stream is not None:  # None when in an option request
             # NOTE(kgriffs): Fixup resp.stream so that it is non-blocking
diff --git a/falcon/util/misc.py b/falcon/util/misc.py
index cd68c1de1..ffa1c8cc5 100644
--- a/falcon/util/misc.py
+++ b/falcon/util/misc.py
@@ -337,6 +337,20 @@ def get_argnames(func: Callable[..., Any]) -> List[str]:
     return args
 
 
+@functools.lru_cache
+def _has_arg_name_cached(func: Callable[..., Any], name: str) -> bool:
+    return name in get_argnames(func)
+
+
+def _has_arg_name(func: Callable[..., Any], name: str) -> bool:
+    try:
+        return _has_arg_name_cached(func, name)
+    except TypeError:
+        # NOTE(vytas): Most probably the exception was thrown by the LRU cache
+        #   indicating that the func object was unhashable.
+        return name in get_argnames(func)
+
+
 def secure_filename(filename: str, max_length: Optional[int] = None) -> str:
     """Sanitize the provided `filename` to contain only ASCII characters.
 


Test patch:
diff --git a/tests/asgi/test_ws.py b/tests/asgi/test_ws.py
index 49510c2f0..b68be9178 100644
--- a/tests/asgi/test_ws.py
+++ b/tests/asgi/test_ws.py
@@ -1,5 +1,6 @@
 import asyncio
 from collections import deque
+from contextlib import nullcontext
 import functools
 import os
 
@@ -13,6 +14,7 @@
 from falcon.asgi.ws import WebSocket
 from falcon.asgi.ws import WebSocketOptions
 from falcon.testing.helpers import _WebSocketState as ClientWebSocketState
+from falcon.util.deprecation import DeprecatedWarning
 
 try:
     import cbor2
@@ -1457,3 +1459,64 @@ async def on_websocket(self, req, ws, user):
                 received.append(await ws.receive_text())
 
     assert len(received) == 6
+
+
+@pytest.mark.parametrize('sink', ['no-ws-arg', 'with-ws-arg', 'unhashable-obj'])
+async def test_sinks(conductor, sink):
+    async def say_hello(req, resp_or_ws, name):
+        greeting = f'Hello, {name}!'
+        if req.is_websocket:
+            await resp_or_ws.accept()
+            await resp_or_ws.send_text(greeting)
+
+            await resp_or_ws.receive_text()  # Receive /bye
+        else:
+            resp_or_ws.content_type = falcon.MEDIA_TEXT
+            resp_or_ws.text = greeting
+
+    async def sink_no_ws_arg(req, resp_or_ws, name):
+        assert resp_or_ws is not None
+        await say_hello(req, resp_or_ws, name)
+
+    async def sink_with_ws_arg(req, resp, name, ws=None):
+        if ws:
+            assert resp is None
+            await say_hello(req, ws, name)
+        else:
+            assert resp is not None
+            await say_hello(req, resp, name)
+
+    class Unhashable(list):
+        async def __call__(self, req, resp, name, ws=None):
+            await sink_with_ws_arg(req, resp, name, ws=ws)
+
+    sinks = {
+        'no-ws-arg': sink_no_ws_arg,
+        'with-ws-arg': sink_with_ws_arg,
+        'unhashable-obj': Unhashable(),
+    }
+    conductor.app.add_sink(sinks[sink], r'/ws/(?P<name>\w+)')
+
+    async with conductor as c:
+        resp = await c.simulate_get('/ws/World')
+        assert resp.status_code == 200
+        assert resp.text == 'Hello, World!'
+
+        expectation = (
+            pytest.warns(DeprecatedWarning) if sink == 'no-ws-arg' else nullcontext()
+        )
+        with expectation:
+            async with c.simulate_ws('/ws/World') as ws:
+                assert await ws.receive_text() == 'Hello, World!'
+                await ws.send_text('/bye')
+
+
+async def test_static_route(conductor, tmp_path):
+    conductor.app.add_static_route('/static', tmp_path)
+
+    async with conductor as c:
+        with pytest.raises(falcon.WebSocketDisconnected) as exc_info:
+            async with c.simulate_ws('/static/report.json'):
+                pass
+
+        assert exc_info.value.code == 3400
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 28e708eea..84910b857 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1342,6 +1342,33 @@ def __call__(self, a, b):
     assert misc.get_argnames(functools.partial(foo, 42)) == ['b', 'c']
 
 
+def test_has_arg_name():
+    def foo(a, b, c):
+        pass
+
+    class Bar:
+        def __call__(self, a, b):
+            pass
+
+    class Baz(list):
+        def __call__(self, a, b):
+            pass
+
+    assert misc._has_arg_name(foo, 'a')
+    assert misc._has_arg_name(foo, 'b')
+    assert misc._has_arg_name(foo, 'c')
+
+    bar = Bar()
+    assert misc._has_arg_name(bar, 'a')
+    assert misc._has_arg_name(bar, 'b')
+    assert not misc._has_arg_name(bar, 'c')
+
+    baz = Baz()
+    assert misc._has_arg_name(baz, 'a')
+    assert misc._has_arg_name(baz, 'b')
+    assert not misc._has_arg_name(baz, 'c')
+
+
 class TestContextType:
     class CustomContextType(structures.Context):
         def __init__(self):