Defaults ws max_size on server to 16MB (encode#995)

* Defaults ws max_size on server to 16MB * Renamed test * Lint * Lint 2 * Renamed to ws_max_size * Added test for Config * Added click option * Flake8 * Updated docs in deployment * Updated docs in index * Added usage note in settings.md * Test with new test run_server * MOre scenarios * Lint * Update tests/protocols/test_websocket.py Co-authored-by: Marcelo Trylesinski <[email protected]> Co-authored-by: Tom Christie <[email protected]> Co-authored-by: Marcelo Trylesinski <[email protected]>
sephioh · Oct 29, 2022 · 2932a59 · 2932a59
1 parent f5b92fb
commit 2932a59
Show file tree

Hide file tree

Showing 8 changed files with 75 additions and 0 deletions.
diff --git a/docs/deployment.md b/docs/deployment.md
@@ -49,6 +49,8 @@ Options:
   --ws [auto|none|websockets|wsproto]
                                   WebSocket protocol implementation.
                                   [default: auto]
+  --ws-max-size INTEGER           WebSocket max size message in bytes
+                                  [default: 16777216]
   --lifespan [auto|on|off]        Lifespan implementation.  [default: auto]
   --interface [auto|asgi3|asgi2|wsgi]
                                   Select ASGI3, ASGI2, or WSGI as the
@@ -101,6 +103,7 @@ Options:
   --help                          Show this message and exit.
 ```
 
+
 See the [settings documentation](settings.md) for more details on the supported options for running uvicorn.
 
 ## Running programmatically

diff --git a/docs/index.md b/docs/index.md
@@ -119,6 +119,8 @@ Options:
   --ws [auto|none|websockets|wsproto]
                                   WebSocket protocol implementation.
                                   [default: auto]
+  --ws-max-size INTEGER           WebSocket max size message in bytes
+                                  [default: 16777216]
   --lifespan [auto|on|off]        Lifespan implementation.  [default: auto]
   --interface [auto|asgi3|asgi2|wsgi]
                                   Select ASGI3, ASGI2, or WSGI as the

diff --git a/docs/settings.md b/docs/settings.md
@@ -42,6 +42,7 @@ you should put `uvicorn.run` into `if __name__ == '__main__'` clause in the main
 * `--loop <str>` - Set the event loop implementation. The uvloop implementation provides greater performance, but is not compatible with Windows or PyPy. **Options:** *'auto', 'asyncio', 'uvloop'.* **Default:** *'auto'*.
 * `--http <str>` - Set the HTTP protocol implementation. The httptools implementation provides greater performance, but it not compatible with PyPy, and requires compilation on Windows. **Options:** *'auto', 'h11', 'httptools'.* **Default:** *'auto'*.
 * `--ws <str>` - Set the WebSockets protocol implementation. Either of the `websockets` and `wsproto` packages are supported. Use `'none'` to deny all websocket requests. **Options:** *'auto', 'none', 'websockets', 'wsproto'.* **Default:** *'auto'*.
+* `--ws-max-size <int>` - Set the WebSockets max message size, in bytes. Please note that this can be used only with the default `websockets` protocol.
 * `--lifespan <str>` - Set the Lifespan protocol implementation. **Options:** *'auto', 'on', 'off'.* **Default:** *'auto'*.
 
 ## Application Interface

diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py
@@ -16,6 +16,7 @@
     ClientPerMessageDeflateFactory = None
 
 
+ONLY_WEBSOCKETPROTOCOL = [p for p in [WebSocketProtocol] if p is not None]
 WS_PROTOCOLS = [p for p in [WSProtocol, WebSocketProtocol] if p is not None]
 pytestmark = pytest.mark.skipif(
     websockets is None, reason="This test needs the websockets module"
@@ -457,3 +458,53 @@ async def get_subprotocol(url):
     async with run_server(config):
         accepted_subprotocol = await get_subprotocol("ws://127.0.0.1:8000")
         assert accepted_subprotocol == subprotocol
+
+
+MAX_WS_BYTES = 1024 * 1024 * 16
+MAX_WS_BYTES_PLUS1 = MAX_WS_BYTES + 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("protocol_cls", ONLY_WEBSOCKETPROTOCOL)
+@pytest.mark.parametrize(
+    "client_size_sent, server_size_max, expected_result",
+    [
+        (MAX_WS_BYTES, MAX_WS_BYTES, 0),
+        (MAX_WS_BYTES_PLUS1, MAX_WS_BYTES, 1009),
+        (10, 10, 0),
+        (11, 10, 1009),
+    ],
+    ids=[
+        "max=defaults sent=defaults",
+        "max=defaults sent=defaults+1",
+        "max=10 sent=10",
+        "max=10 sent=11",
+    ],
+)
+async def test_send_binary_data_to_server_bigger_than_default(
+    protocol_cls, client_size_sent, server_size_max, expected_result
+):
+    class App(WebSocketResponse):
+        async def websocket_connect(self, message):
+            await self.send({"type": "websocket.accept"})
+
+        async def websocket_receive(self, message):
+            _bytes = message.get("bytes")
+            await self.send({"type": "websocket.send", "bytes": _bytes})
+
+    async def send_text(url):
+        async with websockets.connect(url, max_size=client_size_sent) as websocket:
+            await websocket.send(b"\x01" * client_size_sent)
+            return await websocket.recv()
+
+    config = Config(
+        app=App, ws=protocol_cls, lifespan="off", ws_max_size=server_size_max
+    )
+    async with run_server(config):
+        if expected_result == 0:
+            data = await send_text("ws://127.0.0.1:8000")
+            assert data == b"\x01" * client_size_sent
+        else:
+            with pytest.raises(websockets.ConnectionClosedError) as e:
+                data = await send_text("ws://127.0.0.1:8000")
+            assert e.value.code == expected_result
diff --git a/tests/test_config.py b/tests/test_config.py
@@ -273,3 +273,9 @@ def test_config_log_level(log_level):
     assert logging.getLogger("uvicorn.access").level == log_level
     assert logging.getLogger("uvicorn.asgi").level == log_level
     assert config.log_level == log_level
+
+
+def test_ws_max_size():
+    config = Config(app=asgi_app, ws_max_size=1000)
+    config.load()
+    assert config.ws_max_size == 1000
diff --git a/uvicorn/config.py b/uvicorn/config.py
@@ -131,6 +131,7 @@ def __init__(
         loop="auto",
         http="auto",
         ws="auto",
+        ws_max_size=16 * 1024 * 1024,
         lifespan="auto",
         env_file=None,
         log_config=LOGGING_CONFIG,
@@ -170,6 +171,7 @@ def __init__(
         self.loop = loop
         self.http = http
         self.ws = ws
+        self.ws_max_size = ws_max_size
         self.lifespan = lifespan
         self.log_config = log_config
         self.log_level = log_level

diff --git a/uvicorn/main.py b/uvicorn/main.py
@@ -113,6 +113,13 @@ def print_version(ctx, param, value):
     help="WebSocket protocol implementation.",
     show_default=True,
 )
+@click.option(
+    "--ws-max-size",
+    type=int,
+    default=16777216,
+    help="WebSocket max size message in bytes",
+    show_default=True,
+)
 @click.option(
     "--lifespan",
     type=LIFESPAN_CHOICES,
@@ -289,6 +296,7 @@ def main(
     loop: str,
     http: str,
     ws: str,
+    ws_max_size: int,
     lifespan: str,
     interface: str,
     debug: bool,
@@ -330,6 +338,7 @@ def main(
         "loop": loop,
         "http": http,
         "ws": ws,
+        "ws_max_size": ws_max_size,
         "lifespan": lifespan,
         "env_file": env_file,
         "log_config": LOGGING_CONFIG if log_config is None else log_config,

diff --git a/uvicorn/protocols/websockets/websockets_impl.py b/uvicorn/protocols/websockets/websockets_impl.py
@@ -59,6 +59,7 @@ def __init__(self, config, server_state, _loop=None):
             ws_handler=self.ws_handler,
             ws_server=self.ws_server,
             extensions=[ServerPerMessageDeflateFactory()],
+            max_size=self.config.ws_max_size,
         )
 
     def connection_made(self, transport):