add first pass to implement dave

Author: DA-344

discord/client.py

@@ -50,7 +50,7 @@
 from .invite import Invite
 from .iterators import EntitlementIterator, GuildIterator
 from .mentions import AllowedMentions
-from .monetization import SKU, Entitlement
+from .monetization import SKU
 from .object import Object
 from .soundboard import SoundboardSound
 from .stage_instance import StageInstance

discord/gateway.py

@@ -41,7 +41,6 @@
 
 from . import utils
 from .activity import BaseActivity
-from .enums import SpeakingState
 from .errors import ConnectionClosed, InvalidArgument
 
 if TYPE_CHECKING:
@@ -55,8 +54,6 @@
 __all__ = (
     "DiscordWebSocket",
     "KeepAliveHandler",
-    "VoiceKeepAliveHandler",
-    "DiscordVoiceWebSocket",
     "ReconnectWebSocket",
 )
 
@@ -228,31 +225,6 @@ def ack(self) -> None:
             _log.warning(self.behind_msg, self.shard_id, self.latency)
 
 
-class VoiceKeepAliveHandler(KeepAliveHandler):
-    if TYPE_CHECKING:
-        ws: DiscordVoiceWebSocket
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.recent_ack_latencies = deque(maxlen=20)
-        self.msg = "Keeping shard ID %s voice websocket alive with timestamp %s."
-        self.block_msg = "Shard ID %s voice heartbeat blocked for more than %s seconds"
-        self.behind_msg = "High socket latency, shard ID %s heartbeat is %.1fs behind"
-
-    def get_payload(self):
-        return {
-            "op": self.ws.HEARTBEAT,
-            "d": {"t": int(time.time() * 1000), "seq_ack": self.ws.seq_ack},
-        }
-
-    def ack(self):
-        ack_time = time.perf_counter()
-        self._last_ack = ack_time
-        self._last_recv = ack_time
-        self.latency = ack_time - self._last_send
-        self.recent_ack_latencies.append(self.latency)
-
-
 class DiscordClientWebSocketResponse(aiohttp.ClientWebSocketResponse):
     async def close(self, *, code: int = 4000, message: bytes = b"") -> bool:
         return await super().close(code=code, message=message)

discord/voice/client.py

@@ -241,6 +241,14 @@ def average_latency(self) -> float:
         ws = self.ws
         return float("inf") if not ws else ws.average_latency
 
