From eda7b087a9770ff0f426d0349c2f9952c9fa1964 Mon Sep 17 00:00:00 2001
From: Evie Viau <evie@uwueviee.live>
Date: Sun, 20 Feb 2022 10:07:49 -0500
Subject: [PATCH] Drop connection and not crash if websocket handshake fails

---
 src/main.rs | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 215adaf..763d387 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -64,8 +64,16 @@ async fn accept_conn(peer: SocketAddr, stream: TcpStream, redis_client: Client,
 
 async fn handle_conn(peer: SocketAddr, stream: TcpStream, redis_client: Client, shared_secret: String) -> tokio_tungstenite::tungstenite::Result<()> {
     let ws_stream = tokio_tungstenite::accept_async(stream)
-        .await
-        .expect("Failed to complete the websocket handshake!");
+        .await;
+
+    if ws_stream.is_err() {
+        println!("Failed to complete the websocket handshake! Dropping {}!", peer);
+
+        return Ok(());
+    }
+
+    let ws_stream = ws_stream.unwrap();
+
     println!("Connected to peer: {}!", &peer);
 
     let (mut ws_sender, mut ws_receiver) = ws_stream.split();