feat: validate server message (#345)

* chore: introduce `ServerMessage` and `ClientMessage` types

`ServerMessage` represents all the different message types that a
server might send to a client
`ClientMessage` represents all the different message types that a
client might send to a server

these are there to validate the payload types, enforcing the
correctness of the payloads

* fix: improve messages schema

now all of messages should be valid

* feat: introduce `ServerMessage` parsing

now, instead of parsing and dispatching the messages inside the
`AsyncRealtimeChannel._trigger` function, first parse the message
using pydantic's `ServerMessageAdapter` type adapter, and then
dispatch the callbacks inside `_handle_message`, based on which
message was parsed.

this required the removal of the `_on` function, which was used to
both define callbacks to messages, but also to define callbacks to
event refs, which was very confusing. this behavior was used precisely
to define callbacks that are called on acks.

now, AsyncPush'es are saved in
AsyncRealtimeChannel.messages_waiting_for_ack, and later the callbacks
registered by `AsyncPush.receive` (now saved directly in the push
object) are called when a `phx_reply` type message is received.

some logic was introduce to avoid this message buffer from growing
forever indefinetly, by removing these messages from the dictionary
once the callback is called, or once the timeout is reached.

Author: o-santi

realtime/_async/channel.py

@@ -7,11 +7,12 @@
 from urllib.parse import urlencode, urlparse, urlunparse
 
 import websockets
+from pydantic import ValidationError
 from websockets import connect
 from websockets.asyncio.client import ClientConnection
 
 from ..exceptions import NotConnectedError
-from ..message import Message
+from ..message import Message, ServerMessageAdapter
 from ..transformers import http_endpoint_url
 from ..types import (
     DEFAULT_HEARTBEAT_INTERVAL,
@@ -99,13 +100,14 @@ async def _listen(self) -> None:
             async for msg in self._ws_connection:
                 logger.info(f"receive: {msg!r}")
 
-                message = Message.model_validate_json(msg)
-                channel = self.channels.get(message.topic)
-
-                if channel:
-                    channel._trigger(
-                        message.event, dict(**message.payload), message.ref
-                    )
+                try:
+                    message = ServerMessageAdapter.validate_json(msg)
+                except ValidationError as e:
+                    logger.error(f"Unrecognized message format {msg!r}\n{e}")
+                    continue
+                logger.info(f"parsed message as {message}")
+                if channel := self.channels.get(message.topic):
+                    channel._handle_message(message)
         except websockets.exceptions.ConnectionClosedError as e:
             await self._on_connect_error(e)
 

realtime/_async/client.py

@@ -21,26 +21,11 @@
 
 
 class AsyncRealtimePresence:
-    def __init__(self, channel, opts: Optional[PresenceOpts] = None):
-        self.channel = channel
+    def __init__(self):
         self.state: RealtimePresenceState = {}
-        self.pending_diffs: List[RawPresenceDiff] = []
-        self.join_ref: Optional[str] = None
         self.on_join_callback: Optional[PresenceOnJoinCallback] = None
         self.on_leave_callback: Optional[PresenceOnLeaveCallback] = None
         self.on_sync_callback: Optional[Callable[[], None]] = None
-        self.on_auth_success_callback: Optional[Callable[[], None]] = None
-        self.on_auth_failure_callback: Optional[Callable[[], None]] = None
-        # Initialize with default events if not provided
-        events = (
-            opts.events
-            if opts
-            else PresenceEvents(state="presence_state", diff="presence_diff")
-        )
-        # Set up event listeners for presence state and diff
-        self.channel._on(events.state, callback=self._on_state_event)
-        self.channel._on(events.diff, callback=self._on_diff_event)
-        self.channel._on("phx_auth", callback=self._on_auth_event)
 
     def on_join(self, callback: PresenceOnJoinCallback):
         self.on_join_callback = callback
@@ -51,67 +36,43 @@ def on_leave(self, callback: PresenceOnLeaveCallback):
     def on_sync(self, callback: Callable[[], None]):
         self.on_sync_callback = callback
 
-    def on_auth_success(self, callback: Callable[[], None]):
-        self.on_auth_success_callback = callback
+    def _on_state_event(self, payload: RawPresenceState):
+        state = AsyncRealtimePresence._transform_state(payload)
+        self.state = self._sync_state(state)
 
-    def on_auth_failure(self, callback: Callable[[], None]):
-        self.on_auth_failure_callback = callback
-
-    def _on_state_event(self, payload: RawPresenceState, *args):
-        self.join_ref = self.channel.join_ref
-        self.state = self._sync_state(self.state, payload)
-
-        for diff in self.pending_diffs:
-            self.state = self._sync_diff(self.state, diff)
-        self.pending_diffs = []
         if self.on_sync_callback:
             self.on_sync_callback()
 
-    def _on_diff_event(self, payload: RawPresenceDiff, *args):
-        if self.in_pending_sync_state():
-            self.pending_diffs.append(payload)
-        else:
-            self.state = self._sync_diff(self.state, payload)
-            if self.on_sync_callback:
-                self.on_sync_callback()
-
-    def _on_auth_event(self, payload: Dict[str, Any], *args):
-        if payload.get("status") == "ok":
-            if self.on_auth_success_callback:
-                self.on_auth_success_callback()
-        else:
-            if self.on_auth_failure_callback:
-                self.on_auth_failure_callback()
+    def _on_diff_event(self, payload: RawPresenceDiff):
+        joins = AsyncRealtimePresence._transform_state(payload["joins"])
+        leaves = AsyncRealtimePresence._transform_state(payload["leaves"])
+        self.state = self._sync_diff(joins, leaves)
+        if self.on_sync_callback:
+            self.on_sync_callback()
 
     def _sync_state(
         self,
-        current_state: RealtimePresenceState,
-        new_state: Union[RawPresenceState, RealtimePresenceState],
+        new_state: RealtimePresenceState,
     ) -> RealtimePresenceState:
-        state = {key: list(value) for key, value in current_state.items()}
-        transformed_state = AsyncRealtimePresence._transform_state(new_state)
-
-        joins: Dict[str, Any] = {}
-        leaves: Dict[str, Any] = {
-            k: v for k, v in state.items() if k not in transformed_state
-        }
+        joins = {}
+        leaves = {k: v for k, v in self.state.items() if k not in new_state}
 
-        for key, value in transformed_state.items():
-            current_presences = state.get(key, [])
+        for key, value in new_state.items():
+            current_presences = self.state.get(key, [])
 
             if len(current_presences) > 0:
-                new_presence_refs = {presence.get("presence_ref") for presence in value}
+                new_presence_refs = {presence["presence_ref"] for presence in value}
                 cur_presence_refs = {
-                    presence.get("presence_ref") for presence in current_presences
+                    presence["presence_ref"] for presence in current_presences
                 }
 
                 joined_presences = [
-                    p for p in value if p.get("presence_ref") not in cur_presence_refs
+                    p for p in value if p["presence_ref"] not in cur_presence_refs
                 ]
                 left_presences = [
                     p
                     for p in current_presences
-                    if p.get("presence_ref") not in new_presence_refs
+                    if p["presence_ref"] not in new_presence_refs
                 ]
 
                 if joined_presences:
@@ -121,19 +82,14 @@ def _sync_state(
             else:
                 joins[key] = value
 
-        return self._sync_diff(state, {"joins": joins, "leaves": leaves})
+        return self._sync_diff(joins, leaves)
 
     def _sync_diff(
-        self,
-        state: RealtimePresenceState,
-        diff: Union[RawPresenceDiff, PresenceDiff],
+        self, joins: RealtimePresenceState, leaves: RealtimePresenceState
     ) -> RealtimePresenceState:
-        joins = AsyncRealtimePresence._transform_state(diff.get("joins", {}))
-        leaves = AsyncRealtimePresence._transform_state(diff.get("leaves", {}))
-
         for key, new_presences in joins.items():
-            current_presences = state.get(key, [])
-            state[key] = new_presences
+            current_presences = self.state.get(key, [])
+            self.state[key] = new_presences
 
             if len(current_presences) > 0:
                 joined_presence_refs = {
@@ -144,13 +100,13 @@ def _sync_diff(
                     for presence in current_presences
                     if presence.get("presence_ref") not in joined_presence_refs
                 )
-                state[key] = cur_presences + state[key]
+                self.state[key] = cur_presences + self.state[key]
 
             if self.on_join_callback:
                 self.on_join_callback(key, current_presences, new_presences)
 
         for key, left_presences in leaves.items():
-            current_presences = state.get(key, [])
+            current_presences = self.state.get(key, [])
 
             if len(current_presences) == 0:
                 break
@@ -163,22 +119,19 @@ def _sync_diff(
                 for presence in current_presences
                 if presence.get("presence_ref") not in presence_refs_to_remove
             ]
-            state[key] = current_presences
+            self.state[key] = current_presences
 
             if self.on_leave_callback:
                 self.on_leave_callback(key, current_presences, left_presences)
 
             if len(current_presences) == 0:
-                del state[key]
-
-        return state
+                del self.state[key]
 
-    def in_pending_sync_state(self) -> bool:
-        return self.join_ref is None or self.join_ref != self.channel.join_ref
+        return self.state
 
     @staticmethod
     def _transform_state(
-        state: Union[RawPresenceState, RealtimePresenceState],
+        state: RawPresenceState,
     ) -> RealtimePresenceState:
         """
         Transform the raw presence state into a standardized RealtimePresenceState format.
@@ -216,16 +169,12 @@ def _transform_state(
         """
         new_state: RealtimePresenceState = {}
         for key, presences in state.items():
-            if isinstance(presences, dict) and "metas" in presences:
-                new_state[key] = []
-
-                for presence in presences["metas"]:
-                    if "phx_ref_prev" in presence:
-                        del presence["phx_ref_prev"]
-                    new_presence: Presence = {"presence_ref": presence.pop("phx_ref")}
-                    new_presence.update(presence)
-                    new_state[key].append(new_presence)
-
-            else:
-                new_state[key] = presences
+            new_state[key] = []
+
+            for presence in presences["metas"]:
+                if "phx_ref_prev" in presence:
+                    del presence["phx_ref_prev"]
+                new_presence: Presence = {"presence_ref": presence.pop("phx_ref")}
+                new_presence.update(presence)
+                new_state[key].append(new_presence)
         return new_state