Repository: aaugustin/websockets
Pull request #1041 (2021-09-05 18:57:04)

Add limits on reads


Patch:
diff --git a/docs/conf.py b/docs/conf.py
index d9e3cd5..ffe61f7 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -74,8 +74,9 @@ autodoc_typehints_description_target = "documented"
 
 # Workaround for https://github.com/sphinx-doc/sphinx/issues/9560
 from sphinx.domains.python import PythonDomain
-assert PythonDomain.object_types['data'].roles == ('data', 'obj')
-PythonDomain.object_types['data'].roles = ('data', 'class', 'obj')
+
+assert PythonDomain.object_types["data"].roles == ("data", "obj")
+PythonDomain.object_types["data"].roles = ("data", "class", "obj")
 
 intersphinx_mapping = {"python": ("https://docs.python.org/3", None)}
 
diff --git a/docs/howto/faq.rst b/docs/howto/faq.rst
index ec657bb..a32e5ec 100644
--- a/docs/howto/faq.rst
+++ b/docs/howto/faq.rst
@@ -11,9 +11,6 @@ FAQ
 
     .. _developing with asyncio: https://docs.python.org/3/library/asyncio-dev.html
 
-.. contents::
-   :local:
-
 Server side
 -----------
 
@@ -221,11 +218,13 @@ Putting an ``await`` statement in a ``for`` or a ``while`` loop isn't enough
 to yield control. Awaiting a coroutine may yield control, but there's no
 guarantee that it will.
 
-For example, ``send()`` only yields control when send buffers are full, which
-never happens in most practical cases.
+For example, :meth:`~legacy.protocol.WebSocketCommonProtocol.send` only yields
+control when send buffers are full, which never happens in most practical
+cases.
 
-If you run a loop that contains only synchronous operations and a ``send()``
-call, you must yield control explicitly with :func:`asyncio.sleep`::
+If you run a loop that contains only synchronous operations and
+a :meth:`~legacy.protocol.WebSocketCommonProtocol.send` call, you must yield
+control explicitly with :func:`asyncio.sleep`::
 
     async def producer(websocket):
         message = generate_next_message()
@@ -248,7 +247,8 @@ blocks the event loop and prevents asyncio from operating normally.
 
 This may lead to messages getting send but not received, to connection
 timeouts, and to unexpected results of shotgun debugging e.g. adding an
-unnecessary call to ``send()`` makes the program functional.
+unnecessary call to :meth:`~legacy.protocol.WebSocketCommonProtocol.send`
+makes the program functional.
 
 Both sides
 ----------
@@ -352,8 +352,8 @@ See the discussion of :doc:`timeouts <../topics/timeouts>` for details.
 If websockets' default timeout of 20 seconds is too short for your use case,
 you can adjust it with the ``ping_timeout`` argument.
 
-How do I set a timeout on ``recv()``?
-.....................................
+How do I set a timeout on :meth:`~legacy.protocol.WebSocketCommonProtocol.recv`?
+................................................................................
 
 Use :func:`~asyncio.wait_for`::
 
diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst
index 8b6325a..84ab6a3 100644
--- a/docs/project/changelog.rst
+++ b/docs/project/changelog.rst
@@ -81,23 +81,47 @@ New features
 
 * Added ``open_timeout`` to :func:`~client.connect`.
 
+* Documented how to integrate with `Django <https://www.djangoproject.com/>`_.
+
+* Documented how to deploy websockets in production, with several options.
+
+* Documented how to authenticate connections.
+
+* Documented how to broadcast messages to many connections.
+
 Improvements
 ............
 
-* Improved logging.
-
-* Provided additional information in :exc:`~exceptions.ConnectionClosed`
-  exceptions.
+* Improved logging. See the :doc:`logging guide <../topics/logging>`.
 
 * Optimized default compression settings to reduce memory usage.
 
+* Optimized processing of client-to-server messages when the C extension isn't
+  available.
+
+* Supported relative redirects in :func:`~client.connect`.
+
+* Handled TCP connection drops during the opening handshake.
+
 * Made it easier to customize authentication with
   :meth:`~auth.BasicAuthWebSocketServerProtocol.check_credentials`.
 
