Copy dave support from serenity-rs#291

Author: Kronicaler

Cargo.toml

@@ -23,6 +23,7 @@ bytes = { optional = true, version = "1" }
 chacha20poly1305 = { optional = true, version = "0.10.1" }
 crypto-common = { optional = true, features = ["std"], version = "0.1" }
 dashmap = { optional = true, version = "6.1.0" }
+davey = { optional = true, version = "0.1.2" }
 derivative = "2"
 discortp = { default-features = false, features = [
     "discord",
@@ -50,7 +51,7 @@ serenity = { default-features = false, optional = true, git = "https://github.co
     "voice",
     "gateway",
 ] }
-serenity-voice-model = { optional = true, version = "0.2" }
+serenity-voice-model = { optional = true, git = "https://github.com/serenity-rs/voice-model.git", branch = "next" }
 socket2 = { optional = true, version = "0.5" }
 streamcatcher = { optional = true, version = "1" }
 stream_lib = { default-features = false, optional = true, version = "0.5.2" }
@@ -104,6 +105,7 @@ driver = [
     "dep:bytes",
     "dep:chacha20poly1305",
     "dep:crypto-common",
+    "dep:davey",
     "dep:discortp",
     "dep:flume",
     "dep:nohash-hasher",

benches/mixing-task.rs

@@ -1,13 +1,7 @@
 use std::error::Error;
 
 use criterion::{
-    black_box,
-    criterion_group,
-    criterion_main,
-    BatchSize,
-    Bencher,
-    BenchmarkId,
-    Criterion,
+    black_box, criterion_group, criterion_main, BatchSize, Bencher, BenchmarkId, Criterion,
 };
 use flume::{Receiver, Sender, TryRecvError};
 use songbird::{
@@ -20,14 +14,10 @@ use songbird::{
             task_message::*,
             CryptoState,
         },
-        Bitrate,
-        DummyMixer,
-        Listeners,
-        MockScheduler,
+        Bitrate, DummyMixer, Listeners, MockScheduler,
     },
     input::{cached::Compressed, codecs::*, Input, RawAdapter},
-    tracks,
-    Config,
+    tracks, Config,
 };
 use std::{io::Cursor, net::UdpSocket, sync::Arc};
 use tokio::runtime::{Handle, Runtime};

src/config.rs

@@ -3,11 +3,7 @@ use crate::driver::{Channels, DecodeMode, SampleRate};
 #[cfg(feature = "driver")]
 use crate::{
     driver::{
-        get_default_scheduler,
-        retry::Retry,
-        tasks::disposal::DisposalThread,
-        CryptoMode,
-        MixMode,
+        get_default_scheduler, retry::Retry, tasks::disposal::DisposalThread, CryptoMode, MixMode,
         Scheduler,
     },
     input::codecs::*,

src/driver/connection/error.rs

@@ -25,6 +25,12 @@ pub enum Error {
     CryptoModeInvalid,
     /// Selected crypto mode was not offered by server.
     CryptoModeUnavailable,
+    /// Failed to create a DAVE key package.
+    #[cfg(feature = "driver")]
+    DaveCreateKeyPackageError(davey::errors::CreateKeyPackageError),
+    /// An error occurred during initialization of the DAVE session.
+    #[cfg(feature = "driver")]
+    DaveInitializationError(davey::errors::InitError),
     /// An indicator that an endpoint URL was invalid.
     EndpointUrl,
     /// Discord failed to correctly respond to IP discovery.
@@ -100,6 +106,8 @@ impl fmt::Display for Error {
             Self::CryptoInvalidLength => write!(f, "server supplied key of wrong length"),
             Self::CryptoModeInvalid => write!(f, "server changed negotiated encryption mode"),
             Self::CryptoModeUnavailable => write!(f, "server did not offer chosen encryption mode"),
+            Self::DaveCreateKeyPackageError(e) => e.fmt(f),
+            Self::DaveInitializationError(e) => e.fmt(f),
             Self::EndpointUrl => write!(f, "endpoint URL received from gateway was invalid"),
             Self::IllegalDiscoveryResponse => {
                 write!(f, "IP discovery/NAT punching response was invalid")
@@ -122,6 +130,8 @@ impl StdError for Error {
             | Error::CryptoInvalidLength
             | Error::CryptoModeInvalid
             | Error::CryptoModeUnavailable
+            | Error::DaveCreateKeyPackageError(_)
+            | Error::DaveInitializationError(_)
             | Error::EndpointUrl
             | Error::IllegalDiscoveryResponse
             | Error::IllegalIp

src/driver/connection/mod.rs

@@ -8,27 +8,31 @@ use super::{
         message::*,
         ws::{self as ws_task, AuxNetwork},
     },
-    Config,
-    CryptoMode,
+    Config, CryptoMode,
 };
 use crate::{
     constants::*,
     model::{
         payload::{Identify, Resume, SelectProtocol},
-        Event as GatewayEvent,
-        ProtocolData,
+        Event as GatewayEvent, ProtocolData,
     },
     ws::WsStream,
     ConnectionInfo,
 };
 use discortp::discord::{IpDiscoveryPacket, IpDiscoveryType, MutableIpDiscoveryPacket};
 use error::{Error, Result};
 use flume::Sender;
+use serenity_voice_model::payload::DaveMlsKeyPackage;
 use socket2::Socket;
 #[cfg(feature = "receive")]
 use std::sync::Arc;
-use std::{net::IpAddr, str::FromStr};
-use tokio::{net::UdpSocket, spawn, time::timeout};
+use std::{
+    net::IpAddr,
+    num::NonZeroU16,
+    str::FromStr,
+    sync::{atomic::AtomicU16, Arc},
+};
+use tokio::{net::UdpSocket, spawn, sync::RwLock, time::timeout};
 use tracing::{debug, info, instrument};
 use url::Url;
 
@@ -72,11 +76,12 @@ impl Connection {
                 session_id: info.session_id.clone(),
                 token: info.token.clone(),
                 user_id: info.user_id.into(),
+                max_dave_protocol_version: Some(davey::DAVE_PROTOCOL_VERSION),
             }))
             .await?;
 
         loop {
-            let Some(value) = client.recv_json().await? else {
+            let Some(value) = client.recv_event().await? else {
                 continue;
             };
 
@@ -181,7 +186,10 @@ impl Connection {
                 .await?;
         }
 
-        let cipher = init_cipher(&mut client, chosen_crypto, &ws_msg_tx).await?;
+        let (cipher, dave_session, dave_protocol_version) =
+            init_cipher(&mut client, &info, chosen_crypto, &ws_msg_tx).await?;
+        let dave_session = Arc::new(RwLock::new(dave_session));
+        let dave_protocol_version = Arc::new(dave_protocol_version);
 
         info!("Connected to: {}", info.endpoint);
 
@@ -213,6 +221,8 @@ impl Connection {
             cipher: cipher.clone(),
             #[cfg(not(feature = "receive"))]
             cipher,
+            dave_session: dave_session.clone(),
+            dave_protocol_version: dave_protocol_version.clone(),
             crypto_state: chosen_crypto.into(),
             #[cfg(feature = "receive")]
             udp_rx: udp_receiver_msg_tx,
@@ -237,6 +247,14 @@ impl Connection {
             hello.heartbeat_interval,
             idx,
             info.clone(),
+            #[cfg(not(feature = "receive"))]
+            dave_session,
+            #[cfg(not(feature = "receive"))]
+            dave_protocol_version,
+            #[cfg(feature = "receive")]
+            dave_session.clone(),
+            #[cfg(feature = "receive")]
+            dave_protocol_version.clone(),
             #[cfg(feature = "receive")]
             ssrc_tracker.clone(),
         );
@@ -252,6 +270,8 @@ impl Connection {
             config.clone(),
             udp_rx,
             ssrc_tracker,
+            dave_session,
+            dave_protocol_version,
         ));
 
         Ok(Connection {
@@ -290,7 +310,7 @@ impl Connection {
         let mut resumed = None;
 
         loop {
-            let Some(value) = client.recv_json().await? else {
+            let Some(value) = client.recv_event().await? else {
                 continue;
             };
 
@@ -344,11 +364,12 @@ fn generate_url(endpoint: &mut String) -> Result<Url> {
 #[inline]
 async fn init_cipher(
     client: &mut WsStream,
+    info: &ConnectionInfo,
     mode: CryptoMode,
     tx: &Sender<WsMessage>,
-) -> Result<Cipher> {
+) -> Result<(Cipher, Option<davey::DaveSession>, AtomicU16)> {
     loop {
-        let Some(value) = client.recv_json().await? else {
+        let Some(value) = client.recv_event().await? else {
             continue;
         };
 
@@ -358,9 +379,37 @@ async fn init_cipher(
                     return Err(Error::CryptoModeInvalid);
                 }
 
-                return mode
-                    .cipher_from_key(&desc.secret_key)
-                    .map_err(|_| Error::CryptoInvalidLength);
+                let dave_session =
+                    if let Some(version) = NonZeroU16::new(desc.dave_protocol_version) {
+                        let mut session = davey::DaveSession::new(
+                            version,
+                            info.user_id.get(),
+                            info.channel_id
+                                .expect("channel ID must be set in connection info")
+                                .get(),
+                            None,
+                        )
+                        .map_err(Error::DaveInitializationError)?;
+
+                        client
+                            .send_binary(&GatewayEvent::DaveMlsKeyPackage(DaveMlsKeyPackage {
+                                key_package: session
+                                    .create_key_package()
+                                    .map_err(Error::DaveCreateKeyPackageError)?,
+                            }))
+                            .await?;
+
+                        Some(session)
+                    } else {
+                        None
+                    };
+
+                return Ok((
+                    mode.cipher_from_key(&desc.secret_key)
+                        .map_err(|_| Error::CryptoInvalidLength)?,
+                    dave_session,
+                    AtomicU16::new(desc.dave_protocol_version),
+                ));
             },
             other => {
                 // Discord can and will send user-specific payload packets during this time

src/driver/crypto.rs

@@ -91,8 +91,9 @@ impl CryptoMode {
             EncryptionAlgorithm::Aes256Gcm => Aes256Gcm::new_from_slice(key)
                 .map(Box::new)
                 .map(Cipher::Aes256Gcm),
-            EncryptionAlgorithm::XChaCha20Poly1305 =>
-                XChaCha20Poly1305::new_from_slice(key).map(Cipher::XChaCha20Poly1305),
+            EncryptionAlgorithm::XChaCha20Poly1305 => {
+                XChaCha20Poly1305::new_from_slice(key).map(Cipher::XChaCha20Poly1305)
+            },
         }
     }
 
@@ -255,8 +256,9 @@ impl From<CryptoMode> for CryptoState {
     fn from(val: CryptoMode) -> Self {
         match val {
             CryptoMode::Aes256Gcm => CryptoState::Aes256Gcm(Wrapping(rand::random::<u32>())),
-            CryptoMode::XChaCha20Poly1305 =>
-                CryptoState::XChaCha20Poly1305(Wrapping(rand::random::<u32>())),
+            CryptoMode::XChaCha20Poly1305 => {
+                CryptoState::XChaCha20Poly1305(Wrapping(rand::random::<u32>()))
+            },
         }
     }
 }

src/driver/mod.rs

@@ -31,12 +31,8 @@ pub(crate) use crypto::CryptoState;
 pub use decode_mode::*;
 pub use mix_mode::MixMode;
 pub use scheduler::{
-    get_default_scheduler,
-    Config as SchedulerConfig,
-    Error as SchedulerError,
-    LiveStatBlock,
-    Mode as SchedulerMode,
-    Scheduler,
+    get_default_scheduler, Config as SchedulerConfig, Error as SchedulerError, LiveStatBlock,
+    Mode as SchedulerMode, Scheduler,
 };
 #[cfg(test)]
 pub use test_config::*;
@@ -49,10 +45,7 @@ use crate::{
     events::EventData,
     input::Input,
     tracks::{Track, TrackHandle},
-    Config,
-    ConnectionInfo,
-    Event,
-    EventHandler,
+    Config, ConnectionInfo, Event, EventHandler,
 };
 /// Opus encoder bitrate settings.
 pub use audiopus::{self as opus, Bitrate};

src/driver/tasks/error.rs

@@ -20,6 +20,7 @@ pub type Result<T> = std::result::Result<T, Error>;
 #[non_exhaustive]
 pub enum Error {
     Crypto(CryptoError),
+    DaveEncrypt(davey::errors::EncryptError),
     #[cfg(any(feature = "receive", test))]
     /// Received an illegal voice packet on the voice UDP socket.
     IllegalVoicePacket,
@@ -58,6 +59,12 @@ impl From<CryptoError> for Error {
     }
 }
 
+impl From<davey::errors::EncryptError> for Error {
+    fn from(value: davey::errors::EncryptError) -> Self {
+        Self::DaveEncrypt(value)
+    }
+}
+
 impl From<IoError> for Error {
     fn from(e: IoError) -> Error {
         Error::Io(e)
@@ -100,3 +107,43 @@ impl From<WsError> for Error {
         Error::Other
     }
 }
+
+#[derive(Debug)]
+#[non_exhaustive]
+pub enum DaveReinitError {
+    Init(davey::errors::InitError),
+    Reinit(davey::errors::ReinitError),
+    Reset(davey::errors::ResetError),
+    CreateKeyPackage(davey::errors::CreateKeyPackageError),
+    Ws(WsError),
+}
+
+impl From<davey::errors::InitError> for DaveReinitError {
+    fn from(value: davey::errors::InitError) -> Self {
+        Self::Init(value)
+    }
+}
+
+impl From<davey::errors::ReinitError> for DaveReinitError {
+    fn from(value: davey::errors::ReinitError) -> Self {
+        Self::Reinit(value)
+    }
+}
+
+impl From<davey::errors::ResetError> for DaveReinitError {
+    fn from(value: davey::errors::ResetError) -> Self {
+        Self::Reset(value)
+    }
+}
+
+impl From<davey::errors::CreateKeyPackageError> for DaveReinitError {
+    fn from(value: davey::errors::CreateKeyPackageError) -> Self {
+        Self::CreateKeyPackage(value)
+    }
+}
+
+impl From<WsError> for DaveReinitError {
+    fn from(value: WsError) -> Self {
+        Self::Ws(value)
+    }
+}

src/driver/tasks/message/mixer.rs

@@ -9,12 +9,18 @@ use crate::{
     input::{AudioStreamError, Compose, Parsed},
 };
 use flume::Sender;
-use std::{net::UdpSocket, sync::Arc};
+use std::{
+    net::UdpSocket,
+    sync::{atomic::AtomicU16, Arc},
+};
 use symphonia_core::{errors::Error as SymphoniaError, formats::SeekedTo};
+use tokio::sync::RwLock;
 
 pub struct MixerConnection {
     pub cipher: Cipher,
     pub crypto_state: CryptoState,
+    pub dave_session: Arc<RwLock<Option<davey::DaveSession>>>,
+    pub dave_protocol_version: Arc<AtomicU16>,
     #[cfg(feature = "receive")]
     pub udp_rx: Sender<UdpRxMessage>,
     pub udp_tx: UdpSocket,