Make it harder to generate invalid HTTP messages.

Author: aaugustin

src/websockets/datastructures.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from collections.abc import Iterable, Iterator, Mapping, MutableMapping
 from typing import Any, Protocol
 
@@ -24,6 +25,10 @@ def __str__(self) -> str:
         return super().__str__()
 
 
+# Obsolete line folding for header values is not supported.
+is_valid_header_value = re.compile(r"[\x09\x20-\x7e\x80-\xff]*").fullmatch
+
+
 class Headers(MutableMapping[str, str]):
     """
     Efficient data structure for manipulating HTTP headers.
@@ -107,6 +112,8 @@ def __getitem__(self, key: str) -> str:
             raise MultipleValuesError(key)
 
     def __setitem__(self, key: str, value: str) -> None:
+        if not is_valid_header_value(str(value)):
+            raise InvalidHeaderValue(value)
         self._dict.setdefault(key.lower(), []).append(value)
         self._list.append((key, value))
 
@@ -181,3 +188,7 @@ def __getitem__(self, key: str) -> str: ...
 keys and values are :class:`str`.
 
 """
+
+
+# At the bottom to break an import cycle.
+from .exceptions import InvalidHeaderValue  # noqa: E402

tests/test_datastructures.py

@@ -1,6 +1,7 @@
 import unittest
 
 from websockets.datastructures import *
+from websockets.exceptions import InvalidHeaderValue
 
 
 class MultipleValuesErrorTests(unittest.TestCase):
@@ -121,6 +122,10 @@ def test_setitem_case_insensitive(self):
         self.headers["upgrade"] = "websocket"
         self.assertEqual(self.headers["Upgrade"], "websocket")
 
+    def test_setitem_invalid_value(self):
+        with self.assertRaises(InvalidHeaderValue):
+            self.headers["X-Invalid"] = "multi\r\nline"
+
     def test_delitem(self):
         del self.headers["Connection"]
         with self.assertRaises(KeyError):

tests/test_server.py

@@ -145,7 +145,7 @@ def test_send_response_without_accept_or_reject(self, _formatdate):
                 Headers(
                     {
                         "Connection": "close",
-                        "Content-Length": 6,
+                        "Content-Length": "6",
                         "Content-Type": "text/plain",
                     }
                 ),