Increases the maximum websocket message size limit to 500 MB in matlab-proxy and matlab-proxy-manager.

fixes #57

Author: krisctl

matlab_proxy/app.py

@@ -586,8 +586,9 @@ async def matlab_view(req):
         and reqH.get(UPGRADE, "").lower() == "websocket"
         and req.method == "GET"
     ):
-        ws_server = web.WebSocketResponse()
-
+        ws_server = web.WebSocketResponse(
+            max_msg_size=constants.MAX_WEBSOCKET_MESSAGE_SIZE_IN_MB, compress=True
+        )
         await ws_server.prepare(req)
 
         async with aiohttp.ClientSession(
@@ -600,6 +601,8 @@ async def matlab_view(req):
             try:
                 async with client_session.ws_connect(
                     matlab_base_url + req.path_qs,
+                    max_msg_size=constants.MAX_WEBSOCKET_MESSAGE_SIZE_IN_MB,  # max websocket message size from MATLAB to browser
+                    compress=12,  # enable websocket messages compression
                 ) as ws_client:
 
                     async def wsforward(ws_from, ws_to):
@@ -631,13 +634,24 @@ async def wsforward(ws_from, ws_to):
                                 await ws_to.close(
                                     code=ws_to.close_code, message=msg.extra
                                 )
+                            elif mt == aiohttp.WSMsgType.ERROR:
+                                logger.error(f"WebSocket error received: {msg}")
+                                if "exceeds limit" in str(msg.data):
+                                    logger.error(
+                                        f"Message too large: {msg.data}. Please refresh browser tab to reconnect."
+                                    )
+                                break
                             else:
                                 raise ValueError(f"Unexpected message type: {msg}")
 
                     await asyncio.wait(
                         [
-                            asyncio.create_task(wsforward(ws_server, ws_client)),
-                            asyncio.create_task(wsforward(ws_client, ws_server)),
+                            asyncio.create_task(
+                                wsforward(ws_server, ws_client)
+                            ),  # browser to MATLAB
+                            asyncio.create_task(
+                                wsforward(ws_client, ws_server)
+                            ),  # MATLAB to browser
                         ],
                         return_when=asyncio.FIRST_COMPLETED,
                     )

matlab_proxy/constants.py

@@ -6,6 +6,7 @@
 CONNECTOR_SECUREPORT_FILENAME: Final[str] = "connector.securePort"
 VERSION_INFO_FILE_NAME: Final[str] = "VersionInfo.xml"
 MAX_HTTP_REQUEST_SIZE: Final[int] = 500_000_000  # 500MB
+MAX_WEBSOCKET_MESSAGE_SIZE_IN_MB: Final[int] = 500_000_000  # 500MB
 MATLAB_LOGS_FILE_NAME: Final[str] = "matlab_logs.txt"
 USER_CODE_OUTPUT_FILE_NAME: Final[str] = "startup_code_output.txt"
 

matlab_proxy_manager/web/app.py

@@ -14,6 +14,7 @@
 import matlab_proxy.util.mwi.environment_variables as mwi_env
 import matlab_proxy.util.system as mwi_sys
 import matlab_proxy_manager.lib.api as mpm_lib
+import matlab_proxy.constants as mp_constants
 from matlab_proxy_manager.utils import constants, helpers, logger
 from matlab_proxy_manager.utils import environment_variables as mpm_env
 from matlab_proxy_manager.utils.auth import authenticate_access_decorator
@@ -342,7 +343,9 @@ async def _forward_websocket_request(
     Returns:
         web.WebSocketResponse: The response from the backend server
     """
-    ws_server = web.WebSocketResponse()
+    ws_server = web.WebSocketResponse(
+        max_msg_size=mp_constants.MAX_WEBSOCKET_MESSAGE_SIZE_IN_MB, compress=True
+    )
     await ws_server.prepare(req)
 
     async with aiohttp.ClientSession(
@@ -351,7 +354,11 @@ async def _forward_websocket_request(
         connector=aiohttp.TCPConnector(ssl=False),
     ) as client_session:
         try:
-            async with client_session.ws_connect(proxy_url) as ws_client:
+            async with client_session.ws_connect(
+                proxy_url,
+                max_msg_size=mp_constants.MAX_WEBSOCKET_MESSAGE_SIZE_IN_MB,  # max websocket message size from MATLAB to browser
+                compress=12,  # enable websocket messages compression
+            ) as ws_client:
 
                 async def ws_forward(ws_src, ws_dest):
                     async for msg in ws_src:
@@ -389,6 +396,13 @@ async def ws_forward(ws_src, ws_dest):
                             await ws_dest.close(
                                 code=ws_dest.close_code, message=msg.extra
                             )
+                        elif msg_type == aiohttp.WSMsgType.ERROR:
+                            log.error(f"WebSocket error received: {msg}")
+                            if "exceeds limit" in str(msg.data):
+                                log.error(
+                                    f"Message too large: {msg.data}. Please refresh browser tab to reconnect."
+                                )
+                            break
                         else:
                             raise ValueError(f"Unexpected message type: {msg}")