websockets: implement notification masking

Signed-off-by:  Eric Callahan <arksine.code@gmail.com>
This commit is contained in:
Eric Callahan 2022-04-07 18:01:50 -04:00
parent d5a55047ab
commit 1133c296b1
No known key found for this signature in database
GPG Key ID: 7027245FBBDDF59A
1 changed files with 11 additions and 4 deletions

View File

@ -20,6 +20,7 @@ from typing import (
Optional,
Callable,
Coroutine,
Tuple,
Type,
TypeVar,
Union,
@ -323,7 +324,7 @@ class WebsocketManager(APITransport):
notify_name = event_name.split(':')[-1]
def notify_handler(*args):
self.notify_websockets(notify_name, *args)
self.notify_websockets(notify_name, args)
self.server.register_event_handler(
event_name, notify_handler)
@ -450,12 +451,15 @@ class WebsocketManager(APITransport):
def notify_websockets(self,
name: str,
data: Any = SENTINEL
data: Union[List, Tuple] = [],
mask: List[int] = []
) -> None:
msg: Dict[str, Any] = {'jsonrpc': "2.0", 'method': "notify_" + name}
if data != SENTINEL:
msg['params'] = [data]
if data:
msg['params'] = data
for ws in list(self.websockets.values()):
if ws.uid in mask:
continue
ws.queue_message(msg)
def get_count(self) -> int:
@ -598,6 +602,9 @@ class WebSocket(WebSocketHandler, Subscribable):
self.queue_message(msg)
return fut
def send_notification(self, name: str, data: List) -> None:
self.wsm.notify_websockets(name, data, [self._uid])
def resolve_pending_response(
self, response_id: int, result: Any
) -> bool: