[IMP] bus, tools: add rate limiting to incoming websocket messages

This commit is part of the websocket integration in Odoo.
It focuses on implementing rate limiting to the incoming
websocket messages. When opening a connection, a burst of messages
is allowed. When requests are received too fast, an exception is
raised and the websocket is closed.

Two config parameters are added to customize the rate limtier
behavior:
   - websocket_rate_limit_delay: Integer specifying the seconds that
     should space out two requests.
   - websocket_rate_limit_burst: Integer specifying how many websocket
     frames can be accepted in excess of the specified rate.

Part-of: odoo/odoo#75510

addons/bus/tests/__init__.py

@@ -2,3 +2,4 @@ from . import common
 from . import test_assetsbundle
 from . import test_health
 from . import test_websocket_caryall
+from . import test_websocket_rate_limiting

addons/bus/tests/common.py

@@ -12,7 +12,7 @@ except ImportError:
 
 import odoo.tools
 from odoo.tests import HOST, common
-from odoo.addons.bus.websocket import CloseCode, WebsocketConnectionHandler
+from ..websocket import CloseCode, WebsocketConnectionHandler
 
 
 class WebsocketCase(common.HttpCase):

addons/bus/tests/test_websocket_rate_limiting.py

@@ -0,0 +1,62 @@
+# Part of Odoo. See LICENSE file for full copyright and licensing details.
+
+import json
+import time
+from websocket._exceptions import WebSocketProtocolException
+
+from odoo.tests import common
+from .common import WebsocketCase
+from ..websocket import CloseCode, Websocket
+
+@common.tagged('post_install', '-at_install')
+class TestWebsocketRateLimiting(WebsocketCase):
+    def test_rate_limiting_base_ok(self):
+        ws = self.websocket_connect()
+
+        for _ in range(Websocket.RL_BURST + 1):
+            ws.send(json.dumps({'event_name': 'test_rate_limiting'}))
+            time.sleep(Websocket.RL_DELAY)
+
+    def test_rate_limiting_base_ko(self):
+        ws = self.websocket_connect()
+
+        # Websocket client's close codes are not up to date. Indeed, the
+        # 1013 close code results in a protocol exception while it is a
+        # valid, registered close code ("TRY LATER") :
+        # https://www.iana.org/assignments/websocket/websocket.xhtml
+        with self.assertRaises(WebSocketProtocolException) as cm:
+            for _ in range(Websocket.RL_BURST + 1):
+                ws.send(json.dumps({'event_name': 'test_rate_limiting'}))
+            self.assert_close_with_code(ws, CloseCode.TRY_LATER)
+        self.assertEqual(str(cm.exception), 'Invalid close opcode.')
+
+    def test_rate_limiting_opening_burst(self):
+        ws = self.websocket_connect()
+
+        # first RL_BURST requests are accepted.
+        for _ in range(Websocket.RL_BURST):
+            ws.send(json.dumps({'event_name': 'test_rate_limiting'}))
+
+        # sending at a correct rate after burst should be accepted.
+        for _ in range(2):
+            time.sleep(Websocket.RL_DELAY)
+            ws.send(json.dumps({'event_name': 'test_rate_limiting'}))
+
+    def test_rate_limiting_start_ok_end_ko(self):
+        ws = self.websocket_connect()
+
+        # first requests are legit and should be accepted
+        for _ in range(Websocket.RL_BURST + 1):
+            ws.send(json.dumps({'event_name': 'test_rate_limiting'}))
+            time.sleep(Websocket.RL_DELAY)
+
+        # Websocket client's close codes are not up to date. Indeed, the
+        # 1013 close code results in a protocol exception while it is a
+        # valid, registered close code ("TRY LATER") :
+        # https://www.iana.org/assignments/websocket/websocket.xhtml
+        with self.assertRaises(WebSocketProtocolException) as cm:
+            # those requests are illicit and should not be accepted.
+            for _ in range(Websocket.RL_BURST * 2):
+                ws.send(json.dumps({'event_name': 'test_rate_limiting'}))
+            self.assert_close_with_code(ws, CloseCode.TRY_LATER)
+        self.assertEqual(str(cm.exception), 'Invalid close opcode.')

addons/bus/websocket.py

@@ -9,6 +9,7 @@ import struct
 import selectors
 import threading
 import time
+from collections import deque
 from contextlib import suppress
 from enum import IntEnum
 from itertools import count
@@ -95,6 +96,13 @@ class ProtocolError(WebsocketException):
     """
 
 
+class RateLimitExceededException(Exception):
+    """
+    Raised when a client exceeds the number of request in a given
+    time.
+    """
+
+
 # ------------------------------------------------------
 # WEBSOCKET
 # ------------------------------------------------------
@@ -201,6 +209,10 @@ class Websocket:
     # or received within CONNECTION_TIMEOUT - 15 seconds.
     CONNECTION_TIMEOUT = 60
     INACTIVITY_TIMEOUT = CONNECTION_TIMEOUT - 15
+    # How many requests can be made in excess of the given rate.
+    RL_BURST = int(config['websocket_rate_limit_burst'])
+    # How many seconds between each request.
+    RL_DELAY = float(config['websocket_rate_limit_delay'])
 
     def __init__(self, socket):
         self._socket = socket
@@ -211,6 +223,8 @@ class Websocket:
         self._selector.register(self._socket, selectors.EVENT_READ)
         self._selector.register(self._outgoing_frame_queue, selectors.EVENT_READ)
         self._timeout_manager = TimeoutManager()
+        # Used for rate limiting.
+        self._incoming_frame_timestamps = deque(maxlen=type(self).RL_BURST)
         self.state = ConnectionState.OPEN
         type(self)._instances.add(self)
 
@@ -318,6 +332,7 @@ class Websocket:
             payload[3::4] = payload[3::4].translate(d)
             return payload
 
+        self._limit_rate()
         first_byte, second_byte = recv_bytes(2)
         fin, rsv1, rsv2, rsv3 = (is_bit_set(first_byte, n) for n in range(4))
         try:
@@ -503,11 +518,28 @@ class Websocket:
             code = CloseCode.INCONSISTENT_DATA
         elif isinstance(exc, PayloadTooLargeException):
             code = CloseCode.MESSAGE_TOO_BIG
+        elif isinstance(exc, RateLimitExceededException):
+            code = CloseCode.TRY_LATER
         if code is CloseCode.SERVER_ERROR:
             reason = None
             _logger.error(exc, exc_info=True)
         self.disconnect(code, reason)
 
+    def _limit_rate(self):
+        """
+        This method is a simple rate limiter designed not to allow
+        more than one request by `RL_DELAY` seconds. `RL_BURST` specify
+        how many requests can be made in excess of the given rate at the
+        begining. When requests are received too fast, raises the
+        `RateLimitExceededException`.
+        """
+        now = time.time()
+        if len(self._incoming_frame_timestamps) >= type(self).RL_BURST:
+            elapsed_time = now - self._incoming_frame_timestamps[0]
+            if elapsed_time < type(self).RL_DELAY * type(self).RL_BURST:
+                raise RateLimitExceededException()
+        self._incoming_frame_timestamps.append(now)
+
     @classmethod
     def _kick_all(cls):
         """ Disconnect all the websocket instances. """

odoo/tools/config.py

@@ -78,6 +78,8 @@ class configmanager(object):
             'reportgz': False,
             'root_path': None,
             'websocket_keep_alive_timeout': 600,
+            'websocket_rate_limit_burst': 10,
+            'websocket_rate_limit_delay': 0.2,
         }
 
         # Not exposed in the configuration file.