Use X-Forwarded-Proto for WebSockets scheme when the proxy provides…

… it (#2258)

* Fix X-Forwarded-Proto when the proxy already sets it to "ws" or "wss"

Minor fix for #2043

Traefik already sets the X-Forwarded-Proto headers to ws or wss for websockets. https://github.com/traefik/traefik/blob/c1ef7429771104e79f2e87b236b21495cb5765f0/pkg/middlewares/forwardedheaders/forwarded_header.go#L150

This change should make sure we don't overwrite those values.

* Fix the logic

* Update test_proxy_headers.py

Test whether passing "wss" in X-Forwarded-Proto works

* Simplify the logic

(probably more ways to write this... lmk which you prefer)

* Update tests and min implementation

* Remove new line

---------

Co-authored-by: Marcelo Trylesinski <[email protected]>

tests/middleware/test_proxy_headers.py

@@ -115,7 +115,18 @@ async def test_proxy_headers_invalid_x_forwarded_for() -> None:
 
 
 @pytest.mark.anyio
+@pytest.mark.parametrize(
+    "x_forwarded_proto,addr",
+    [
+        ("http", "ws://1.2.3.4:0"),
+        ("https", "wss://1.2.3.4:0"),
+        ("ws", "ws://1.2.3.4:0"),
+        ("wss", "wss://1.2.3.4:0"),
+    ],
+)
 async def test_proxy_headers_websocket_x_forwarded_proto(
+    x_forwarded_proto: str,
+    addr: str,
     ws_protocol_cls: "Type[WSProtocol | WebSocketProtocol]",
     http_protocol_cls: "Type[H11Protocol | HttpToolsProtocol]",
     unused_tcp_port: int,
@@ -138,7 +149,7 @@ async def websocket_app(scope, receive, send):
 
     async with run_server(config):
         url = f"ws://127.0.0.1:{unused_tcp_port}"
-        headers = {"X-Forwarded-Proto": "https", "X-Forwarded-For": "1.2.3.4"}
+        headers = {"X-Forwarded-Proto": x_forwarded_proto, "X-Forwarded-For": "1.2.3.4"}
         async with websockets.client.connect(url, extra_headers=headers) as websocket:
             data = await websocket.recv()
-            assert data == "wss://1.2.3.4:0"
+            assert data == addr

uvicorn/middleware/proxy_headers.py

@@ -63,9 +63,7 @@ async def __call__(
                         headers[b"x-forwarded-proto"].decode("latin1").strip()
                     )
                     if scope["type"] == "websocket":
-                        scope["scheme"] = (
-                            "wss" if x_forwarded_proto == "https" else "ws"
-                        )
+                        scope["scheme"] = x_forwarded_proto.replace("http", "ws")
                     else:
                         scope["scheme"] = x_forwarded_proto