1 files changed, 41 insertions, 7 deletions

diff --git a/light-client/src/network.rs b/light-client/src/network.rs
index dc6e894f..e3cd1eb5 100644
--- a/light-client/src/network.rs
+++ b/light-client/src/network.rs
@@ -1,22 +1,56 @@
+use anyhow::Result;
 use hurrycurry_protocol::{PacketC, PacketS, BINCODE_CONFIG};
 use log::warn;
 use std::{collections::VecDeque, net::TcpStream};
-use tungstenite::{stream::MaybeTlsStream, Message, WebSocket};
+use tungstenite::{
+    client::{uri_mode, IntoClientRequest},
+    client_tls_with_config,
+    handshake::client::Request,
+    stream::{MaybeTlsStream, Mode},
+    Message, WebSocket,
+};
 
 pub struct Network {
     sock: WebSocket<MaybeTlsStream<TcpStream>>,
-    queue_in: VecDeque<PacketC>,
-    queue_out: VecDeque<PacketS>,
+    pub queue_in: VecDeque<PacketC>,
+    pub queue_out: VecDeque<PacketS>,
 }
 
 impl Network {
-    pub fn connect(addr: &str) -> Self {
-        let (sock, _resp) = tungstenite::connect(addr).unwrap();
-        Self {
+    pub fn connect(addr: &str) -> Result<Self> {
+        let (parts, _) = addr.into_client_request().unwrap().into_parts();
+        let uri = parts.uri.clone();
+
+        let mut builder = Request::builder()
+            .uri(uri.clone())
+            .method(parts.method.clone())
+            .version(parts.version);
+        *builder.headers_mut().unwrap() = parts.headers.clone();
+        let request = builder.body(()).unwrap();
+
+        let uri = request.uri();
+        let mode = uri_mode(uri)?;
+        let host = request.uri().host().unwrap();
+        let host = if host.starts_with('[') {
+            &host[1..host.len() - 1]
+        } else {
+            host
+        };
+        let port = uri.port_u16().unwrap_or(match mode {
+            Mode::Plain => 27032,
+            Mode::Tls => 443,
+        });
+        let stream = TcpStream::connect((host, port))?;
+        stream.set_nodelay(true).unwrap();
+        // stream.set_nonblocking(true).unwrap();
+
+        let (sock, _) = client_tls_with_config(request, stream, None, None).unwrap();
+
+        Ok(Self {
             sock,
             queue_in: VecDeque::new(),
             queue_out: VecDeque::new(),
-        }
+        })
     }
     pub fn poll(&mut self) {
         self.queue_in.extend(match self.sock.read() {