Add regex support in ServerProtocol(origins=...)

Author: daH005

src/websockets/asyncio/server.py

@@ -599,7 +599,7 @@ def handler(websocket):
             See :meth:`~asyncio.loop.create_server` for details.
         port: TCP port the server listens on.
             See :meth:`~asyncio.loop.create_server` for details.
-        origins: Acceptable values of the ``Origin`` header, for defending
+        origins: Acceptable values of the ``Origin`` header, including regular expressions, for defending
             against Cross-Site WebSocket Hijacking attacks. Include :obj:`None`
             in the list if the lack of an origin is acceptable.
         extensions: List of supported extensions, in order in which they

src/websockets/server.py

@@ -5,6 +5,7 @@
 import email.utils
 import http
 import warnings
+import re
 from collections.abc import Generator, Sequence
 from typing import Any, Callable, cast
 
@@ -49,7 +50,7 @@ class ServerProtocol(Protocol):
     Sans-I/O implementation of a WebSocket server connection.
 
     Args:
-        origins: Acceptable values of the ``Origin`` header; include
+        origins: Acceptable values of the ``Origin`` header, including regular expressions; include
             :obj:`None` in the list if the lack of an origin is acceptable.
             This is useful for defending against Cross-Site WebSocket
             Hijacking attacks.
@@ -309,7 +310,15 @@ def process_origin(self, headers: Headers) -> Origin | None:
         if origin is not None:
             origin = cast(Origin, origin)
         if self.origins is not None:
-            if origin not in self.origins:
+            valid = False
+            for acceptable_origin in self.origins:
+                if isinstance(acceptable_origin, re.Pattern):
+                    valid = acceptable_origin.match(origin)
+                else:
+                    valid = acceptable_origin == origin
+                if valid:
+                    break
+            if not valid:
                 raise InvalidOrigin(origin)
         return origin
 

src/websockets/sync/server.py

@@ -399,7 +399,7 @@ def handler(websocket):
             You may call :func:`socket.create_server` to create a suitable TCP
             socket.
         ssl: Configuration for enabling TLS on the connection.
-        origins: Acceptable values of the ``Origin`` header, for defending
+        origins: Acceptable values of the ``Origin`` header, including regular expressions, for defending
             against Cross-Site WebSocket Hijacking attacks. Include :obj:`None`
             in the list if the lack of an origin is acceptable.
         extensions: List of supported extensions, in order in which they

src/websockets/typing.py

@@ -2,6 +2,7 @@
 
 import http
 import logging
+import re
 import typing
 from typing import Any, NewType, Optional, Union
 
@@ -45,8 +46,8 @@
 Types accepted where an :class:`~http.HTTPStatus` is expected."""
 
 
-Origin = NewType("Origin", str)
-"""Value of a ``Origin`` header."""
+Origin = Union[str, re.Pattern]
+"""Value of a ``Origin`` header or a regular expression."""
 
 
 Subprotocol = NewType("Subprotocol", str)

tests/test_server.py

@@ -1,5 +1,6 @@
 import http
 import logging
+import re
 import sys
 import unittest
 from unittest.mock import patch
@@ -623,6 +624,36 @@ def test_unsupported_origin(self):
             "invalid Origin header: https://original.example.com",
         )
 
+    def test_supported_origin_by_regex(self):
+        """Handshake succeeds when checking origins and the origin is supported by a regular expression."""
+        server = ServerProtocol(
+            origins=["https://example.com", re.compile(r"https://other.*")]
+        )
+        request = make_request()
+        request.headers["Origin"] = "https://other.example.com"
+        response = server.accept(request)
+        server.send_response(response)
+
+        self.assertHandshakeSuccess(server)
+        self.assertEqual(server.origin, "https://other.example.com")
+
+    def test_unsupported_origin_by_regex(self):
+        """Handshake succeeds when checking origins and the origin is unsupported by a regular expression."""
+        server = ServerProtocol(
+            origins=["https://example.com", re.compile(r"https://other.*")]
+        )
+        request = make_request()
+        request.headers["Origin"] = "https://original.example.com"
+        response = server.accept(request)
+        server.send_response(response)
+
+        self.assertEqual(response.status_code, 403)
+        self.assertHandshakeError(
+            server,
+            InvalidOrigin,
+            "invalid Origin header: https://original.example.com",
+        )
+
     def test_no_origin_accepted(self):
         """Handshake succeeds when the lack of an origin is accepted."""
         server = ServerProtocol(origins=[None])