From f8e794a864e2ca31a05423d53e8f066fac608bce Mon Sep 17 00:00:00 2001 From: tsm-odoo Date: Wed, 3 Aug 2022 12:21:05 +0000 Subject: [PATCH] [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 | 1 + addons/bus/tests/common.py | 2 +- .../bus/tests/test_websocket_rate_limiting.py | 62 +++++++++++++++++++ addons/bus/websocket.py | 32 ++++++++++ odoo/tools/config.py | 2 + 5 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 addons/bus/tests/test_websocket_rate_limiting.py diff --git a/addons/bus/tests/__init__.py b/addons/bus/tests/__init__.py index c8f25a068a1..aff0b730980 100644 --- a/addons/bus/tests/__init__.py +++ b/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 diff --git a/addons/bus/tests/common.py b/addons/bus/tests/common.py index 9fdbafbbf67..0e04ae9b30b 100644 --- a/addons/bus/tests/common.py +++ b/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): diff --git a/addons/bus/tests/test_websocket_rate_limiting.py b/addons/bus/tests/test_websocket_rate_limiting.py new file mode 100644 index 00000000000..4c864fe669d --- /dev/null +++ b/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.') diff --git a/addons/bus/websocket.py b/addons/bus/websocket.py index 4fe39af53fd..8f4613efdc7 100644 --- a/addons/bus/websocket.py +++ b/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. """ diff --git a/odoo/tools/config.py b/odoo/tools/config.py index 10da83acc3b..de101967e39 100644 --- a/odoo/tools/config.py +++ b/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.