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.