-* Supported relative redirects in :func:`~client.connect`.
+* Provided additional information in :exc:`~exceptions.ConnectionClosed`
+  exceptions.
+
+* Clarified several exceptions or log messages.
+
+* Restructured documentation.
 
 * Improved API documentation.
 
+* Extended FAQ.
+
+Bug fixes
+.........
+
+* Avoided a crash when receiving a ping while the connection is closing.
+
 9.1
 ---
 
diff --git a/docs/reference/client.rst b/docs/reference/client.rst
index eaa6cdd..dc31dc0 100644
--- a/docs/reference/client.rst
+++ b/docs/reference/client.rst
@@ -33,6 +33,8 @@ Client
 
         .. autoattribute:: id
 
+        .. autoattribute:: logger
+
         .. autoproperty:: local_address
 
         .. autoproperty:: remote_address
diff --git a/docs/reference/common.rst b/docs/reference/common.rst
index 3b9f34a..f5422bc 100644
--- a/docs/reference/common.rst
+++ b/docs/reference/common.rst
@@ -24,6 +24,8 @@ Both sides
 
         .. autoattribute:: id
 
+        .. autoattribute:: logger
+
         .. autoproperty:: local_address
 
         .. autoproperty:: remote_address
diff --git a/docs/reference/server.rst b/docs/reference/server.rst
index 0a5a060..1864594 100644
--- a/docs/reference/server.rst
+++ b/docs/reference/server.rst
@@ -50,6 +50,8 @@ Server
 
         .. autoattribute:: id
 
+        .. autoattribute:: logger
+
         .. autoproperty:: local_address
 
         .. autoproperty:: remote_address
diff --git a/docs/topics/authentication.rst b/docs/topics/authentication.rst
index 44c9c61..8beeea9 100644
--- a/docs/topics/authentication.rst
+++ b/docs/topics/authentication.rst
@@ -120,7 +120,7 @@ WebSocket server.
      the old credentials, which may be expired, resulting in an HTTP 401. Then
      the next connection succeeds. Perhaps errors clear the cache.
 
-     When tokens are short-lived on single-use, this bug produces an
+     When tokens are short-lived or single-use, this bug produces an
      interesting effect: every other WebSocket connection fails.
 
    * Safari 14 ignores credentials entirely.
diff --git a/src/websockets/http11.py b/src/websockets/http11.py
index b82a0bf..052719c 100644
--- a/src/websockets/http11.py
+++ b/src/websockets/http11.py
@@ -7,8 +7,18 @@ from typing import Callable, Generator, Optional
 from . import datastructures, exceptions
 
 
+# Maximum total size of headers is around 256 * 4 KiB = 1 MiB
 MAX_HEADERS = 256
-MAX_LINE = 4110
+
+# We can use the same limit for the request line and header lines:
+# "GET <4096 bytes> HTTP/1.1\r\n" = 4111 bytes
+# "Set-Cookie: <4097 bytes>\r\n" = 4111 bytes
+# (RFC requires 4096 bytes; for some reason Firefox supports 4097 bytes.)
+MAX_LINE = 4111
+
+# Support for HTTP response bodies is intended to read an error message
+# returned by a server. It isn't designed to perform large file transfers.
+MAX_BODY = 2 ** 20  # 1 MiB
 
 
 def d(value: bytes) -> str:
@@ -60,7 +70,7 @@ class Request:
     @classmethod
     def parse(
         cls,
-        read_line: Callable[[], Generator[None, None, bytes]],
+        read_line: Callable[[int], Generator[None, None, bytes]],
     ) -> Generator[None, None, Request]:
         """
         Parse a WebSocket handshake request.
@@ -157,9 +167,9 @@ class Response:
     @classmethod
     def parse(
         cls,
-        read_line: Callable[[], Generator[None, None, bytes]],
+        read_line: Callable[[int], Generator[None, None, bytes]],
         read_exact: Callable[[int], Generator[None, None, bytes]],
-        read_to_eof: Callable[[], Generator[None, None, bytes]],
+        read_to_eof: Callable[[int], Generator[None, None, bytes]],
     ) -> Generator[None, None, Response]:
         """
         Parse a WebSocket handshake response.
