diff --git a/src/wslink/backends/aiohttp/__init__.py b/src/wslink/backends/aiohttp/__init__.py index d6ad19d..ca6880a 100644 --- a/src/wslink/backends/aiohttp/__init__.py +++ b/src/wslink/backends/aiohttp/__init__.py @@ -17,6 +17,7 @@ MAX_MSG_SIZE = int(os.environ.get("WSLINK_MAX_MSG_SIZE", "4194304")) HEART_BEAT = int(os.environ.get("WSLINK_HEART_BEAT", "30")) # 30 seconds HTTP_HEADERS = os.environ.get("WSLINK_HTTP_HEADERS") # path to json file +WS_COMPRESS = bool(int(os.environ.get("WSLINK_WS_COMPRESS", "1"))) # permessage-deflate if HTTP_HEADERS and Path(HTTP_HEADERS).exists(): HTTP_HEADERS = json.loads(Path(HTTP_HEADERS).read_text()) @@ -27,7 +28,7 @@ def reload_settings(): - global MSG_OVERHEAD, MAX_MSG_SIZE, HEART_BEAT, HTTP_HEADERS # noqa:PLW0603 + global MSG_OVERHEAD, MAX_MSG_SIZE, HEART_BEAT, HTTP_HEADERS, WS_COMPRESS # noqa:PLW0603 MSG_OVERHEAD = int(os.environ.get("WSLINK_MSG_OVERHEAD", MSG_OVERHEAD)) MAX_MSG_SIZE = int(os.environ.get("WSLINK_MAX_MSG_SIZE", MAX_MSG_SIZE)) @@ -35,6 +36,7 @@ def reload_settings(): os.environ.get("WSLINK_HEART_BEAT", HEART_BEAT or 30) ) # 30 seconds HTTP_HEADERS = os.environ.get("WSLINK_HTTP_HEADERS", HTTP_HEADERS) + WS_COMPRESS = bool(int(os.environ.get("WSLINK_WS_COMPRESS", int(WS_COMPRESS)))) # Allow to skip heart beat if HEART_BEAT < 1: @@ -135,6 +137,11 @@ def runner(self): def site(self): return self._site + @property + def ws_compress(self): + # Explicit config wins, otherwise fall back to WSLINK_WS_COMPRESS + return bool(self.config.get("ws_compress", WS_COMPRESS)) + def get_port(self): """Return the actual port used by the server""" return self.runner.addresses[0][1] @@ -252,7 +259,9 @@ async def disconnectClients(self): async def handleWsRequest(self, request): client_id = str(uuid.uuid4()).replace("-", "") current_ws = aiohttp_web.WebSocketResponse( - max_msg_size=MSG_OVERHEAD + MAX_MSG_SIZE, heartbeat=HEART_BEAT + max_msg_size=MSG_OVERHEAD + MAX_MSG_SIZE, + heartbeat=HEART_BEAT, + compress=self.web_app.ws_compress, ) self.connections[client_id] = current_ws diff --git a/src/wslink/server.py b/src/wslink/server.py index 997b747..e3f8f2e 100644 --- a/src/wslink/server.py +++ b/src/wslink/server.py @@ -86,6 +86,12 @@ def add_arguments(parser): dest="nows", help="If provided, disables the websocket endpoint", ) + parser.add_argument( + "--no-ws-compress", + action="store_true", + dest="nowscompress", + help="If provided, disables websocket compression (permessage-deflate)", + ) parser.add_argument( "--fs-endpoints", default="", @@ -219,6 +225,10 @@ def start_webserver( server_config["ws"] = {} server_config["ws"][options.ws] = wslinkServer + # Only override WSLINK_WS_COMPRESS when explicitly requested + if options.nowscompress: + server_config["ws_compress"] = False + # Configure default static route if --content requested if len(options.content) > 0: server_config["static"] = {} diff --git a/tests/test_aiohttp_routes.py b/tests/test_aiohttp_routes.py index d7accd6..eb092a2 100644 --- a/tests/test_aiohttp_routes.py +++ b/tests/test_aiohttp_routes.py @@ -5,6 +5,7 @@ from aiohttp import web from wslink.backends.aiohttp import WebAppServer +from wslink.websocket import ServerProtocol def _text_handler(text): @@ -89,3 +90,45 @@ async def test_routes_added_before_start_take_precedence(www): finally: await server.stop() await task + + +async def _ws_compress(port): + url = f"ws://localhost:{port}/ws" + async with ( + aiohttp.ClientSession() as session, + session.ws_connect(url, compress=15) as ws, + ): + # Negotiated permessage-deflate window bits (0 => disabled) + compress = ws.compress + await ws.close() + return compress + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("env", "config", "expected"), + [ + ("1", {}, True), + ("0", {}, False), + ("1", {"ws_compress": False}, False), + ("0", {"ws_compress": True}, True), + ], +) +async def test_ws_compression(monkeypatch, env, config, expected): + monkeypatch.setenv("WSLINK_WS_COMPRESS", env) + server = WebAppServer( + { + "host": "localhost", + "port": 0, + "timeout": 0, + "handle_signals": False, + "ws": {"ws": ServerProtocol()}, + **config, + } + ) + port, task = await _start(server) + try: + assert bool(await _ws_compress(port)) == expected + finally: + await server.stop() + await task