+    @property
+    def privacy_code(self) -> str | None:
+        """Returns the current voice session's privacy code, only available if the call has upgraded to use the
+        DAVE protocol
+        """
+        session = self._connection.dave_session
+        return session and session.voice_privacy_code
+
     async def disconnect(self, *, force: bool = False) -> None:
         """|coro|
 
@@ -288,6 +296,10 @@ def is_paused(self) -> bool:
     # audio related
 
     def _get_voice_packet(self, data: Any) -> bytes:
+
+        session = self._connection.dave_session
+        packet = session.encrypt_opus(data) if session and session.ready else data
+
         header = bytearray(12)
 
         # formulate rtp header
@@ -298,7 +310,7 @@ def _get_voice_packet(self, data: Any) -> bytes:
         struct.pack_into(">I", header, 8, self.ssrc)
 
         encrypt_packet = getattr(self, f"_encrypt_{self.mode}")
-        return encrypt_packet(header, data)
+        return encrypt_packet(header, packet)
 
     # encryption methods
 

discord/voice/enums.py

@@ -42,6 +42,19 @@ class OpCodes(Enum):
     client_connect = 10
     client_disconnect = 11
 
+    # dave protocol stuff
+    dave_prepare_transition = 21
+    dave_execute_transition = 22
+    dave_transition_ready = 23
+    dave_prepare_epoch = 24
+    mls_external_sender_package = 25
+    mls_key_package = 26
+    mls_proposals = 27
+    mls_commit_welcome = 28
+    mls_commit_transition = 29
+    mls_welcome = 30
+    mls_invalid_commit_welcome = 31
+
     def __eq__(self, other: object) -> bool:
         if isinstance(other, int):
             return self.value == other

discord/voice/gateway.py

@@ -36,6 +36,7 @@
 
 import aiohttp
 
+import davey
 from discord import utils
 from discord.enums import SpeakingState
 from discord.errors import ConnectionClosed
@@ -117,6 +118,7 @@ def __init__(
         self.seq_ack: int = -1
         self.state: VoiceConnectionState = state
         self.ssrc_map: dict[str, dict[str, Any]] = {}
+        self.known_users: dict[int, Any] = {}
 
         if hook:
             self._hook = hook or state.ws_hook  # type: ignore
@@ -137,9 +139,22 @@ def session_id(self) -> str | None:
     def session_id(self, value: str | None) -> None:
         self.state.session_id = value
 
+    @property
+    def dave_session(self) -> davey.DaveSession | None:
+        return self.state.dave_session
+
+    @property
+    def self_id(self) -> int:
+        return self._connection.self_id
+
     async def _hook(self, *args: Any) -> Any:
         pass
 
+    async def send_as_bytes(self, op: int, data: bytes) -> None:
+        packet = bytes(op) + data
+        _log.debug("Sending voice websocket binary frame: op: %s data: %s", op, str(data))
+        await self.ws.send_bytes(packet)
+
     async def send_as_json(self, data: Any) -> None:
         _log.debug("Sending voice websocket frame: %s.", data)
         await self.ws.send_str(utils._to_json(data))
@@ -163,6 +178,7 @@ async def received_message(self, msg: Any, /):
         op = msg["op"]
         data = msg.get("d", {})  # this key should ALWAYS be given, but guard anyways
         self.seq_ack = msg.get("seq", self.seq_ack)  # keep the seq_ack updated
+        state = self.state
 
         if op == OpCodes.ready:
             await self.ready(data)
@@ -179,18 +195,99 @@ async def received_message(self, msg: Any, /):
                 "successfully RESUMED.",
             )
         elif op == OpCodes.session_description:
-            self.state.mode = data["mode"]
+            state.mode = data["mode"]
+            state.dave_protocol_version = data["dave_protocol_version"]
             await self.load_secret_key(data)
+            await state.reinit_dave_session()
         elif op == OpCodes.hello:
             interval = data["heartbeat_interval"] / 1000.0
             self._keep_alive = KeepAliveHandler(
                 ws=self,
                 interval=min(interval, 5),
             )
             self._keep_alive.start()
+        elif self.dave_session:
+            if op == OpCodes.dave_prepare_transition:
+                _log.info("Preparing to upgrade to a DAVE connection for channel %s", state.channel_id)
+                state.dave_pending_transition = data
+
+                transition_id = data["transition_id"]
+
+                if transition_id == 0:
+                    await state.execute_dave_transition(data["transition_id"])
+                else:
+                    if data["protocol_version"] == 0:
+                        self.dave_session.set_passthrough_mode(True, 120)
+                    await self.send_dave_transition_ready(transition_id)
+            elif op == OpCodes.dave_execute_transition:
+                _log.info("Upgrading to DAVE connection for channel %s", state.channel_id)
+                await state.execute_dave_transition(data["transition_id"])
+            elif op == OpCodes.dave_prepare_epoch:
+                epoch = data["epoch"]
+                _log.debug("Preparing for DAVE epoch in channel %s: %s", state.channel_id, epoch)
+                # if epoch is 1 then a new MLS group is to be created for the proto version
+                if epoch == 1:
+                    state.dave_protocol_version = data["protocol_version"]
+                    await state.reinit_dave_session()
+        else:
+            _log.debug("Unhandled op code: %s with data %s", op, data)
 
         await utils.maybe_coroutine(self._hook, self, msg)
 
+    async def received_binary_message(self, msg: bytes) -> None:
+        self.seq_ack = struct.unpack_from(">H", msg, 0)[0]
+        op = msg[2]
+        _log.debug("Voice websocket binary frame received: %d bytes, seq: %s, op: %s", len(msg), self.seq_ack, op)
+
+        state = self.state
+
+        if not self.dave_session:
+            return
+
+        if op == OpCodes.mls_external_sender_package:
+            self.dave_session.set_external_sender(msg[3:])
+        elif op == OpCodes.mls_proposals:
+            op_type = msg[3]
+            result = self.dave_session.process_proposals(
+                davey.ProposalsOperationType.append if op_type == 0 else davey.ProposalsOperationType.revoke,
+                msg[4:],
+            )
+
+            if isinstance(result, davey.CommitWelcome):
+                await self.send_as_bytes(
+                    OpCodes.mls_key_package.value,
+                    (result.commit + result.welcome) if result.welcome else result.commit,
+                )
+            _log.debug("Processed MLS proposals for current dave session")
+        elif op == OpCodes.mls_commit_transition:
+            transt_id = struct.unpack_from(">H", msg, 3)[0]
+            try:
+                self.dave_session.process_commit(msg[5:])
+                if transt_id != 0:
+                    state.dave_pending_transition = {
+                        "transition_id": transt_id,
+                        "protocol_version": state.dave_protocol_version,
+                    }
+                    await self.send_dave_transition_ready(transt_id)
+                _log.debug("Processed MLS commit for transition %s", transt_id)
+            except Exception as exc:
+                _log.debug("An exception ocurred while processing a MLS commit, this should be safe to ignore: %s", exc)
+                await state.recover_dave_from_invalid_commit(transt_id)
+        elif op == OpCodes.mls_welcome:
+            transt_id = struct.unpack_from(">H", msg, 3)[0]
+            try:
+                self.dave_session.process_welcome(msg[5:])
+                if transt_id != 0:
+                    state.dave_pending_transition = {
+                        "transition_id": transt_id,
+                        "protocol_version": state.dave_protocol_version,
+                    }
+                    await self.send_dave_transition_ready(transt_id)
+                _log.debug("Processed MLS welcome for transition %s", transt_id)
+            except Exception as exc:
+                _log.debug("An exception ocurred while processing a MLS welcome, this should be safe to ignore: %s", exc)
+                await state.recover_dave_from_invalid_commit(transt_id)
+
     async def ready(self, data: dict[str, Any]) -> None:
         state = self.state
 
@@ -232,6 +329,7 @@ async def select_protocol(self, ip: str, port: int, mode: str) -> None:
                     "port": port,
                     "mode": mode,
                 },
+                "dave_protocol_version": self.state.dave_protocol_version,
             },
         }
         await self.send_as_json(payload)
@@ -292,6 +390,8 @@ async def poll_event(self) -> None:
 
         if msg.type is aiohttp.WSMsgType.TEXT:
             await self.received_message(utils._from_json(msg.data))
+        elif msg.type is aiohttp.WSMsgType.BINARY:
+            await self.received_binary_message(msg.data)
         elif msg.type is aiohttp.WSMsgType.ERROR:
             _log.debug("Received %s", msg)
             raise ConnectionClosed(self.ws, shard_id=None) from msg.data
@@ -355,6 +455,16 @@ async def identify(self) -> None:
                 "user_id": str(state.user.id),
                 "session_id": self.session_id,
                 "token": self.token,
+                "max_dave_protocol_version": self.state.max_dave_proto_version,
+            },
+        }
+        await self.send_as_json(payload)
+
+    async def send_dave_transition_ready(self, transition_id: int) -> None:
+        payload = {
+            "op": int(OpCodes.dave_transition_ready),
+            "d": {
+                "transition_id": transition_id,
             },
         }
         await self.send_as_json(payload)