Add default headers to WebSockets implementations

docs/settings.md

@@ -83,7 +83,7 @@ connecting IPs in the `forwarded-allow-ips` configuration.
 * `--date-header` / `--no-date-header` - Enable/Disable default `Date` header.
 
 !!! note
-    The `--no-server-header` flag doesn't have effect on the WebSockets implementations.
+    The `--no-server-header` and `--no-date-header` flags doesn't have effect on the WebSockets implementations.
 
 ## HTTPS
 
@@ -103,4 +103,4 @@ connecting IPs in the `forwarded-allow-ips` configuration.
 
 ## Timeouts
 
-* `--timeout-keep-alive <int>` - Close Keep-Alive connections if no new data is received within this timeout. **Default:** *5*.
+* `--timeout-keep-alive <int>` - Close Keep-Alive connections if no new data is received within this timeout. **Default:** *5*.

pyproject.toml

@@ -33,15 +33,17 @@ dependencies = [
     "typing-extensions;python_version < '3.8'",
 ]
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [project.optional-dependencies]
 standard = [
     "colorama>=0.4;sys_platform == 'win32'",
     "httptools>=0.4.0",
     "python-dotenv>=0.13",
     "PyYAML>=5.1",
     "uvloop>=0.14.0,!=0.15.0,!=0.15.1; sys_platform != 'win32' and (sys_platform != 'cygwin' and platform_python_implementation != 'PyPy')",
-    "watchfiles>=0.13",
-    "websockets>=10.0",
+    "watchfiles>=0.13"
 ]
 
 [project.scripts]
@@ -60,3 +62,4 @@ path = "uvicorn/__init__.py"
 include = [
     "/uvicorn",
 ]
+

requirements.txt

@@ -4,8 +4,8 @@
 asgiref==3.5.2
 
 # Explicit optionals
-wsproto==1.1.0
-
+wsproto==1.2.0
+git+https://github.com/aaugustin/websockets.git@main
 # Packaging
 build==0.8.0
 twine==4.0.1

tests/protocols/test_websocket.py

@@ -10,6 +10,7 @@
 
 try:
     import websockets
+    import websockets.exceptions
     from websockets.extensions.permessage_deflate import ClientPerMessageDeflateFactory
 
     from uvicorn.protocols.websockets.websockets_impl import WebSocketProtocol
@@ -18,9 +19,9 @@
     WebSocketProtocol = None
     ClientPerMessageDeflateFactory = None
 
-
 ONLY_WEBSOCKETPROTOCOL = [p for p in [WebSocketProtocol] if p is not None]
 WS_PROTOCOLS = [p for p in [WSProtocol, WebSocketProtocol] if p is not None]
+ONLY_WS_PROTOCOL = [p for p in [WSProtocol] if p is not None]
 pytestmark = pytest.mark.skipif(
     websockets is None, reason="This test needs the websockets module"
 )
@@ -658,3 +659,102 @@ async def send_text(url):
         await send_text("ws://127.0.0.1:8000")
 
     assert frames == [b"abc", b"abc", b"abc"]
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("ws_protocol_cls", WS_PROTOCOLS)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_default_server_headers(ws_protocol_cls, http_protocol_cls):
+    class App(WebSocketResponse):
+        async def websocket_connect(self, message):
+            await self.send({"type": "websocket.accept"})
+
+    async def open_connection(url):
+        async with websockets.connect(url) as websocket:
+            return websocket.response_headers
+
+    config = Config(app=App, ws=ws_protocol_cls, http=http_protocol_cls, lifespan="off")
+    async with run_server(config):
+        headers = await open_connection("ws://127.0.0.1:8000")
+        assert headers.get("server") == "uvicorn" and "date" in headers
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("ws_protocol_cls", WS_PROTOCOLS)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_no_server_headers(ws_protocol_cls, http_protocol_cls):
+    class App(WebSocketResponse):
+        async def websocket_connect(self, message):
+            await self.send({"type": "websocket.accept"})
+
+    async def open_connection(url):
+        async with websockets.connect(url) as websocket:
+            return websocket.response_headers
+
+    config = Config(
+        app=App,
+        ws=ws_protocol_cls,
+        http=http_protocol_cls,
+        lifespan="off",
+        server_header=False,
+    )
+    async with run_server(config):
+        headers = await open_connection("ws://127.0.0.1:8000")
+        assert "server" not in headers
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("ws_protocol_cls", ONLY_WS_PROTOCOL)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_no_date_header(ws_protocol_cls, http_protocol_cls):
+    class App(WebSocketResponse):
+        async def websocket_connect(self, message):
+            await self.send({"type": "websocket.accept"})
+
+    async def open_connection(url):
+        async with websockets.connect(url) as websocket:
+            return websocket.response_headers
+
+    config = Config(
+        app=App,
+        ws=ws_protocol_cls,
+        http=http_protocol_cls,
+        lifespan="off",
+        date_header=False,
+    )
+    async with run_server(config):
+        headers = await open_connection("ws://127.0.0.1:8000")
+        assert "date" not in headers
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("ws_protocol_cls", ONLY_WS_PROTOCOL)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_multiple_server_header(ws_protocol_cls, http_protocol_cls):
+    class App(WebSocketResponse):
+        async def websocket_connect(self, message):
+            await self.send(
+                {
+                    "type": "websocket.accept",
+                    "headers": [
+                        (b"Server", b"over-ridden"),
+                        (b"Server", b"another-value"),
+                    ],
+                }
+            )
+
+    async def open_connection(url):
+        async with websockets.connect(url) as websocket:
+            return websocket.response_headers
+
+    config = Config(
+        app=App,
+        ws=ws_protocol_cls,
+        http=http_protocol_cls,
+        lifespan="off",
+    )
+    async with run_server(config):
+        headers = await open_connection("ws://127.0.0.1:8000")
+        assert all(
+            x in headers.get_all("Server") for x in ["over-ridden", "another-value"]
+        )

