diff options
| -rw-r--r-- | kvmd/apps/kvmd/server.py | 18 | ||||
| -rw-r--r-- | kvmd/htserver.py | 20 |
2 files changed, 21 insertions, 17 deletions
diff --git a/kvmd/apps/kvmd/server.py b/kvmd/apps/kvmd/server.py index f40381a9..aa643d9f 100644 --- a/kvmd/apps/kvmd/server.py +++ b/kvmd/apps/kvmd/server.py @@ -57,7 +57,7 @@ from ...htserver import make_json_response from ...htserver import make_json_exception from ...htserver import send_ws_event from ...htserver import broadcast_ws_event -from ...htserver import parse_ws_event +from ...htserver import process_ws_messages from ...htserver import HttpServer from ...plugins import BasePlugin @@ -278,21 +278,7 @@ class KvmdServer(HttpServer): # pylint: disable=too-many-arguments,too-many-ins ]) await send_ws_event(ws, "loop", {}) - - async for msg in ws: - if msg.type != aiohttp.web.WSMsgType.TEXT: - break - try: - (event_type, event) = parse_ws_event(msg.data) - except Exception as err: - get_logger(0).error("Can't parse JSON event from websocket: %r", err) - else: - handler = self.__ws_handlers.get(event_type) - if handler: - await handler(ws, event) - else: - get_logger(0).error("Unknown websocket event: %r", msg.data) - + await process_ws_messages(ws, self.__ws_handlers) return ws finally: await self.__remove_ws_client(client) diff --git a/kvmd/htserver.py b/kvmd/htserver.py index 34b38846..00fe24ea 100644 --- a/kvmd/htserver.py +++ b/kvmd/htserver.py @@ -38,6 +38,7 @@ from aiohttp.web import Request from aiohttp.web import Response from aiohttp.web import StreamResponse from aiohttp.web import WebSocketResponse +from aiohttp.web import WSMsgType from aiohttp.web import Application from aiohttp.web import run_app from aiohttp.web import normalize_path_middleware @@ -220,7 +221,7 @@ async def broadcast_ws_event(wss: List[WebSocketResponse], event_type: str, even ], return_exceptions=True) -def parse_ws_event(msg: str) -> Tuple[str, Dict]: +def _parse_ws_event(msg: str) -> Tuple[str, Dict]: data = json.loads(msg) if not isinstance(data, dict): raise RuntimeError("Top-level event structure is not a dict") @@ -233,6 +234,23 @@ def parse_ws_event(msg: str) -> Tuple[str, Dict]: return (event_type, event) +async def process_ws_messages(ws: WebSocketResponse, handlers: Dict[str, Callable]) -> None: + logger = get_logger(1) + async for msg in ws: + if msg.type != WSMsgType.TEXT: + break + try: + (event_type, event) = _parse_ws_event(msg.data) + except Exception as err: + logger.error("Can't parse JSON event from websocket: %r", err) + else: + handler = handlers.get(event_type) + if handler: + await handler(ws, event) + else: + logger.error("Unknown websocket event: %r", msg.data) + + # ===== _REQUEST_AUTH_INFO = "_kvmd_auth_info" |