Tungstenite: Match scheme for "ws" or "wss" (#554)

Match scheme for "ws" or "wss" and choose whether to connect with TLS connector for tungstenite

src/gateway/backends/tungstenite.rs

@@ -9,8 +9,10 @@ use futures_util::{
 };
 use tokio::net::TcpStream;
 use tokio_tungstenite::{
-    connect_async_tls_with_config, tungstenite, Connector, MaybeTlsStream, WebSocketStream,
+    connect_async_tls_with_config, connect_async_with_config, tungstenite, Connector,
+    MaybeTlsStream, WebSocketStream,
 };
+use url::Url;
 
 use crate::gateway::{GatewayMessage, RawGatewayMessage};
 
@@ -32,6 +34,21 @@ impl TungsteniteBackend {
     pub async fn connect(
         websocket_url: &str,
     ) -> Result<(TungsteniteSink, TungsteniteStream), TungsteniteBackendError> {
+        let websocket_url_parsed =
+            Url::parse(websocket_url).map_err(|_| TungsteniteBackendError::TungsteniteError {
+                error: tungstenite::error::Error::Url(
+                    tungstenite::error::UrlError::UnsupportedUrlScheme,
+                ),
+            })?;
+        if websocket_url_parsed.scheme() == "ws" {
+            let (websocket_stream, _) =
+                match connect_async_with_config(websocket_url, None, false).await {
+                    Ok(websocket_stream) => websocket_stream,
+                    Err(e) => return Err(TungsteniteBackendError::TungsteniteError { error: e }),
+                };
+
+            Ok(websocket_stream.split())
+        } else if websocket_url_parsed.scheme() == "wss" {
             let certs = webpki_roots::TLS_SERVER_ROOTS;
             let roots = rustls::RootCertStore {
                 roots: certs
@@ -64,6 +81,13 @@ impl TungsteniteBackend {
             };
 
             Ok(websocket_stream.split())
+        } else {
+            Err(TungsteniteBackendError::TungsteniteError {
+                error: tungstenite::error::Error::Url(
+                    tungstenite::error::UrlError::UnsupportedUrlScheme,
+                ),
+            })
+        }
     }
 }