diff --git a/lib/tfw/networking/__init__.py b/lib/tfw/networking/__init__.py index 3dac2c0..6dd15f9 100644 --- a/lib/tfw/networking/__init__.py +++ b/lib/tfw/networking/__init__.py @@ -1,7 +1,7 @@ # Copyright (C) 2018 Avatao.com Innovative Learning Kft. # All Rights Reserved. See LICENSE file for details. -from .serialization import serialize_tfw_msg, deserialize_tfw_msg +from .serialization import serialize_tfw_msg, deserialize_tfw_msg, validate_message from .zmq_connector_base import ZMQConnectorBase # from .controller_connector import ControllerConnector # TODO: readd once controller stuff is resolved from .message_sender import MessageSender diff --git a/lib/tfw/networking/serialization.py b/lib/tfw/networking/serialization.py index f24b38a..0586309 100644 --- a/lib/tfw/networking/serialization.py +++ b/lib/tfw/networking/serialization.py @@ -24,6 +24,12 @@ The purpose of this module is abstracting away this low level behaviour. import json +def validate_message(message): + if 'key' not in message: + return False + return True + + def serialize_tfw_msg(message): return _serialize_all(message['key'], message) diff --git a/lib/tfw/networking/server/zmq_websocket_handler.py b/lib/tfw/networking/server/zmq_websocket_handler.py index 4d0d994..fee2ff1 100644 --- a/lib/tfw/networking/server/zmq_websocket_handler.py +++ b/lib/tfw/networking/server/zmq_websocket_handler.py @@ -6,7 +6,7 @@ from abc import ABC, abstractmethod from tornado.websocket import WebSocketHandler -from tfw.networking import deserialize_tfw_msg +from tfw.networking import deserialize_tfw_msg, validate_message from tfw.config.logs import logging LOG = logging.getLogger(__name__) @@ -32,15 +32,17 @@ class ZMQWebSocketHandler(WebSocketHandler, ABC): keyhandlers = {'mirror': self.mirror} message = deserialize_tfw_msg(*msg_parts) - key = message['key'] LOG.debug('Received on pull socket: %s', message) + if not validate_message(message): + return + self.handle_trigger(message) - if key not in keyhandlers: + if message['key'] not in keyhandlers: for instance in ZMQWebSocketHandler.instances: instance.write_message(message) else: try: - keyhandlers[key](message) + keyhandlers[message['key']](message) except KeyError: LOG.error('Invalid mirror message format! Ignoring.') @@ -50,7 +52,8 @@ class ZMQWebSocketHandler(WebSocketHandler, ABC): def on_message(self, message): LOG.debug('Received on WebSocket: %s', message) - self.send_message(self.make_eventhandler_response(message)) + if validate_message(message): + self.send_message(self.make_eventhandler_response(message)) @abstractmethod def make_eventhandler_response(self, message):