diff --git a/moonraker/app.py b/moonraker/app.py index 1b40dc3..382c439 100644 --- a/moonraker/app.py +++ b/moonraker/app.py @@ -58,6 +58,7 @@ if TYPE_CHECKING: # 50 MiB Max Standard Body Size MAX_BODY_SIZE = 50 * 1024 * 1024 +MAX_WS_CONNS_DEFAULT = 50 EXCLUDED_ARGS = ["_", "token", "access_token", "connection_id"] AUTHORIZED_EXTS = [".png", ".jpg"] DEFAULT_KLIPPY_LOG_PATH = "/tmp/klippy.log" @@ -169,6 +170,9 @@ class MoonrakerApp: self.registered_base_handlers: List[str] = [] self.max_upload_size = config.getint('max_upload_size', 1024) self.max_upload_size *= 1024 * 1024 + max_ws_conns = config.getint( + 'max_websocket_connections', MAX_WS_CONNS_DEFAULT + ) # SSL config self.cert_path: pathlib.Path = self._get_path_option( @@ -193,6 +197,7 @@ class MoonrakerApp: 'websocket_ping_interval': 10, 'websocket_ping_timeout': 30, 'server': self.server, + 'max_websocket_connections': max_ws_conns, 'default_handler_class': AuthorizedErrorHandler, 'default_handler_args': {}, 'log_function': self.log_request, diff --git a/moonraker/websockets.py b/moonraker/websockets.py index 892806a..bb34d7d 100644 --- a/moonraker/websockets.py +++ b/moonraker/websockets.py @@ -650,6 +650,8 @@ class BaseSocketClient(Subscribable): raise NotImplementedError("Children must implement close_socket()") class WebSocket(WebSocketHandler, BaseSocketClient): + connection_count: int = 0 + def initialize(self) -> None: self.on_create(self.settings['server']) self.ip_addr: str = self.request.remote_ip or "" @@ -663,6 +665,7 @@ class WebSocket(WebSocketHandler, BaseSocketClient): return self._user_info def open(self, *args, **kwargs) -> None: + self.__class__.connection_count += 1 self.set_nodelay(True) self._connected_time = self.eventloop.get_loop_time() agent = self.request.headers.get("User-Agent", "") @@ -686,6 +689,7 @@ class WebSocket(WebSocketHandler, BaseSocketClient): def on_close(self) -> None: self.is_closed = True + self.__class__.connection_count -= 1 self.message_buf = [] now = self.eventloop.get_loop_time() pong_elapsed = now - self.last_pong_time @@ -725,7 +729,12 @@ class WebSocket(WebSocketHandler, BaseSocketClient): return True # Check Authorized User - def prepare(self): + def prepare(self) -> None: + max_conns = self.settings["max_websocket_connections"] + if self.__class__.connection_count >= max_conns: + raise self.server.error( + "Maximum Number of Websocket Connections Reached" + ) auth: AuthComp = self.server.lookup_component('authorization', None) if auth is not None: try: