From 6ae56ac2ccdbb2d1adc6ccc675dc52f1cac8d995 Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Sun, 25 Feb 2024 02:18:40 +0000
Subject: [PATCH] increase max ws response size for instances with large agent
 counts

---
 api/tacticalrmm/core/management/commands/get_mesh_exe_url.py   | 3 ++-
 api/tacticalrmm/core/management/commands/initial_mesh_setup.py | 3 ++-
 api/tacticalrmm/core/mesh_utils.py                             | 3 ++-
 api/tacticalrmm/core/utils.py                                  | 3 ++-
 api/tacticalrmm/tacticalrmm/constants.py                       | 2 ++
 5 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/api/tacticalrmm/core/management/commands/get_mesh_exe_url.py b/api/tacticalrmm/core/management/commands/get_mesh_exe_url.py
index 07a548bf..b1e8944b 100644
--- a/api/tacticalrmm/core/management/commands/get_mesh_exe_url.py
+++ b/api/tacticalrmm/core/management/commands/get_mesh_exe_url.py
@@ -5,13 +5,14 @@ import websockets
 from django.core.management.base import BaseCommand
 
 from core.utils import get_mesh_ws_url
+from tacticalrmm.constants import WS_MAX_SIZE
 
 
 class Command(BaseCommand):
     help = "Sets up initial mesh central configuration"
 
     async def websocket_call(self, uri):
-        async with websockets.connect(uri) as websocket:
+        async with websockets.connect(uri, max_size=WS_MAX_SIZE) as websocket:
             # Get Invitation Link
             await websocket.send(
                 json.dumps(
diff --git a/api/tacticalrmm/core/management/commands/initial_mesh_setup.py b/api/tacticalrmm/core/management/commands/initial_mesh_setup.py
index 3ca36357..3a836c31 100644
--- a/api/tacticalrmm/core/management/commands/initial_mesh_setup.py
+++ b/api/tacticalrmm/core/management/commands/initial_mesh_setup.py
@@ -6,13 +6,14 @@ from django.conf import settings
 from django.core.management.base import BaseCommand
 
 from core.utils import get_core_settings, get_mesh_ws_url
+from tacticalrmm.constants import WS_MAX_SIZE
 
 
 class Command(BaseCommand):
     help = "Sets up initial mesh central configuration"
 
     async def websocket_call(self, uri):
-        async with websockets.connect(uri) as websocket:
+        async with websockets.connect(uri, max_size=WS_MAX_SIZE) as websocket:
             # Get Device groups to see if it exists
             await websocket.send(json.dumps({"action": "meshes"}))
 
diff --git a/api/tacticalrmm/core/mesh_utils.py b/api/tacticalrmm/core/mesh_utils.py
index 81ccc517..2ff62bbb 100644
--- a/api/tacticalrmm/core/mesh_utils.py
+++ b/api/tacticalrmm/core/mesh_utils.py
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any
 import websockets
 
 from accounts.utils import is_superuser
+from tacticalrmm.constants import WS_MAX_SIZE
 from tacticalrmm.helpers import make_random_password
 from tacticalrmm.logger import logger
 
@@ -34,7 +35,7 @@ def mesh_action(
     *, payload: dict[str, Any], uri: str, wait=True
 ) -> dict[str, Any] | None:
     async def _do(payload, uri: str):
-        async with websockets.connect(uri) as ws:
+        async with websockets.connect(uri, max_size=WS_MAX_SIZE) as ws:
             await ws.send(json.dumps(payload))
             if wait:
                 async for message in ws:
diff --git a/api/tacticalrmm/core/utils.py b/api/tacticalrmm/core/utils.py
index 1c7c5da4..df0d2c7e 100644
--- a/api/tacticalrmm/core/utils.py
+++ b/api/tacticalrmm/core/utils.py
@@ -16,6 +16,7 @@ from tacticalrmm.constants import (
     AGENT_TBL_PEND_ACTION_CNT_CACHE_PREFIX,
     CORESETTINGS_CACHE_KEY,
     ROLE_CACHE_PREFIX,
+    WS_MAX_SIZE,
     AgentPlat,
     MeshAgentIdent,
 )
@@ -99,7 +100,7 @@ def get_mesh_ws_url() -> str:
 
 
 async def get_mesh_device_id(uri: str, device_group: str) -> None:
-    async with websockets.connect(uri) as ws:
+    async with websockets.connect(uri, max_size=WS_MAX_SIZE) as ws:
         payload = {"action": "meshes", "responseid": "meshctrl"}
         await ws.send(json.dumps(payload))
 
diff --git a/api/tacticalrmm/tacticalrmm/constants.py b/api/tacticalrmm/tacticalrmm/constants.py
index 4d952af7..4292b9aa 100644
--- a/api/tacticalrmm/tacticalrmm/constants.py
+++ b/api/tacticalrmm/tacticalrmm/constants.py
@@ -32,6 +32,8 @@ AGENT_OUTAGES_LOCK = "agent-outages-task-lock-key"
 ORPHANED_WIN_TASK_LOCK = "orphaned-win-task-lock-key"
 SYNC_MESH_PERMS_TASK_LOCK = "sync-mesh-perms-lock-key"
 
+WS_MAX_SIZE = 100 * 2**20
+
 
 class GoArch(models.TextChoices):
     AMD64 = "amd64", "amd64"