websocket handling

CHANGES.rst

@@ -165,6 +165,12 @@ Unreleased
 -   ``request.values`` does not include ``form`` for GET requests (even
     though GET bodies are undefined). This prevents bad caching proxies
     from caching form data instead of query strings. :pr:`2037`
+-   The development server adds the underlying socket to ``environ`` as
+    ``werkzeug.socket``. This is non-standard and specific to the dev
+    server, other servers may expose this under their own key. It is
+    useful for handling a WebSocket upgrade request. :issue:`2052`
+-   URL matching assumes ``websocket=True`` mode for WebSocket upgrade
+    requests. :issue:`2052`
 
 
 Version 1.0.2

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)

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:

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": "",

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(