Support separate limits for messages and fragments.

Fix #1602.

Author: aaugustin

docs/faq/misc.rst

@@ -22,6 +22,7 @@ you must disable:
 
 * Compression: set ``compression=None``
 * Keepalive: set ``ping_interval=None``
+* Limits: set ``max_size=None``
 * UTF-8 decoding: send ``bytes`` rather than ``str``
 
 Then, please consider whether websockets is the bottleneck of the performance

docs/project/changelog.rst

@@ -35,6 +35,7 @@ notice.
 Improvements
 ............
 
+* Allowed setting separate limits for messages and fragments with ``max_size``.
 * Added support for HTTP/1.0 proxies.
 
 15.0.1

src/websockets/asyncio/client.py

@@ -236,7 +236,9 @@ class connect:
  close_timeout: Timeout for closing the connection in seconds.
  :obj:`None` disables the timeout.
  max_size: Maximum size of incoming messages in bytes.
- :obj:`None` disables the limit.
+ :obj:`None` disables the limit. You may pass a ``(max_message_size,
+ max_fragment_size)`` tuple to set different limits for messages and
+ fragments when you expect long messages sent in short fragments.
  max_queue: High-water mark of the buffer where frames are received.
  It defaults to 16 frames. The low-water mark defaults to ``max_queue
  // 4``. You may pass a ``(high, low)`` tuple to set the high-water
