Remove Websocket & Time crates (#172)

* Don't bring in time crate through chrono * Removed websocket crate, use blocking tungstenite instead * Removed leftover comment
wolf4ood · Feb 6, 2023 · 1f73b81 · 1f73b81
1 parent 16f469b
commit 1f73b81
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 22 deletions.
diff --git a/gremlin-client/Cargo.toml b/gremlin-client/Cargo.toml
@@ -38,10 +38,12 @@ serde = "1.0"
 serde_json = "1.0"
 serde_derive="1.0"
 r2d2 = "0.8.3"
-chrono = "0.4"
+#Avoids bringing in time crate (https://github.com/time-rs/time/issues/293)
+chrono = { version = "0.4", default-features = false}
 lazy_static = "1.3.0"
 base64 = "0.13.1"
 native-tls = "0.2.3"
+tungstenite = { version = "0.18.0", features = ["native-tls"] }
 async-tungstenite = { version = "0.18", optional = true, default-features=false}
 async-std =  { version = "1.4.0", optional = true, features = ["unstable","attributes"] }
 async-trait = { version = "0.1.10", optional = true }
@@ -62,11 +64,6 @@ pin-project-lite = { version = "0.2", optional = true}
 tokio = { version = "1", optional=true, features = ["full"] }
 
 
-[dependencies.websocket]
-version="0.26"
-default-features = false 
-features = ["sync","sync-ssl"]
-
 [dependencies.uuid]
 features = ["serde", "v4"]
 version = "1.1.2"

diff --git a/gremlin-client/src/connection.rs b/gremlin-client/src/connection.rs
@@ -1,8 +1,15 @@
+use std::net::TcpStream;
+
 use crate::{GraphSON, GremlinError, GremlinResult};
 use native_tls::TlsConnector;
-use websocket::{stream::sync::NetworkStream, sync::Client, ClientBuilder, OwnedMessage};
+use tungstenite::{
+    client::{uri_mode, IntoClientRequest},
+    client_tls_with_config,
+    stream::{MaybeTlsStream, Mode, NoDelay},
+    Connector, Message, WebSocket,
+};
 
-struct ConnectionStream(Client<Box<dyn NetworkStream + Send>>);
+struct ConnectionStream(WebSocket<MaybeTlsStream<TcpStream>>);
 
 impl std::fmt::Debug for ConnectionStream {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
@@ -13,30 +20,49 @@ impl std::fmt::Debug for ConnectionStream {
 impl ConnectionStream {
     fn connect(options: ConnectionOptions) -> GremlinResult<Self> {
         let connector = match options.tls_options.as_ref() {
-            Some(option) => Some(
+            Some(option) => Some(Connector::NativeTls(
                 option
                     .tls_connector()
                     .map_err(|e| GremlinError::Generic(e.to_string()))?,
-            ),
+            )),
             _ => None,
         };
 
-        let client = ClientBuilder::new(&options.websocket_url())
-            .map_err(|e| GremlinError::Generic(e.to_string()))?
-            .connect(connector)?;
+        let request = options
+            .websocket_url()
+            .into_client_request()
+            .map_err(|e| GremlinError::Generic(e.to_string()))?;
+        let uri = request.uri();
+        let mode = uri_mode(uri).map_err(|e| GremlinError::Generic(e.to_string()))?;
+        let host = request
+            .uri()
+            .host()
+            .ok_or_else(|| GremlinError::Generic("No Hostname".into()))?;
+        let port = uri.port_u16().unwrap_or(match mode {
+            Mode::Plain => 80,
+            Mode::Tls => 443,
+        });
+        let mut stream = TcpStream::connect((host, port))
+            .map_err(|e| GremlinError::Generic(format!("Unable to connect {e:?}")))?;
+        NoDelay::set_nodelay(&mut stream, true)
+            .map_err(|e| GremlinError::Generic(e.to_string()))?;
+
+        let (client, _response) =
+            client_tls_with_config(options.websocket_url(), stream, None, connector)
+                .map_err(|e| GremlinError::Generic(e.to_string()))?;
 
         Ok(ConnectionStream(client))
     }
 
     fn send(&mut self, payload: Vec<u8>) -> GremlinResult<()> {
         self.0
-            .send_message(&OwnedMessage::Binary(payload))
+            .write_message(Message::Binary(payload))
             .map_err(GremlinError::from)
     }
 
     fn recv(&mut self) -> GremlinResult<Vec<u8>> {
-        match self.0.recv_message()? {
-            OwnedMessage::Binary(binary) => Ok(binary),
+        match self.0.read_message()? {
+            Message::Binary(binary) => Ok(binary),
             _ => unimplemented!(),
         }
     }

diff --git a/gremlin-client/src/error.rs b/gremlin-client/src/error.rs
@@ -2,10 +2,6 @@ use crate::structure::GValue;
 
 use thiserror::Error;
 
-use websocket::WebSocketError;
-
-#[cfg(feature = "async_gremlin")]
-use async_tungstenite::tungstenite;
 #[cfg(feature = "async_gremlin")]
 use mobc;
 
@@ -16,7 +12,7 @@ pub enum GremlinError {
     Generic(String),
 
     #[error(transparent)]
-    WebSocket(#[from] WebSocketError),
+    WebSocket(#[from] tungstenite::error::Error),
 
     #[error(transparent)]
     Pool(#[from] r2d2::Error),
@@ -38,7 +34,7 @@ pub enum GremlinError {
 
     #[cfg(feature = "async_gremlin")]
     #[error(transparent)]
-    WebSocketAsync(#[from] tungstenite::error::Error),
+    WebSocketAsync(#[from] async_tungstenite::tungstenite::Error),
     #[cfg(feature = "async_gremlin")]
     #[error(transparent)]
     ChannelSend(#[from] futures::channel::mpsc::SendError),