Merge pull request #38 from avatao-content/rate_limit_magic

Rate limit magic

lib/tfw/components/fsm_managing_event_handler.py

@@ -2,18 +2,19 @@
 # All Rights Reserved. See LICENSE file for details.
 
 from tfw import EventHandlerBase
-from tfw.crypto import KeyManager, sign_message
+from tfw.crypto import KeyManager, sign_message, verify_message
 from tfw.config.logs import logging
 
 LOG = logging.getLogger(__name__)
 
 
 class FSMManagingEventHandler(EventHandlerBase):
-    def __init__(self, key, fsm_type):
+    def __init__(self, key, fsm_type, require_signature=False):
         super().__init__(key)
         self.fsm = fsm_type()
         self._fsm_updater = FSMUpdater(self.fsm)
         self.auth_key = KeyManager().auth_key
+        self._require_signature = require_signature
 
         self.command_handlers = {
             'trigger': self.handle_trigger,
@@ -22,8 +23,8 @@ class FSMManagingEventHandler(EventHandlerBase):
 
     def handle_event(self, message):
         try:
-            data = message['data']
-            message['data'] = self.command_handlers[data['command']](data)
+            message = self.command_handlers[message['data']['command']](message)
+            if message:
                 fsm_update_message = self._fsm_updater.generate_fsm_update()
                 sign_message(self.auth_key, message)
                 sign_message(self.auth_key, fsm_update_message)
@@ -32,14 +33,19 @@ class FSMManagingEventHandler(EventHandlerBase):
         except KeyError:
             LOG.error('IGNORING MESSAGE: Invalid message received: %s', message)
 
-    def handle_trigger(self, data):
-        trigger = data['value']
-        self.fsm.step(trigger)
-        return data
+    def handle_trigger(self, message):
+        trigger = message['data']['value']
+        if self._require_signature:
+            if not verify_message(self.auth_key, message):
+                LOG.error('Ignoring unsigned trigger command: %s', message)
+                return None
+        if self.fsm.step(trigger):
+            return message
+        return None
 
-    def handle_update(self, data):
+    def handle_update(self, message):
         # pylint: disable=no-self-use
-        return data
+        return message
 
 
 class FSMUpdater:

lib/tfw/decorators/rate_limiter.py

@@ -1,25 +1,100 @@
 # Copyright (C) 2018 Avatao.com Innovative Learning Kft.
 # All Rights Reserved. See LICENSE file for details.
 
-from functools import wraps
+from functools import wraps, partial
 from time import time, sleep
 
+from tfw.decorators.lazy_property import lazy_property
+
 
 class RateLimiter:
+    """
+    Decorator class for rate limiting, blocking.
+
+    When applied to a function this decorator will apply rate limiting
+    if the function is invoked more frequently than rate_per_seconds.
+
+    By default rate limiting means sleeping until the next invocation time
+    as per __init__ parameter rate_per_seconds.
+
+    Note that this decorator BLOCKS THE THREAD it is being executed on,
+    so it is only acceptable for stuff running on a separate thread.
+
+    If this is no good for you please refer to AsyncRateLimiter in this module,
+    which is designed not to block and use the IOLoop it is being called from.
+    """
     def __init__(self, rate_per_second):
+        """
+        :param rate_per_second: max frequency the decorated method should be
+                                invoked with
+        """
         self.min_interval = 1 / float(rate_per_second)
+        self.fun = None
         self.last_call = time()
 
+    def action(self, seconds_to_next_call):
+        if seconds_to_next_call:
+            sleep(seconds_to_next_call)
+        self.fun()
+
     def __call__(self, fun):
         @wraps(fun)
         def wrapper(*args, **kwargs):
-            self._limit_rate()
-            fun(*args, **kwargs)
+            self.fun = partial(fun, *args, **kwargs)
+            limit_seconds = self._limit_rate()
+            self.action(limit_seconds)
         return wrapper
 
     def _limit_rate(self):
-        since_last_call = time() - self.last_call
-        to_next_call = self.min_interval - since_last_call
+        seconds_since_last_call = time() - self.last_call
+        seconds_to_next_call = self.min_interval - seconds_since_last_call
+
+        if seconds_to_next_call > 0:
+            return seconds_to_next_call
         self.last_call = time()
-        if to_next_call > 0:
-            sleep(to_next_call)
+        return 0
+
+
+class AsyncRateLimiter(RateLimiter):
+    """
+    Decorator class for rate limiting, non-blocking.
+
+    The semantics of the rate limiting:
+        - unlike RateLimiter this decorator never blocks, instead it adds an async
+          callback version of the decorated function to the IOLoop
+          (to be executed after the rate limiting has expired).
+        - the timing works similarly to RateLimiter
+    """
+    def __init__(self, rate_per_second, ioloop_factory):
+        """
+        :param rate_per_second: max frequency the decorated method should be
+                                invoked with
+        :param ioloop_factory: callable that should return an instance of the
+                               IOLoop of the application
+        """
+        self._ioloop_factory = ioloop_factory
+        self._ioloop = None
+        self._last_callback = None
+
+        self._make_action_thread_safe()
+        super().__init__(rate_per_second=rate_per_second)
+
+    def _make_action_thread_safe(self):
+        self.action = partial(self.ioloop.add_callback, self.action)
+
+    @lazy_property
+    def ioloop(self):
+        return self._ioloop_factory()
+
+    def action(self, seconds_to_next_call):
+        if self._last_callback:
+            self.ioloop.remove_timeout(self._last_callback)
+
+        self._last_callback = self.ioloop.call_later(
+            seconds_to_next_call,
+            self.fun_with_debounce
+        )
+
+    def fun_with_debounce(self):
+        self.last_call = time()
+        self.fun()

lib/tfw/fsm_base.py

@@ -62,5 +62,7 @@ class FSMBase(Machine, CallbackMixin):
             try:
                 self.trigger(trigger)
                 self.trigger_history.append(trigger)
+                return True
             except (AttributeError, MachineError):
                 LOG.debug('FSM failed to execute nonexistent trigger: "%s"', trigger)
+        return False

lib/tfw/networking/server/tfw_server.py

@@ -9,6 +9,7 @@ from tornado.web import Application
 from tfw.networking.event_handlers import ServerUplinkConnector
 from tfw.networking.server import EventHandlerConnector
 from tfw.networking import MessageSender
+from tfw.crypto import KeyManager, verify_message, sign_message
 from tfw.config.logs import logging
 from .zmq_websocket_proxy import ZMQWebSocketProxy
 
@@ -24,6 +25,7 @@ class TFWServer:
     def __init__(self):
         self._event_handler_connector = EventHandlerConnector()
         self._uplink_connector = ServerUplinkConnector()
+        self._auth_key = KeyManager().auth_key
 
         self.application = Application([(
             r'/ws', ZMQWebSocketProxy, {
@@ -37,13 +39,16 @@ class TFWServer:
     def handle_trigger(self, message):
         if 'trigger' in message:
             LOG.debug('Executing handler for trigger "%s"', message.get('trigger', ''))
-            self._uplink_connector.send_to_eventhandler({
+            fsm_eh_command = {
                 'key': 'fsm',
                 'data': {
                     'command': 'trigger',
-                    'value': message.get('trigger', '')
+                    'value': message['trigger']
                 }
-            })
+            }
+            if verify_message(self._auth_key, message):
+                sign_message(self._auth_key, fsm_eh_command)
+            self._uplink_connector.send_to_eventhandler(fsm_eh_command)
 
     def handle_recover(self, message):
         if message['key'] == 'recover':