Driver: Support tokio-websockets (#226)

* Driver: Support `tokio-websockets`

* Fix bad feature flag

* Fix CI & examples features

* Use tungstenite in twilight example

* Error if none or both ws features are enabled

* Match `twilight-gateway` features

src/driver/tasks/ws.rs

@@ -20,6 +20,7 @@ use tokio::{
     select,
     time::{sleep_until, Instant},
 };
+#[cfg(feature = "tungstenite")]
 use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
 use tracing::{debug, info, instrument, trace, warn};
 
@@ -248,6 +249,7 @@ pub(crate) async fn runner(mut interconnect: Interconnect, mut aux: AuxNetwork)
 
 fn ws_error_is_not_final(err: &WsError) -> bool {
     match err {
+        #[cfg(feature = "tungstenite")]
         WsError::WsClosed(Some(frame)) => match frame.code {
             CloseCode::Library(l) =>
                 if let Some(code) = VoiceCloseCode::from_u16(l) {
@@ -257,6 +259,16 @@ fn ws_error_is_not_final(err: &WsError) -> bool {
                 },
             _ => true,
         },
+        #[cfg(feature = "tws")]
+        WsError::WsClosed(Some(code)) => match (*code).into() {
+            code @ 4000..=4999_u16 =>
+                if let Some(code) = VoiceCloseCode::from_u16(code) {
+                    code.should_resume()
+                } else {
+                    true
+                },
+            _ => true,
+        },
         e => {
             debug!("Error sending/receiving ws {:?}.", e);
             true

src/events/context/data/disconnect.rs

@@ -4,6 +4,7 @@ use crate::{
     model::{CloseCode as VoiceCloseCode, FromPrimitive},
     ws::Error as WsError,
 };
+#[cfg(feature = "tungstenite")]
 use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
 
 /// Voice connection details gathered at termination or failure.
@@ -108,10 +109,16 @@ impl From<&ConnectionError> for DisconnectReason {
 impl From<&WsError> for DisconnectReason {
     fn from(e: &WsError) -> Self {
         Self::WsClosed(match e {
+            #[cfg(feature = "tungstenite")]
             WsError::WsClosed(Some(frame)) => match frame.code {
                 CloseCode::Library(l) => VoiceCloseCode::from_u16(l),
                 _ => None,
             },
+            #[cfg(feature = "tws")]
+            WsError::WsClosed(Some(code)) => match (*code).into() {
+                code @ 4000..=4999_u16 => VoiceCloseCode::from_u16(code),
+                _ => None,
+            },
             _ => None,
         })
     }

src/ws.rs

@@ -5,6 +5,7 @@ use tokio::{
     net::TcpStream,
     time::{timeout, Duration},
 };
+#[cfg(feature = "tungstenite")]
 use tokio_tungstenite::{
     tungstenite::{
         error::Error as TungsteniteError,
@@ -14,16 +15,32 @@ use tokio_tungstenite::{
     MaybeTlsStream,
     WebSocketStream,
 };
+#[cfg(feature = "tws")]
+use tokio_websockets::{
+    CloseCode,
+    Error as TwsError,
+    Limits,
+    MaybeTlsStream,
+    Message,
+    WebSocketStream,
+};
 use tracing::{debug, instrument};
 use url::Url;
 
+#[cfg(any(
+    all(feature = "tws", feature = "tungstenite"),
+    all(not(feature = "tws"), not(feature = "tungstenite"))
+))]
+compile_error!("specify one of `features = [\"tungstenite\"]` (recommended w/ serenity) or `features = [\"tws\"]` (recommended w/ twilight)");
+
 pub struct WsStream(WebSocketStream<MaybeTlsStream<TcpStream>>);
 
 impl WsStream {
     #[instrument]
     pub(crate) async fn connect(url: Url) -> Result<Self> {
-        let (stream, _) = tokio_tungstenite::connect_async_with_config(
-            url.to_string(),
+        #[cfg(feature = "tungstenite")]
+        let (stream, _) = tokio_tungstenite::connect_async_with_config::<Url>(
+            url,
             Some(Config {
                 max_message_size: None,
                 max_frame_size: None,
@@ -32,6 +49,13 @@ impl WsStream {
             true,
         )
         .await?;
+        #[cfg(feature = "tws")]
+        let (stream, _) = tokio_websockets::ClientBuilder::new()
+            .limits(Limits::unlimited())
+            .uri(url.as_str())
+            .unwrap() // Any valid URL is a valid URI.
+            .connect()
+            .await?;
 
         Ok(Self(stream))
     }
@@ -53,11 +77,12 @@ impl WsStream {
     }
 
     pub(crate) async fn send_json(&mut self, value: &Event) -> Result<()> {
-        Ok(crate::json::to_string(value)
-            .map(Message::Text)
-            .map_err(Error::from)
-            .map(|m| self.0.send(m))?
-            .await?)
+        let res = crate::json::to_string(value);
+        #[cfg(feature = "tungstenite")]
+        let res = res.map(Message::Text);
+        #[cfg(feature = "tws")]
+        let res = res.map(Message::text);
+        Ok(res.map_err(Error::from).map(|m| self.0.send(m))?.await?)
     }
 }
 
@@ -71,9 +96,15 @@ pub enum Error {
     /// As a result, only text messages are expected.
     UnexpectedBinaryMessage(Vec<u8>),
 
+    #[cfg(feature = "tungstenite")]
     Ws(TungsteniteError),
+    #[cfg(feature = "tws")]
+    Ws(TwsError),
 
+    #[cfg(feature = "tungstenite")]
     WsClosed(Option<CloseFrame<'static>>),
+    #[cfg(feature = "tws")]
+    WsClosed(Option<CloseCode>),
 }
 
 impl From<JsonError> for Error {
@@ -82,16 +113,25 @@ impl From<JsonError> for Error {
     }
 }
 
+#[cfg(feature = "tungstenite")]
 impl From<TungsteniteError> for Error {
     fn from(e: TungsteniteError) -> Error {
         Error::Ws(e)
     }
 }
 
+#[cfg(feature = "tws")]
+impl From<TwsError> for Error {
+    fn from(e: TwsError) -> Self {
+        Error::Ws(e)
+    }
+}
+
 #[inline]
 #[allow(unused_unsafe)]
 pub(crate) fn convert_ws_message(message: Option<Message>) -> Result<Option<Event>> {
-    Ok(match message {
+    #[cfg(feature = "tungstenite")]
+    return Ok(match message {
         // SAFETY:
         // simd-json::serde::from_str may leave an &mut str in a non-UTF state on failure.
         // The below is safe as we have taken ownership of the inner `String`, and if
@@ -112,5 +152,33 @@ pub(crate) fn convert_ws_message(message: Option<Message>) -> Result<Option<Even
         },
         // Ping/Pong message behaviour is internally handled by tungstenite.
         _ => None,
-    })
+    });
+    #[cfg(feature = "tws")]
+    return Ok(if let Some(message) = message {
+        if message.is_text() {
+            let mut payload = message.as_text().unwrap().to_owned();
+            // SAFETY:
+            // simd-json::serde::from_str may leave an &mut str in a non-UTF state on failure.
+            // The below is safe as we have created an owned copy of the payload `&str`, and if
+            // failure occurs we forcibly re-validate its contents before logging.
+            (unsafe { crate::json::from_str(payload.as_mut_str()) })
+                .map_err(|e| {
+                    let safe_payload = String::from_utf8_lossy(payload.as_bytes());
+                    debug!("Unexpected JSON: {e}. Payload: {safe_payload}");
+                    e
+                })
+                .ok()
+        } else if message.is_binary() {
+            return Err(Error::UnexpectedBinaryMessage(
+                message.into_payload().to_vec(),
+            ));
+        } else if message.is_close() {
+            return Err(Error::WsClosed(message.as_close().map(|(c, _)| c)));
+        } else {
+            // ping/pong; will also be internally handled by tokio-websockets
+            None
+        }
+    } else {
+        None
+    });
 }