Add default headers to WebSockets implementations

docs/settings.md

@@ -90,7 +90,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-date-header` flag doesn't have effect on the WebSockets implementations.
 
 ## HTTPS
 

pyproject.toml

@@ -42,7 +42,7 @@ standard = [
     "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",
+    "websockets>=10.4",
 ]
 
 [project.scripts]

tests/protocols/test_websocket.py

@@ -5,11 +5,13 @@
 
 from tests.protocols.test_http import HTTP_PROTOCOLS
 from tests.utils import run_server
+from uvicorn import Server
 from uvicorn.config import Config
 from uvicorn.protocols.websockets.wsproto_impl import WSProtocol
 
 try:
     import websockets
+    import websockets.exceptions
     from websockets.extensions.permessage_deflate import ClientPerMessageDeflateFactory
 
     from uvicorn.protocols.websockets.websockets_impl import WebSocketProtocol
@@ -18,9 +20,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 +660,197 @@ 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_in_ws(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"]
+        )
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("ws_protocol_cls", WS_PROTOCOLS)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_multiple_arbitrary_headers_with_same_name(
+    ws_protocol_cls, http_protocol_cls
+):
+    class App(WebSocketResponse):
+        async def websocket_connect(self, message):
+            await self.send(
+                {
+                    "type": "websocket.accept",
+                    "headers": [
+                        (b"Potato", b"cool-potato"),
+                        (b"Potato", b"super-cool-potato"),
+                    ],
+                }
+            )
+
+    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("Potato") for x in ["cool-potato", "super-cool-potato"]
+        )
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("ws_protocol_cls", ONLY_WEBSOCKETPROTOCOL)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_multiple_server_header_in_websockets(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 len(headers.get_all("Server")) == 1
+        assert headers.get("Server") == "uvicorn"
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("ws_protocol_cls", ONLY_WEBSOCKETPROTOCOL)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_server_shutdown_when_connection_active_in_websockets(
+    ws_protocol_cls, http_protocol_cls
+):
+    class App(WebSocketResponse):
+        async def websocket_connect(self, message):
+            await self.send({"type": "websocket.accept"})
+
+    config = Config(
+        app=App,
+        ws=ws_protocol_cls,
+        http=http_protocol_cls,
+        lifespan="off",
+    )
+    server = Server(config=config)
+    cancel_handle = asyncio.ensure_future(server.serve(sockets=None))
+    await asyncio.sleep(0.1)
+    async with websockets.connect("ws://127.0.0.1:8000"):
+        ws_conn = list(server.server_state.connections)[0]
+        ws_conn.shutdown()
+        assert ws_conn.ws_server.closing is True
+        assert ws_conn.transport.is_closing()
+    await server.shutdown()
+    cancel_handle.cancel()

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

@@ -19,12 +19,13 @@
     get_local_addr,
     get_path_with_query_string,
     get_remote_addr,
+    get_server_header,
     is_ssl,
 )
 from uvicorn.server import ServerState
 
 if sys.version_info < (3, 8):
-    from typing_extensions import Literal
+    from typing_extensions import Literal  # pragma: no cover
 else:
     from typing import Literal
 
@@ -70,10 +71,17 @@ 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="uvicorn"
+            )
+        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]
@@ -103,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=[],
@@ -265,12 +274,17 @@ 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", []))
+                for name, value in headers:
+                    if name.lower() in [b"server"]:
+                        continue
+                    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

@@ -32,6 +32,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
@@ -253,7 +254,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())