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,10 +310,37 @@ 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 origin_maybe_regex in self.origins:
+                if origin_maybe_regex is None:
+                    continue
+                if self.probably_regex(origin_maybe_regex):
+                    valid = re.match(origin_maybe_regex, origin)
+                else:
+                    valid = origin_maybe_regex == origin
+                if valid:
+                    break
+            if not valid:
                 raise InvalidOrigin(origin)
         return origin
 
+    @staticmethod
+    def probably_regex(maybe_regex: str) -> bool:
+        """
+        Determine if the given string is a regex.
+
+        Args:
+            maybe_regex: A string that may be a regex.
+
+        Returns:
+            True if the string is a regex, False otherwise.
+
+        """
+        common_regex_chars = ['*', '\\', ']', '?', '$', '^', '[', ']', '(', ')']
+        # Use common characters used in regular expressions as a proxy
+        # for if this string is in fact a regex.
+        return any((c in maybe_regex for c in common_regex_chars))
+
     def process_extensions(
         self,
         headers: Headers,

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