Update websocket error handling (#148)

* Update websocket error handling
* Format

Author: SinaKhalili

src/websocket.rs

@@ -14,6 +14,7 @@ use rust_decimal::Decimal;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use serde_json::json;
 use tokio::{
+    io::AsyncWriteExt,
     net::{TcpListener, TcpStream},
     sync::Mutex,
     task::JoinHandle,
@@ -50,13 +51,41 @@ pub async fn start_ws_server(
 }
 
 async fn accept_connection(
-    stream: TcpStream,
+    mut stream: TcpStream,
     ws_client: Arc<PubsubClient>,
     wallet: Arc<Wallet>,
     program_data: &'static ProgramData,
 ) {
     let addr = stream.peer_addr().expect("peer address");
-    let ws_stream = accept_async(stream).await.expect("Ws handshake");
+
+    // Check for WebSocket upgrade header before accept_async consumes the stream
+    let mut buf = [0u8; 1024];
+    let n = stream.peek(&mut buf).await.unwrap_or(0);
+    if !buf[..n]
+        .windows(7)
+        .any(|w| w.eq_ignore_ascii_case(b"upgrade"))
+    {
+        warn!(target: LOG_TARGET, "non-WebSocket request from {}, rejecting", addr);
+        let _ = stream
+            .write_all(
+                b"HTTP/1.1 426 Upgrade Required\r\n\
+              Content-Type: text/plain\r\n\
+              Upgrade: websocket\r\n\
+              Connection: Upgrade\r\n\
+              \r\n\
+              This is a WebSocket endpoint. Use a WebSocket client to connect.\n",
+            )
+            .await;
+        return;
+    }
+
+    let ws_stream = match accept_async(stream).await {
+        Ok(ws) => ws,
+        Err(err) => {
+            warn!(target: LOG_TARGET, "Ws handshake failed from {}: {}", addr, err);
+            return;
+        }
+    };
     info!(target: LOG_TARGET, "accepted Ws connection: {}", addr);
 
     let (mut ws_out, mut ws_in) = ws_stream.split();