From 762b087aeeaea6e9bbea98795294503f21856ae2 Mon Sep 17 00:00:00 2001 From: "Benjamin Hanquin (beha)" Date: Fri, 10 Nov 2023 10:00:42 +0100 Subject: [PATCH] [IMP] bus: privatize internal attributes Internal attributes of our Websocket class are not supposed to be overridden or even used by developers in business modules. Ensure that by name mangling these attributes, to make the use of them harder and make understand those attributes are meant for the internal use within the Websocket class only. Part-of: odoo/odoo#75510 --- addons/bus/tests/test_websocket_caryall.py | 11 +-- addons/bus/websocket.py | 87 +++++++++++----------- 2 files changed, 51 insertions(+), 47 deletions(-) diff --git a/addons/bus/tests/test_websocket_caryall.py b/addons/bus/tests/test_websocket_caryall.py index 5c8e70dc7d3..ab913ba8c2e 100644 --- a/addons/bus/tests/test_websocket_caryall.py +++ b/addons/bus/tests/test_websocket_caryall.py @@ -17,6 +17,7 @@ except ImportError: from odoo.api import Environment from odoo.tests import common, new_test_user from .common import WebsocketCase +from .. import websocket as websocket_module from ..models.bus import dispatch from ..websocket import ( CloseCode, @@ -32,7 +33,7 @@ from ..websocket import ( class TestWebsocketCaryall(WebsocketCase): def test_lifecycle_hooks(self): events = [] - with patch.object(Websocket, '_event_callbacks', defaultdict(set)): + with patch.object(Websocket, '_Websocket__event_callbacks', defaultdict(set)): @Websocket.onopen def onopen(env, websocket): # pylint: disable=unused-variable self.assertIsInstance(env, Environment) @@ -51,10 +52,10 @@ class TestWebsocketCaryall(WebsocketCase): self.assertEqual(events, ['open', 'close']) def test_instances_weak_set(self): - with patch.object(Websocket, "_instances", WeakSet()): + with patch.object(websocket_module, "_websocket_instances", WeakSet()): first_ws = self.websocket_connect() second_ws = self.websocket_connect() - self.assertEqual(len(Websocket._instances), 2) + self.assertEqual(len(websocket_module._websocket_instances), 2) first_ws.close(CloseCode.CLEAN) second_ws.close(CloseCode.CLEAN) self.wait_remaining_websocket_connections() @@ -62,7 +63,7 @@ class TestWebsocketCaryall(WebsocketCase): # collected. Stop it now. self._serve_forever_patch.stop() gc.collect() - self.assertEqual(len(Websocket._instances), 0) + self.assertEqual(len(websocket_module._websocket_instances), 0) def test_timeout_manager_no_response_timeout(self): with freeze_time('2022-08-19') as frozen_time: @@ -202,7 +203,7 @@ class TestWebsocketCaryall(WebsocketCase): self.assertEqual(mock.call_args[0][2], client_last_notification_id) def test_no_cursor_when_no_callback_for_lifecycle_event(self): - with patch.object(Websocket, '_event_callbacks', defaultdict(set)): + with patch.object(Websocket, '_Websocket__event_callbacks', defaultdict(set)): with patch('odoo.addons.bus.websocket.acquire_cursor') as mock: self.websocket_connect() self.assertFalse(mock.called) diff --git a/addons/bus/websocket.py b/addons/bus/websocket.py index d7654f7c36e..58164ec157d 100644 --- a/addons/bus/websocket.py +++ b/addons/bus/websocket.py @@ -210,9 +210,11 @@ class CloseFrame(Frame): super().__init__(Opcode.CLOSE, payload) +_websocket_instances = WeakSet() + + class Websocket: - _instances = WeakSet() - _event_callbacks = defaultdict(set) + __event_callbacks = defaultdict(set) # Maximum size for a message in bytes, whether it is sent as one # frame or many fragmented ones. MESSAGE_MAX_SIZE = 2 ** 20 @@ -229,7 +231,8 @@ class Websocket: def __init__(self, sock, session): # Session linked to the current websocket connection. self._session = session - self._socket = sock + self._db = session.db + self.__socket = sock self._close_sent = False self._close_received = False self._timeout_manager = TimeoutManager() @@ -237,19 +240,19 @@ class Websocket: self._incoming_frame_timestamps = deque(maxlen=type(self).RL_BURST) # Used to notify the websocket that bus notifications are # available. - self._notif_sock_w, self._notif_sock_r = socket.socketpair() + self.__notif_sock_w, self.__notif_sock_r = socket.socketpair() self._channels = set() self._last_notif_sent_id = 0 # Websocket start up - self._selector = ( + self.__selector = ( selectors.PollSelector() if odoo.evented and hasattr(selectors, 'PollSelector') else selectors.DefaultSelector() ) - self._selector.register(self._socket, selectors.EVENT_READ) - self._selector.register(self._notif_sock_r, selectors.EVENT_READ) + self.__selector.register(self.__socket, selectors.EVENT_READ) + self.__selector.register(self.__notif_sock_r, selectors.EVENT_READ) self.state = ConnectionState.OPEN - type(self)._instances.add(self) + _websocket_instances.add(self) self._trigger_lifecycle_event(LifecycleEvent.OPEN) # ------------------------------------------------------ @@ -261,7 +264,7 @@ class Websocket: try: readables = { selector_key[0].fileobj for selector_key in - self._selector.select(type(self).INACTIVITY_TIMEOUT) + self.__selector.select(type(self).INACTIVITY_TIMEOUT) } if self._timeout_manager.has_timed_out() and self.state is ConnectionState.OPEN: self.disconnect( @@ -273,9 +276,9 @@ class Websocket: if not readables: self._send_ping_frame() continue - if self._notif_sock_r in readables: + if self.__notif_sock_r in readables: self._dispatch_bus_notifications() - if self._socket in readables: + if self.__socket in readables: message = self._process_next_message() if message is not None: yield message @@ -298,12 +301,12 @@ class Websocket: @classmethod def onopen(cls, func): - cls._event_callbacks[LifecycleEvent.OPEN].add(func) + cls.__event_callbacks[LifecycleEvent.OPEN].add(func) return func @classmethod def onclose(cls, func): - cls._event_callbacks[LifecycleEvent.CLOSE].add(func) + cls.__event_callbacks[LifecycleEvent.CLOSE].add(func) return func def subscribe(self, channels, last): @@ -324,11 +327,11 @@ class Websocket: return readables = { selector_key[0].fileobj for selector_key in - self._selector.select(0) + self.__selector.select(0) } - if self._notif_sock_r not in readables: + if self.__notif_sock_r not in readables: # Send a random bit to mark the socket as readable. - self._notif_sock_w.send(b'x') + self.__notif_sock_w.send(b'x') # ------------------------------------------------------ # PRIVATE METHODS @@ -356,7 +359,7 @@ class Websocket: """ Pull n bytes from the socket """ data = bytearray() while len(data) < n: - received_data = self._socket.recv(n - len(data)) + received_data = self.__socket.recv(n - len(data)) if not received_data: raise ConnectionClosed() data.extend(received_data) @@ -498,7 +501,7 @@ class Websocket: struct.pack('!BBQ', first_byte, 127, payload_length) ) output.extend(frame.payload) - self._socket.sendall(output) + self.__socket.sendall(output) self._timeout_manager.acknowledge_frame_sent(frame) if not isinstance(frame, CloseFrame): return @@ -508,7 +511,7 @@ class Websocket: return self._terminate() # After sending a control frame indicating the connection # should be closed, a peer does not send any further data. - self._selector.unregister(self._notif_sock_r) + self.__selector.unregister(self.__notif_sock_r) def _send_close_frame(self, code, reason=None): """ Send a close frame. """ @@ -525,17 +528,17 @@ class Websocket: def _terminate(self): """ Close the underlying TCP socket. """ with suppress(OSError, TimeoutError): - self._socket.shutdown(socket.SHUT_WR) + self.__socket.shutdown(socket.SHUT_WR) # Call recv until obtaining a return value of 0 indicating # the other end has performed an orderly shutdown. A timeout # is set to ensure the connection will be closed even if # the other end does not close the socket properly. - self._socket.settimeout(1) - while self._socket.recv(4096): + self.__socket.settimeout(1) + while self.__socket.recv(4096): pass - self._selector.unregister(self._socket) - self._selector.close() - self._socket.close() + self.__selector.unregister(self.__socket) + self.__selector.close() + self.__socket.close() self.state = ConnectionState.CLOSED dispatch.unsubscribe(self) self._trigger_lifecycle_event(LifecycleEvent.CLOSE) @@ -602,24 +605,17 @@ class Websocket: raise RateLimitExceededException() self._incoming_frame_timestamps.append(now) - @classmethod - def _kick_all(cls): - """ Disconnect all the websocket instances. """ - for websocket in cls._instances: - if websocket.state is ConnectionState.OPEN: - websocket.disconnect(CloseCode.GOING_AWAY) - def _trigger_lifecycle_event(self, event_type): """ Trigger a lifecycle event that is, call every function registered for this event type. Every callback is given both the environment and the related websocket. """ - if not type(self)._event_callbacks[event_type]: + if not type(self).__event_callbacks[event_type]: return - with closing(acquire_cursor(self._session.db)) as cr: + with closing(acquire_cursor(self._db)) as cr: env = api.Environment(cr, self._session.uid, self._session.context) - for callback in type(self)._event_callbacks[event_type]: + for callback in type(self).__event_callbacks[event_type]: try: service_model.retrying(functools.partial(callback, env, self), env) except Exception: @@ -644,7 +640,7 @@ class Websocket: if session.uid is not None and not check_session(session, env): raise SessionExpiredException() # Mark the notification request as processed. - self._notif_sock_r.recv(1) + self.__notif_sock_r.recv(1) notifications = env['bus.bus']._poll(self._channels, self._last_notif_sent_id) if not notifications: return @@ -838,11 +834,11 @@ class WebsocketConnectionHandler: try: response = cls._get_handshake_response(request.httprequest.headers) socket = request.httprequest.environ['socket'] - response.call_on_close(functools.partial( - cls._serve_forever, - Websocket(socket, request.session), - request.db, - request.httprequest + session, db, httprequest = request.session, request.db, request.httprequest + response.call_on_close(lambda: cls._serve_forever( + Websocket(socket, session), + db, + httprequest, )) # Force save the session. Session must be persisted to handle # WebSocket authentication. @@ -943,4 +939,11 @@ class WebsocketConnectionHandler: _logger.exception("Exception occurred during websocket request handling") -CommonServer.on_stop(Websocket._kick_all) +def _kick_all(): + """ Disconnect all the websocket instances. """ + for websocket in _websocket_instances: + if websocket.state is ConnectionState.OPEN: + websocket.disconnect(CloseCode.GOING_AWAY) + + +CommonServer.on_stop(_kick_all)