@@ -314,7 +316,7 @@ def __init__(
  ping_timeout: float | None = 20,
  close_timeout: float | None = 10,
  # Limits
- max_size: int | None = 2**20,
+ max_size: int | None | tuple[int | None, int | None] = 2**20,
  max_queue: int | None | tuple[int | None, int | None] = 16,
  write_limit: int | tuple[int, int | None] = 2**15,
  # Logging

src/websockets/asyncio/connection.py

@@ -63,9 +63,11 @@ def __init__(
  self.ping_interval = ping_interval
  self.ping_timeout = ping_timeout
  self.close_timeout = close_timeout
+ self.max_queue: tuple[int | None, int | None]
  if isinstance(max_queue, int) or max_queue is None:
- max_queue = (max_queue, None)
- self.max_queue = max_queue
+ self.max_queue = (max_queue, None)
+ else:
+ self.max_queue = max_queue
  if isinstance(write_limit, int):
  write_limit = (write_limit, None)
  self.write_limit = write_limit

src/websockets/asyncio/server.py

@@ -644,7 +644,9 @@ def handler(websocket):
  close_timeout: Timeout for closing connections in seconds.
  :obj:`None` disables the timeout.
  max_size: Maximum size of incoming messages in bytes.
- :obj:`None` disables the limit.
+ :obj:`None` disables the limit. You may pass a ``(max_message_size,
+ max_fragment_size)`` tuple to set different limits for messages and
+ fragments when you expect long messages sent in short fragments.
  max_queue: High-water mark of the buffer where frames are received.
  It defaults to 16 frames. The low-water mark defaults to ``max_queue
  // 4``. You may pass a ``(high, low)`` tuple to set the high-water
@@ -719,7 +721,7 @@ def __init__(
  ping_timeout: float | None = 20,
  close_timeout: float | None = 10,
  # Limits
- max_size: int | None = 2**20,
+ max_size: int | None | tuple[int | None, int | None] = 2**20,
  max_queue: int | None | tuple[int | None, int | None] = 16,
  write_limit: int | tuple[int, int | None] = 2**15,
  # Logging

src/websockets/client.py

@@ -60,8 +60,10 @@ class ClientProtocol(Protocol):
  subprotocols: List of supported subprotocols, in order of decreasing
  preference.
  state: Initial state of the WebSocket connection.
- max_size: Maximum size of incoming messages in bytes;
- :obj:`None` disables the limit.
+ max_size: Maximum size of incoming messages in bytes.
+ :obj:`None` disables the limit. You may pass a ``(max_message_size,
+ max_fragment_size)`` tuple to set different limits for messages and
+ fragments when you expect long messages sent in short fragments.
  logger: Logger for this connection;
  defaults to ``logging.getLogger("websockets.client")``;
  see the :doc:`logging guide <../../topics/logging>` for details.
@@ -76,7 +78,7 @@ def __init__(
  extensions: Sequence[ClientExtensionFactory] | None = None,
  subprotocols: Sequence[Subprotocol] | None = None,
  state: State = CONNECTING,
- max_size: int | None = 2**20,
+ max_size: int | None | tuple[int | None, int | None] = 2**20,
  logger: LoggerLike | None = None,
  ) -> None:
  super().__init__(

src/websockets/exceptions.py

@@ -394,11 +394,11 @@ def __init__(
  self,
  size_or_message: int | None | str,
  max_size: int | None = None,
- cur_size: int | None = None,
+ current_size: int | None = None,
  ) -> None:
  if isinstance(size_or_message, str):
  assert max_size is None
- assert cur_size is None
+ assert current_size is None
  warnings.warn( # deprecated in 14.0 - 2024-11-09
  "PayloadTooBig(message) is deprecated; "
  "change to PayloadTooBig(size, max_size)",
@@ -410,8 +410,8 @@ def __init__(
  self.size: int | None = size_or_message
  assert max_size is not None
  self.max_size: int = max_size
- self.cur_size: int | None = None
- self.set_current_size(cur_size)
+ self.current_size: int | None = None
+ self.set_current_size(current_size)
 
  def __str__(self) -> str:
  if self.message is not None:
@@ -420,16 +420,16 @@ def __str__(self) -> str:
  message = "frame "
  if self.size is not None:
  message += f"with {self.size} bytes "
- if self.cur_size is not None:
- message += f"after reading {self.cur_size} bytes "
+ if self.current_size is not None:
+ message += f"after reading {self.current_size} bytes "
  message += f"exceeds limit of {self.max_size} bytes"
  return message
 
- def set_current_size(self, cur_size: int | None) -> None:
- assert self.cur_size is None
- if cur_size is not None:
- self.max_size += cur_size
- self.cur_size = cur_size
+ def set_current_size(self, current_size: int | None) -> None:
+ assert self.current_size is None
+ if current_size is not None:
+ self.max_size += current_size
+ self.current_size = current_size
 
 
 class InvalidState(WebSocketException, AssertionError):

src/websockets/protocol.py

@@ -77,8 +77,10 @@ class Protocol:
  Args:
  side: :attr:`~Side.CLIENT` or :attr:`~Side.SERVER`.
  state: Initial state of the WebSocket connection.
- max_size: Maximum size of incoming messages in bytes;
- :obj:`None` disables the limit.
+ max_size: Maximum size of incoming messages in bytes.
+ :obj:`None` disables the limit. You may pass a ``(max_message_size,
+ max_fragment_size)`` tuple to set different limits for messages and
+ fragments when you expect long messages sent in short fragments.
  logger: Logger for this connection; depending on ``side``,
  defaults to ``logging.getLogger("websockets.client")``
  or ``logging.getLogger("websockets.server")``;
@@ -91,7 +93,7 @@ def __init__(
  side: Side,
  *,
  state: State = OPEN,
- max_size: int | None = 2**20,
+ max_size: tuple[int | None, int | None] | int | None = 2**20,
  logger: LoggerLike | None = None,
  ) -> None:
  # Unique identifier. For logs.
@@ -114,11 +116,14 @@ def __init__(
  self.state = state
 
  # Maximum size of incoming messages in bytes.
- self.max_size = max_size
+ if isinstance(max_size, int) or max_size is None:
+ self.max_message_size, self.max_fragment_size = max_size, None
+ else:
+ self.max_message_size, self.max_fragment_size = max_size
 
  # Current size of incoming message in bytes. Only set while reading a
  # fragmented message i.e. a data frames with the FIN bit not set.
- self.cur_size: int | None = None
+ self.current_size: int | None = None
 
  # True while sending a fragmented message i.e. a data frames with the
  # FIN bit not set.
@@ -578,12 +583,19 @@ def parse(self) -> Generator[None]:
  # connection isn't closed cleanly.
  raise EOFError("unexpected end of stream")
 
- if self.max_size is None:
- max_size = None
- elif self.cur_size is None:
- max_size = self.max_size
- else:
- max_size = self.max_size - self.cur_size
+ max_size = None
+
+ if self.max_message_size is not None:
+ if self.current_size is None:
+ max_size = self.max_message_size
+ else:
+ max_size = self.max_message_size - self.current_size
+
+ if self.max_fragment_size is not None:
+ if max_size is None:
+ max_size = self.max_fragment_size
+ else:
+ max_size = min(max_size, self.max_fragment_size)
 
  # During a normal closure, execution ends here on the next
  # iteration of the loop after receiving a close frame. At
@@ -613,7 +625,7 @@ def parse(self) -> Generator[None]:
  self.parser_exc = exc
 
  except PayloadTooBig as exc:
- exc.set_current_size(self.cur_size)
+ exc.set_current_size(self.current_size)
  self.fail(CloseCode.MESSAGE_TOO_BIG, str(exc))
  self.parser_exc = exc
 
@@ -664,18 +676,18 @@ def recv_frame(self, frame: Frame) -> None:
 
  """
  if frame.opcode is OP_TEXT or frame.opcode is OP_BINARY:
- if self.cur_size is not None:
+ if self.current_size is not None:
  raise ProtocolError("expected a continuation frame")
  if not frame.fin:
- self.cur_size = len(frame.data)
+ self.current_size = len(frame.data)
 
  elif frame.opcode is OP_CONT:
- if self.cur_size is None:
+ if self.current_size is None:
  raise ProtocolError("unexpected continuation frame")
  if frame.fin:
- self.cur_size = None
+ self.current_size = None
  else:
- self.cur_size += len(frame.data)
+ self.current_size += len(frame.data)
 
  elif frame.opcode is OP_PING:
  # 5.5.2. Ping: "Upon receipt of a Ping frame, an endpoint MUST
@@ -696,7 +708,7 @@ def recv_frame(self, frame: Frame) -> None:
  assert self.close_sent is not None
  self.close_rcvd_then_sent = False
 
- if self.cur_size is not None:
+ if self.current_size is not None:
  raise ProtocolError("incomplete fragmented message")
 
  # 5.5.1 Close: "If an endpoint receives a Close frame and did

src/websockets/server.py

@@ -65,8 +65,10 @@ class ServerProtocol(Protocol):
  signature as the :meth:`select_subprotocol` method, including a
  :class:`ServerProtocol` instance as first argument.
  state: Initial state of the WebSocket connection.
- max_size: Maximum size of incoming messages in bytes;
- :obj:`None` disables the limit.
+ max_size: Maximum size of incoming messages in bytes.
+ :obj:`None` disables the limit. You may pass a ``(max_message_size,
+ max_fragment_size)`` tuple to set different limits for messages and
+ fragments when you expect long messages sent in short fragments.
  logger: Logger for this connection;
  defaults to ``logging.getLogger("websockets.server")``;
  see the :doc:`logging guide <../../topics/logging>` for details.
@@ -87,7 +89,7 @@ def __init__(
  | None
  ) = None,
  state: State = CONNECTING,
- max_size: int | None = 2**20,
+ max_size: int | None | tuple[int | None, int | None] = 2**20,
  logger: LoggerLike | None = None,
  ) -> None:
  super().__init__(

src/websockets/sync/client.py

@@ -150,7 +150,7 @@ def connect(
  ping_timeout: float | None = 20,
  close_timeout: float | None = 10,
  # Limits
- max_size: int | None = 2**20,
+ max_size: int | None | tuple[int | None, int | None] = 2**20,
  max_queue: int | None | tuple[int | None, int | None] = 16,
  # Logging
  logger: LoggerLike | None = None,
@@ -210,7 +210,9 @@ def connect(
  close_timeout: Timeout for closing the connection in seconds.
  :obj:`None` disables the timeout.
  max_size: Maximum size of incoming messages in bytes.
- :obj:`None` disables the limit.
+ :obj:`None` disables the limit. You may pass a ``(max_message_size,
+ max_fragment_size)`` tuple to set different limits for messages and
+ fragments when you expect long messages sent in short fragments.
  max_queue: High-water mark of the buffer where frames are received.
  It defaults to 16 frames. The low-water mark defaults to ``max_queue
  // 4``. You may pass a ``(high, low)`` tuple to set the high-water