Repository: pallets/werkzeug

Issue:
WebSocket routing and handling
While there's been some effort in adding websocket routing support, there are still a couple of things that are missing:

- When the server receives a ws:// or wss:// request, it still creates a `MapAdapter()` instance that is configured to handle HTTP, so the only way to match a WebSocket route is to manually pass `websocket=True` in the `match()` call. If the incoming request has the `Connection: Upgrade` and `Upgrade: WebSocket` headers in it, then it is a WebSocket request, so the MapAdapter should be initialized to match WebSocket by default.
- To be able to implement WebSocket, the route handler function needs to take control of the underlying socket, so that it can do the WebSocket handshake, send the 101 response and then start sending/receiving WebSocket frames. Gunicorn handles this by adding the `gunicorn.socket` key to the `environ` dictionary. Werkzeug could add a `werkzeug.socket` key to expose its socket.

Adding solutions to the above problems would allow a Werkzeug or Flask application to implement a WebSocket route, for example using [wsproto](https://python-hyper.org/projects/wsproto/en/stable/).


Discussion:

Resolving patch:
diff --git a/examples/wsecho.py b/examples/wsecho.py
new file mode 100644
--- /dev/null
+++ b/examples/wsecho.py
@@ -0,0 +1,79 @@
+"""Shows how you can implement a simple WebSocket echo server using the
+wsproto library.
+"""
+from werkzeug.exceptions import InternalServerError
+from werkzeug.serving import run_simple
+from werkzeug.wrappers import Request
+from werkzeug.wrappers import Response
+from wsproto import ConnectionType
+from wsproto import WSConnection
+from wsproto.events import AcceptConnection
+from wsproto.events import CloseConnection
+from wsproto.events import Message
+from wsproto.events import Ping
+from wsproto.events import Request as WSRequest
+from wsproto.events import TextMessage
+from wsproto.frame_protocol import CloseReason
+
+
+@Request.application
+def websocket(request):
+    # The underlying socket must be provided by the server. Gunicorn and
+    # Werkzeug's dev server are known to support this.
+    stream = request.environ.get("werkzeug.socket")
+
+    if stream is None:
+        stream = request.environ.get("gunicorn.socket")
+
+    if stream is None:
+        raise InternalServerError()
+
+    # Initialize the wsproto connection. Need to recreate the request
+    # data that was read by the WSGI server already.
+    ws = WSConnection(ConnectionType.SERVER)
+    in_data = b"GET %s HTTP/1.1\r\n" % request.path.encode("utf8")
+
+    for header, value in request.headers.items():
+        in_data += f"{header}: {value}\r\n".encode("utf8")
+
+    in_data += b"\r\n"
+    ws.receive_data(in_data)
+    running = True
+
+    while True:
+        out_data = b""
+
+        for event in ws.events():
+            if isinstance(event, WSRequest):
+                out_data += ws.send(AcceptConnection())
+            elif isinstance(event, CloseConnection):
+                out_data += ws.send(event.response())
+                running = False
+            elif isinstance(event, Ping):
+                out_data += ws.send(event.response())
+            elif isinstance(event, TextMessage):
+                # echo the incoming message back to the client
+                if event.data == "quit":
+                    out_data += ws.send(
+                        CloseConnection(CloseReason.NORMAL_CLOSURE, "bye")
+                    )
+                    running = False
+                else:
+                    out_data += ws.send(Message(data=event.data))
+
+        if out_data:
+            stream.send(out_data)
+
+        if not running:
+            break
+
+        in_data = stream.recv(4096)
+        ws.receive_data(in_data)
+
+    # The connection will be closed at this point, but WSGI still
+    # requires a response.
+    return Response("", status=204)
+
+
+if __name__ == "__main__":
+    run_simple("localhost", 5000, websocket)
diff --git a/src/werkzeug/routing.py b/src/werkzeug/routing.py
--- a/src/werkzeug/routing.py
+++ b/src/werkzeug/routing.py
@@ -1634,15 +1634,21 @@ def bind_to_environ(
         wsgi_server_name = get_host(environ).lower()
         scheme = environ["wsgi.url_scheme"]
 
+        if (
+            environ.get("HTTP_CONNECTION", "").lower() == "upgrade"
+            and environ.get("HTTP_UPGRADE", "").lower() == "websocket"
+        ):
+            scheme = "wss" if scheme == "https" else "ws"
+
         if server_name is None:
             server_name = wsgi_server_name
         else:
             server_name = server_name.lower()
 
             # strip standard port to match get_host()
-            if scheme == "http" and server_name.endswith(":80"):
+            if scheme in {"http", "ws"} and server_name.endswith(":80"):
                 server_name = server_name[:-3]
-            elif scheme == "https" and server_name.endswith(":443"):
+            elif scheme in {"https", "wss"} and server_name.endswith(":443"):
                 server_name = server_name[:-4]
 
         if subdomain is None and not self.host_matching:
diff --git a/src/werkzeug/serving.py b/src/werkzeug/serving.py
--- a/src/werkzeug/serving.py
+++ b/src/werkzeug/serving.py
@@ -184,6 +184,7 @@ def shutdown_server():
             "wsgi.multiprocess": self.server.multiprocess,
             "wsgi.run_once": False,
             "werkzeug.server.shutdown": shutdown_server,
+            "werkzeug.socket": self.connection,
             "SERVER_SOFTWARE": self.server_version,
             "REQUEST_METHOD": self.command,
             "SCRIPT_NAME": "",


Test patch:
diff --git a/tests/test_routing.py b/tests/test_routing.py
--- a/tests/test_routing.py
+++ b/tests/test_routing.py
@@ -62,6 +62,17 @@ def test_basic_routing():
     with pytest.raises(r.WebsocketMismatch):
         adapter.match("/foo", websocket=True)
 
+    adapter = map.bind_to_environ(
+        create_environ(
+            "/ws?foo=bar",
+            "http://example.org/",
+            headers=[("Connection", "Upgrade"), ("upgrade", "WebSocket")],
+        )
+    )
+    assert adapter.match("/ws") == ("ws", {})
+    with pytest.raises(r.WebsocketMismatch):
+        adapter.match("/ws", websocket=False)
+
 
 def test_merge_slashes_match():
     url_map = r.Map(