uvicorn/protocols/utils.py

@@ -1,6 +1,6 @@
 import asyncio
 import urllib.parse
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 if TYPE_CHECKING:
     from asgiref.typing import WWWScope
@@ -53,3 +53,10 @@ def get_path_with_query_string(scope: "WWWScope") -> str:
             path_with_query_string, scope["query_string"].decode("ascii")
         )
     return path_with_query_string
+
+
+def get_server_header(default_headers: List[Tuple[bytes, bytes]], override: str) -> str:
+    return next(
+        (i for i in default_headers if i[0] in (b"Server", b"server")),
+        [override.encode()],
+    )[-1].decode()

uvicorn/protocols/websockets/websockets_impl.py

@@ -15,7 +15,12 @@
 
 from uvicorn.config import Config
 from uvicorn.logging import TRACE_LOG_LEVEL
-from uvicorn.protocols.utils import get_local_addr, get_remote_addr, is_ssl
+from uvicorn.protocols.utils import (
+    get_local_addr,
+    get_remote_addr,
+    get_server_header,
+    is_ssl,
+)
 from uvicorn.server import ServerState
 
 if sys.version_info < (3, 8):
@@ -65,10 +70,18 @@ def __init__(
         self.app = config.loaded_app
         self.loop = _loop or asyncio.get_event_loop()
         self.root_path = config.root_path
+        if self.config.server_header:
+            self.server_header = get_server_header(
+                default_headers=server_state.default_headers,
+                override=websockets.server.USER_AGENT,
+            )
+        else:
+            self.server_header = None
 
         # Shared server state
         self.connections = server_state.connections
         self.tasks = server_state.tasks
+        self.default_headers = server_state.default_headers
 
         # Connection state
         self.transport: asyncio.Transport = None  # type: ignore[assignment]
@@ -98,6 +111,7 @@ def __init__(
             max_size=self.config.ws_max_size,
             ping_interval=self.config.ws_ping_interval,
             ping_timeout=self.config.ws_ping_timeout,
+            server_header=self.server_header,
             extensions=extensions,
             logger=logging.getLogger("uvicorn.error"),
             extra_headers=[],
@@ -260,12 +274,19 @@ async def asgi_send(self, message: "ASGISendEvent") -> None:
                 self.accepted_subprotocol = cast(
                     Optional[Subprotocol], message.get("subprotocol")
                 )
-                if "headers" in message:
-                    self.extra_headers.extend(
-                        # ASGI spec requires bytes
-                        # But for compatibility we need to convert it to strings
-                        (name.decode("latin-1"), value.decode("latin-1"))
-                        for name, value in message["headers"]
+                headers = list(message.get("headers", []))
+                _added_names = []
+                for name, value in headers:
+                    if name.lower() in _added_names:
+                        continue
+                    _added_names.append(name.lower())
+                    self.extra_headers.append(
+                        (
+                            # ASGI spec requires bytes
+                            # But for compatibility we need to convert it to strings
+                            name.decode("latin-1"),
+                            value.decode("latin-1"),
+                        )
                     )
                 self.handshake_started_event.set()
 

uvicorn/protocols/websockets/wsproto_impl.py

@@ -27,6 +27,7 @@ def __init__(self, config, server_state, _loop=None):
         # Shared server state
         self.connections = server_state.connections
         self.tasks = server_state.tasks
+        self.default_headers = server_state.default_headers
 
         # Connection state
         self.transport = None
@@ -248,7 +249,7 @@ async def send(self, message):
                 )
                 self.handshake_complete = True
                 subprotocol = message.get("subprotocol")
-                extra_headers = message.get("headers", [])
+                extra_headers = list(message.get("headers", [])) + self.default_headers
                 extensions = []
                 if self.config.ws_per_message_deflate:
                     extensions.append(PerMessageDeflate())