Implement Songbird driver configuration (#1074)

src/driver/config.rs

@@ -1,10 +1,81 @@
-use super::CryptoMode;
+use super::{CryptoMode, DecodeMode};
 
 /// Configuration for the inner Driver.
 ///
-/// At present, this cannot be changed.
-#[derive(Clone, Debug, Default)]
+#[derive(Clone, Debug)]
 pub struct Config {
     /// Selected tagging mode for voice packet encryption.
-    pub crypto_mode: Option<CryptoMode>,
+    ///
+    /// Defaults to [`CryptoMode::Normal`].
+    ///
+    /// Changes to this field will not immediately apply if the
+    /// driver is actively connected, but will apply to subsequent
+    /// sessions.
+    ///
+    /// [`CryptoMode::Normal`]: enum.CryptoMode.html#variant.Normal
+    pub crypto_mode: CryptoMode,
+    /// Configures whether decoding and decryption occur for all received packets.
+    ///
+    /// If voice receiving voice packets, generally you should choose [`DecodeMode::Decode`].
+    /// [`DecodeMode::Decrypt`] is intended for users running their own selective decoding,
+    /// who rely upon [user speaking events], or who need to inspect Opus packets.
+    /// If you're certain you will never need any RT(C)P events, then consider [`DecodeMode::Pass`].
+    ///
+    /// Defaults to [`DecodeMode::Decrypt`]. This is due to per-packet decoding costs,
+    /// which most users will not want to pay, but allowing speaking events which are commonly used.
+    ///
+    /// [`DecodeMode::Decode`]: enum.DecodeMode.html#variant.Decode
+    /// [`DecodeMode::Decrypt`]: enum.DecodeMode.html#variant.Decrypt
+    /// [`DecodeMode::Pass`]: enum.DecodeMode.html#variant.Pass
+    /// [user speaking events]: ../events/enum.CoreEvent.html#variant.SpeakingUpdate
+    pub decode_mode: DecodeMode,
+    /// Number of concurrently active tracks to allocate memory for.
+    ///
+    /// This should be set at, or just above, the maximum number of tracks
+    /// you expect your bot will play at the same time. Exceeding the size of
+    /// the internal queue will trigger a larger memory allocation and copy,
+    /// possibly causing the mixer thread to miss a packet deadline.
+    ///
+    /// Defaults to `1`.
+    ///
+    /// Changes to this field in a running driver will only ever increase
+    /// the capacity of the track store.
+    pub preallocated_tracks: usize,
+}
+
+impl Default for Config {
+    fn default() -> Self {
+        Self {
+            crypto_mode: CryptoMode::Normal,
+            decode_mode: DecodeMode::Decrypt,
+            preallocated_tracks: 1,
+        }
+    }
+}
+
+impl Config {
+    /// Sets this `Config`'s chosen cryptographic tagging scheme.
+    pub fn crypto_mode(mut self, crypto_mode: CryptoMode) -> Self {
+        self.crypto_mode = crypto_mode;
+        self
+    }
+
+    /// Sets this `Config`'s received packet decryption/decoding behaviour.
+    pub fn decode_mode(mut self, decode_mode: DecodeMode) -> Self {
+        self.decode_mode = decode_mode;
+        self
+    }
+
+    /// Sets this `Config`'s number of tracks to preallocate.
+    pub fn preallocated_tracks(mut self, preallocated_tracks: usize) -> Self {
+        self.preallocated_tracks = preallocated_tracks;
+        self
+    }
+
+    /// This is used to prevent changes which would invalidate the current session.
+    pub(crate) fn make_safe(&mut self, previous: &Config, connected: bool) {
+        if connected {
+            self.crypto_mode = previous.crypto_mode;
+        }
+    }
 }

src/driver/connection/mod.rs

@@ -41,8 +41,6 @@ impl Connection {
         interconnect: &Interconnect,
         config: &Config,
     ) -> Result<Connection> {
-        let crypto_mode = config.crypto_mode.unwrap_or(CryptoMode::Normal);
-
         let url = generate_url(&mut info.endpoint)?;
 
         #[cfg(all(feature = "rustls", not(feature = "native")))]
@@ -95,7 +93,7 @@ impl Connection {
         let ready =
             ready.expect("Ready packet expected in connection initialisation, but not found.");
 
-        if !has_valid_mode(&ready.modes, crypto_mode) {
+        if !has_valid_mode(&ready.modes, config.crypto_mode) {
             return Err(Error::CryptoModeUnavailable);
         }
 
@@ -147,14 +145,14 @@ impl Connection {
                     protocol: "udp".into(),
                     data: ProtocolData {
                         address,
-                        mode: crypto_mode.to_request_str().into(),
+                        mode: config.crypto_mode.to_request_str().into(),
                         port: view.get_port(),
                     },
                 }))
                 .await?;
         }
 
-        let cipher = init_cipher(&mut client, crypto_mode).await?;
+        let cipher = init_cipher(&mut client, config.crypto_mode).await?;
 
         info!("Connected to: {}", info.endpoint);
 
@@ -169,6 +167,7 @@ impl Connection {
 
         let mix_conn = MixerConnection {
             cipher: cipher.clone(),
+            crypto_state: config.crypto_mode.into(),
             udp_rx: udp_receiver_msg_tx,
             udp_tx: udp_sender_msg_tx,
         };
@@ -193,7 +192,7 @@ impl Connection {
             interconnect.clone(),
             udp_receiver_msg_rx,
             cipher,
-            crypto_mode,
+            config.clone(),
             udp_rx,
         ));
         tokio::spawn(udp_tx::runner(udp_sender_msg_rx, ssrc, udp_tx));

src/driver/crypto.rs

@@ -1,38 +1,223 @@
 //! Encryption schemes supported by Discord's secure RTP negotiation.
+use byteorder::{NetworkEndian, WriteBytesExt};
+use discortp::{rtp::RtpPacket, MutablePacket};
+use rand::Rng;
+use std::num::Wrapping;
+use xsalsa20poly1305::{
+    aead::{AeadInPlace, Error as CryptoError},
+    Nonce,
+    Tag,
+    XSalsa20Poly1305 as Cipher,
+    NONCE_SIZE,
+    TAG_SIZE,
+};
 
 /// Variants of the XSalsa20Poly1305 encryption scheme.
-///
-/// At present, only `Normal` is supported or selectable.
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 #[non_exhaustive]
-pub enum Mode {
+pub enum CryptoMode {
     /// The RTP header is used as the source of nonce bytes for the packet.
     ///
     /// Equivalent to a nonce of at most 48b (6B) at no extra packet overhead:
     /// the RTP sequence number and timestamp are the varying quantities.
     Normal,
     /// An additional random 24B suffix is used as the source of nonce bytes for the packet.
+    /// This is regenerated randomly for each packet.
     ///
     /// Full nonce width of 24B (192b), at an extra 24B per packet (~1.2 kB/s).
     Suffix,
-    /// An additional random 24B suffix is used as the source of nonce bytes for the packet.
+    /// An additional random 4B suffix is used as the source of nonce bytes for the packet.
+    /// This nonce value increments by `1` with each packet.
     ///
     /// Nonce width of 4B (32b), at an extra 4B per packet (~0.2 kB/s).
     Lite,
 }
 
-impl Mode {
+impl From<CryptoState> for CryptoMode {
+    fn from(val: CryptoState) -> Self {
+        use CryptoState::*;
+        match val {
+            Normal => CryptoMode::Normal,
+            Suffix => CryptoMode::Suffix,
+            Lite(_) => CryptoMode::Lite,
+        }
+    }
+}
+
+impl CryptoMode {
     /// Returns the name of a mode as it will appear during negotiation.
     pub fn to_request_str(self) -> &'static str {
-        use Mode::*;
+        use CryptoMode::*;
         match self {
             Normal => "xsalsa20_poly1305",
             Suffix => "xsalsa20_poly1305_suffix",
             Lite => "xsalsa20_poly1305_lite",
         }
     }
+
+    /// Returns the number of bytes each nonce is stored as within
+    /// a packet.
+    pub fn nonce_size(self) -> usize {
+        use CryptoMode::*;
+        match self {
+            Normal => RtpPacket::minimum_packet_size(),
+            Suffix => NONCE_SIZE,
+            Lite => 4,
+        }
+    }
+
+    /// Returns the number of bytes occupied by the encryption scheme
+    /// which fall before the payload.
+    pub fn payload_prefix_len(self) -> usize {
+        TAG_SIZE
+    }
+
+    /// Returns the number of bytes occupied by the encryption scheme
+    /// which fall after the payload.
+    pub fn payload_suffix_len(self) -> usize {
+        use CryptoMode::*;
+        match self {
+            Normal => 0,
+            Suffix | Lite => self.nonce_size(),
+        }
+    }
+
+    /// Calculates the number of additional bytes required compared
+    /// to an unencrypted payload.
+    pub fn payload_overhead(self) -> usize {
+        self.payload_prefix_len() + self.payload_suffix_len()
+    }
+
+    /// Extracts the byte slice in a packet used as the nonce, and the remaining mutable
+    /// portion of the packet.
+    fn nonce_slice<'a>(self, header: &'a [u8], body: &'a mut [u8]) -> (&'a [u8], &'a mut [u8]) {
+        use CryptoMode::*;
+        match self {
+            Normal => (header, body),
+            Suffix | Lite => {
+                let len = body.len();
+                let (body_left, nonce_loc) = body.split_at_mut(len - self.payload_suffix_len());
+                (&nonce_loc[..self.nonce_size()], body_left)
+            },
+        }
+    }
+
+    /// Decrypts a Discord RT(C)P packet using the given key.
+    ///
+    /// If successful, this returns the number of bytes to be ignored from the
+    /// start and end of the packet payload.
+    #[inline]
+    pub(crate) fn decrypt_in_place(
+        self,
+        packet: &mut impl MutablePacket,
+        cipher: &Cipher,
+    ) -> Result<(usize, usize), CryptoError> {
+        let header_len = packet.packet().len() - packet.payload().len();
+        let (header, body) = packet.packet_mut().split_at_mut(header_len);
+        let (slice_to_use, body_remaining) = self.nonce_slice(header, body);
+
+        let mut nonce = Nonce::default();
+        let nonce_slice = if slice_to_use.len() == NONCE_SIZE {
+            Nonce::from_slice(&slice_to_use[..NONCE_SIZE])
+        } else {
+            let max_bytes_avail = slice_to_use.len();
+            nonce[..self.nonce_size().min(max_bytes_avail)].copy_from_slice(slice_to_use);
+            &nonce
+        };
+
+        let body_start = self.payload_prefix_len();
+        let body_tail = self.payload_suffix_len();
+
+        let (tag_bytes, data_bytes) = body_remaining.split_at_mut(body_start);
+        let tag = Tag::from_slice(tag_bytes);
+
+        Ok(cipher
+            .decrypt_in_place_detached(nonce_slice, b"", data_bytes, tag)
+            .map(|_| (body_start, body_tail))?)
+    }
+
+    /// Encrypts a Discord RT(C)P packet using the given key.
+    ///
+    /// Use of this requires that the input packet has had a nonce generated in the correct location,
+    /// and `payload_len` specifies the number of bytes after the header including this nonce.
+    #[inline]
+    pub fn encrypt_in_place(
+        self,
+        packet: &mut impl MutablePacket,
+        cipher: &Cipher,
+        payload_len: usize,
+    ) -> Result<(), CryptoError> {
+        let header_len = packet.packet().len() - packet.payload().len();
+        let (header, body) = packet.packet_mut().split_at_mut(header_len);
+        let (slice_to_use, body_remaining) = self.nonce_slice(header, &mut body[..payload_len]);
+
+        let mut nonce = Nonce::default();
+        let nonce_slice = if slice_to_use.len() == NONCE_SIZE {
+            Nonce::from_slice(&slice_to_use[..NONCE_SIZE])
+        } else {
+            nonce[..self.nonce_size()].copy_from_slice(slice_to_use);
+            &nonce
+        };
+
+        // body_remaining is now correctly truncated by this point.
+        // the true_payload to encrypt follows after the first TAG_LEN bytes.
+        let tag =
+            cipher.encrypt_in_place_detached(nonce_slice, b"", &mut body_remaining[TAG_SIZE..])?;
+        body_remaining[..TAG_SIZE].copy_from_slice(&tag[..]);
+
+        Ok(())
+    }
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+#[non_exhaustive]
+pub(crate) enum CryptoState {
+    Normal,
+    Suffix,
+    Lite(Wrapping<u32>),
+}
+
+impl From<CryptoMode> for CryptoState {
+    fn from(val: CryptoMode) -> Self {
+        use CryptoMode::*;
+        match val {
+            Normal => CryptoState::Normal,
+            Suffix => CryptoState::Suffix,
+            Lite => CryptoState::Lite(Wrapping(rand::random::<u32>())),
+        }
+    }
 }
 
-// TODO: implement encrypt + decrypt + nonce selection for each.
-// This will probably need some research into correct handling of
-// padding, reported length, SRTP profiles, and so on.
+impl CryptoState {
+    /// Writes packet nonce into the body, if required, returning the new length.
+    pub fn write_packet_nonce(
+        &mut self,
+        packet: &mut impl MutablePacket,
+        payload_end: usize,
+    ) -> usize {
+        let mode = self.kind();
+        let endpoint = payload_end + mode.payload_suffix_len();
+
+        use CryptoState::*;
+        match self {
+            Suffix => {
+                rand::thread_rng().fill(&mut packet.payload_mut()[payload_end..endpoint]);
+            },
+            Lite(mut i) => {
+                (&mut packet.payload_mut()[payload_end..endpoint])
+                    .write_u32::<NetworkEndian>(i.0)
+                    .expect(
+                        "Nonce size is guaranteed to be sufficient to write u32 for lite tagging.",
+                    );
+                i += Wrapping(1);
+            },
+            _ => {},
+        }
+
+        endpoint
+    }
+
+    pub fn kind(&self) -> CryptoMode {
+        CryptoMode::from(*self)
+    }
+}