@@ -234,7 +244,16 @@ class Response:
                 content_length = int(raw_content_length)
 
             if content_length is None:
-                body = yield from read_to_eof()
+                try:
+                    body = yield from read_to_eof(MAX_BODY)
+                except RuntimeError:
+                    raise exceptions.SecurityError(
+                        f"body too large: over {MAX_BODY} bytes"
+                    )
+            elif content_length > MAX_BODY:
+                raise exceptions.SecurityError(
+                    f"body too large: {content_length} bytes"
+                )
             else:
                 body = yield from read_exact(content_length)
 
@@ -255,7 +274,7 @@ class Response:
 
 
 def parse_headers(
-    read_line: Callable[[], Generator[None, None, bytes]],
+    read_line: Callable[[int], Generator[None, None, bytes]],
 ) -> Generator[None, None, datastructures.Headers]:
     """
     Parse HTTP headers.
@@ -306,7 +325,7 @@ def parse_headers(
 
 
 def parse_line(
-    read_line: Callable[[], Generator[None, None, bytes]],
+    read_line: Callable[[int], Generator[None, None, bytes]],
 ) -> Generator[None, None, bytes]:
     """
     Parse a single line.
@@ -322,10 +341,9 @@ def parse_line(
         SecurityError: if the response exceeds a security limit.
 
     """
-    # Security: TODO: add a limit here
-    line = yield from read_line()
-    # Security: this guarantees header values are small (hard-coded = 4 KiB)
-    if len(line) > MAX_LINE:
+    try:
+        line = yield from read_line(MAX_LINE)
+    except RuntimeError:
         raise exceptions.SecurityError("line too long")
     # Not mandatory but safe - https://www.rfc-editor.org/rfc/rfc7230.html#section-3.5
     if not line.endswith(b"\r\n"):
diff --git a/src/websockets/legacy/protocol.py b/src/websockets/legacy/protocol.py
index 618e451..a31a5c7 100644
--- a/src/websockets/legacy/protocol.py
+++ b/src/websockets/legacy/protocol.py
@@ -218,7 +218,8 @@ class WebSocketCommonProtocol(asyncio.Protocol):
             logger = logging.getLogger("websockets.protocol")
         # https://github.com/python/typeshed/issues/5561
         logger = cast(logging.Logger, logger)
-        self.logger = logging.LoggerAdapter(logger, {"websocket": self})
+        self.logger: LoggerLike = logging.LoggerAdapter(logger, {"websocket": self})
+        """Logger for this connection."""
 
         # Track if DEBUG is enabled. Shortcut logging calls if it isn't.
         self.debug = logger.isEnabledFor(logging.DEBUG)
diff --git a/src/websockets/streams.py b/src/websockets/streams.py
index 094cbb5..f861d4b 100644
--- a/src/websockets/streams.py
+++ b/src/websockets/streams.py
@@ -17,7 +17,7 @@ class StreamReader:
         self.buffer = bytearray()
         self.eof = False
 
-    def read_line(self) -> Generator[None, None, bytes]:
+    def read_line(self, m: int) -> Generator[None, None, bytes]:
         """
         Read a LF-terminated line from the stream.
 
@@ -25,8 +25,12 @@ class StreamReader:
 
         The return value includes the LF character.
 
+        Args:
+            m: maximum number bytes to read; this is a security limit.
+
         Raises:
             EOFError: if the stream ends without a LF.
+            RuntimeError: if the stream ends in more than ``m`` bytes.
 
         """
         n = 0  # number of bytes to read
@@ -36,9 +40,13 @@ class StreamReader:
             if n > 0:
                 break
             p = len(self.buffer)
+            if p > m:
+                raise RuntimeError(f"read {p} bytes, expected no more than {m} bytes")
             if self.eof:
                 raise EOFError(f"stream ends after {p} bytes, before end of line")
             yield
+        if n > m:
+            raise RuntimeError(f"read {n} bytes, expected no more than {m} bytes")
         r = self.buffer[:n]
         del self.buffer[:n]
         return r
@@ -66,14 +74,23 @@ class StreamReader:
         del self.buffer[:n]
         return r
 
-    def read_to_eof(self) -> Generator[None, None, bytes]:
+    def read_to_eof(self, m: int) -> Generator[None, None, bytes]:
         """
         Read all bytes from the stream.
 
         This is a generator-based coroutine.
 
+        Args:
+            m: maximum number bytes to read; this is a security limit.
+
+        Raises:
+            RuntimeError: if the stream ends in more than ``m`` bytes.
+
         """
         while not self.eof:
+            p = len(self.buffer)
+            if p > m:
+                raise RuntimeError(f"read {p} bytes, expected no more than {m} bytes")
             yield
         r = self.buffer[:]
         del self.buffer[:]


Test patch:
diff --git a/tests/test_http11.py b/tests/test_http11.py
index afd85f6..61d3779 100644
--- a/tests/test_http11.py
+++ b/tests/test_http11.py
@@ -158,7 +158,8 @@ class ResponseTests(GeneratorTestCase):
         with self.assertRaises(EOFError) as raised:
             next(self.parse())
         self.assertEqual(
-            str(raised.exception), "connection closed while reading HTTP status line"
+            str(raised.exception),
+            "connection closed while reading HTTP status line",
         )
 
     def test_parse_invalid_status_line(self):
@@ -230,6 +231,24 @@ class ResponseTests(GeneratorTestCase):
         response = self.assertGeneratorReturns(gen)
         self.assertEqual(response.body, b"Hello world!\n")
 
+    def test_parse_body_with_content_length_too_long(self):
+        self.reader.feed_data(b"HTTP/1.1 200 OK\r\nContent-Length: 1048577\r\n\r\n")
+        with self.assertRaises(SecurityError) as raised:
+            next(self.parse())
+        self.assertEqual(
+            str(raised.exception),
+            "body too large: 1048577 bytes",
+        )
+
+    def test_parse_body_without_content_length_too_long(self):
+        self.reader.feed_data(b"HTTP/1.1 200 OK\r\n\r\n" + b"a" * 1048577)
+        with self.assertRaises(SecurityError) as raised:
+            next(self.parse())
+        self.assertEqual(
+            str(raised.exception),
+            "body too large: over 1048576 bytes",
+        )
+
     def test_parse_body_with_transfer_encoding(self):
         self.reader.feed_data(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n")
         with self.assertRaises(NotImplementedError) as raised:
@@ -314,8 +333,8 @@ class HeadersTests(GeneratorTestCase):
             next(self.parse_headers())
 
     def test_parse_too_long_line(self):
-        # Header line contains 5 + 4104 + 2 = 4111 bytes.
-        self.reader.feed_data(b"foo: " + b"a" * 4104 + b"\r\n\r\n")
+        # Header line contains 5 + 4105 + 2 = 4112 bytes.
+        self.reader.feed_data(b"foo: " + b"a" * 4105 + b"\r\n\r\n")
         with self.assertRaises(SecurityError):
             next(self.parse_headers())
 
diff --git a/tests/test_streams.py b/tests/test_streams.py
index 8abefbc..fd7c66a 100644
--- a/tests/test_streams.py
+++ b/tests/test_streams.py
@@ -10,24 +10,24 @@ class StreamReaderTests(GeneratorTestCase):
     def test_read_line(self):
         self.reader.feed_data(b"spam\neggs\n")
 
-        gen = self.reader.read_line()
+        gen = self.reader.read_line(32)
         line = self.assertGeneratorReturns(gen)
         self.assertEqual(line, b"spam\n")
 
-        gen = self.reader.read_line()
+        gen = self.reader.read_line(32)
         line = self.assertGeneratorReturns(gen)
         self.assertEqual(line, b"eggs\n")
 
     def test_read_line_need_more_data(self):
         self.reader.feed_data(b"spa")
 
-        gen = self.reader.read_line()
+        gen = self.reader.read_line(32)
         self.assertGeneratorRunning(gen)
         self.reader.feed_data(b"m\neg")
         line = self.assertGeneratorReturns(gen)
         self.assertEqual(line, b"spam\n")
 
-        gen = self.reader.read_line()
+        gen = self.reader.read_line(32)
         self.assertGeneratorRunning(gen)
         self.reader.feed_data(b"gs\n")
         line = self.assertGeneratorReturns(gen)
@@ -37,11 +37,34 @@ class StreamReaderTests(GeneratorTestCase):
         self.reader.feed_data(b"spa")
         self.reader.feed_eof()
 
-        gen = self.reader.read_line()
+        gen = self.reader.read_line(32)
         with self.assertRaises(EOFError) as raised:
             next(gen)
         self.assertEqual(
-            str(raised.exception), "stream ends after 3 bytes, before end of line"
+            str(raised.exception),
+            "stream ends after 3 bytes, before end of line",
+        )
+
+    def test_read_line_too_long(self):
+        self.reader.feed_data(b"spam\neggs\n")
+
+        gen = self.reader.read_line(2)
+        with self.assertRaises(RuntimeError) as raised:
+            next(gen)
+        self.assertEqual(
+            str(raised.exception),
+            "read 5 bytes, expected no more than 2 bytes",
+        )
+
+    def test_read_line_too_long_need_more_data(self):
+        self.reader.feed_data(b"spa")
+
+        gen = self.reader.read_line(2)
+        with self.assertRaises(RuntimeError) as raised:
+            next(gen)
+        self.assertEqual(
+            str(raised.exception),
+            "read 3 bytes, expected no more than 2 bytes",
         )
 
     def test_read_exact(self):
@@ -78,11 +101,12 @@ class StreamReaderTests(GeneratorTestCase):
         with self.assertRaises(EOFError) as raised:
             next(gen)
         self.assertEqual(
-            str(raised.exception), "stream ends after 3 bytes, expected 4 bytes"
+            str(raised.exception),
+            "stream ends after 3 bytes, expected 4 bytes",
         )
 
     def test_read_to_eof(self):
-        gen = self.reader.read_to_eof()
+        gen = self.reader.read_to_eof(32)
 
         self.reader.feed_data(b"spam")
         self.assertGeneratorRunning(gen)
@@ -94,10 +118,21 @@ class StreamReaderTests(GeneratorTestCase):
     def test_read_to_eof_at_eof(self):
         self.reader.feed_eof()
 
-        gen = self.reader.read_to_eof()
+        gen = self.reader.read_to_eof(32)
         data = self.assertGeneratorReturns(gen)
         self.assertEqual(data, b"")
 
+    def test_read_to_eof_too_long(self):
+        gen = self.reader.read_to_eof(2)
+
+        self.reader.feed_data(b"spam")
+        with self.assertRaises(RuntimeError) as raised:
+            next(gen)
+        self.assertEqual(
+            str(raised.exception),
+            "read 4 bytes, expected no more than 2 bytes",
+        )
+
     def test_at_eof_after_feed_data(self):
         gen = self.reader.at_eof()
         self.assertGeneratorRunning(gen)
@@ -137,16 +172,22 @@ class StreamReaderTests(GeneratorTestCase):
         self.reader.feed_eof()
         with self.assertRaises(EOFError) as raised:
             self.reader.feed_data(b"spam")
-        self.assertEqual(str(raised.exception), "stream ended")
+        self.assertEqual(
+            str(raised.exception),
+            "stream ended",
+        )
 
     def test_feed_eof_after_feed_eof(self):
         self.reader.feed_eof()
         with self.assertRaises(EOFError) as raised:
             self.reader.feed_eof()
-        self.assertEqual(str(raised.exception), "stream ended")
+        self.assertEqual(
+            str(raised.exception),
+            "stream ended",
+        )
 
     def test_discard(self):
-        gen = self.reader.read_to_eof()
+        gen = self.reader.read_to_eof(32)
 
         self.reader.feed_data(b"spam")
         self.reader.discard()