diff --git a/src/socket_client.rs b/src/socket_client.rs
index 926e70e..60695af 100644
--- a/src/socket_client.rs
+++ b/src/socket_client.rs
@@ -116,11 +116,7 @@ pub async fn connect_tcp_local<
 ) -> ResultType<Stream> {
     let target_str = target.to_string();
 
-    // if target_str.starts_with("ws://") || target_str.starts_with("wss://") {
-    //     Ok(Stream::WebSocket(
-    //         websocket::WsFramedStream::new(target_str, local, None, ms_timeout).await?,
-    //     ))
-    if true {
+    if target_str.starts_with("ws://") || target_str.starts_with("wss://") {
         Ok(Stream::WebSocket(
             websocket::WsFramedStream::new(target_str, local, None, ms_timeout).await?,
         ))
diff --git a/src/websocket.rs b/src/websocket.rs
index 094b770..2790899 100644
--- a/src/websocket.rs
+++ b/src/websocket.rs
@@ -15,6 +15,7 @@ use tokio::{net::TcpStream, time::timeout};
 use tokio_tungstenite::{
     connect_async, tungstenite::protocol::Message as WsMessage, MaybeTlsStream, WebSocketStream,
 };
+use tungstenite::client::IntoClientRequest;
 use tungstenite::protocol::Role;
 
 #[derive(Clone)]
@@ -70,15 +71,11 @@ impl WsFramedStream {
             })
         } else {
             log::info!("{:?}", url_str);
-            let ws_url = format!("ws://{}", url_str);
-            let (stream, _) = connect_async(ws_url).await?;
+
+            let (stream, _) = connect_async(url_str.into_client_request().unwrap()).await?;
 
             let addr = match stream.get_ref() {
                 MaybeTlsStream::Plain(tcp) => tcp.peer_addr()?,
-                #[cfg(feature = "native-tls")]
-                MaybeTlsStream::NativeTls(tls) => tls.get_ref().peer_addr()?,
-                #[cfg(feature = "rustls")]
-                MaybeTlsStream::Rustls(tls) => tls.get_ref().0.peer_addr()?,
                 _ => return Err(Error::new(ErrorKind::Other, "Unsupported stream type").into()),
             };