Continue listening after failed WebSocket handshake.

client/src/control/ws.rs

@@ -54,40 +54,53 @@ async fn forward_outgoing(
 }
 
 async fn handle(
-  stream: WebSocketStream<TcpStream>,
+  stream: TcpStream,
   remote_address: &SocketAddr,
   message_tx: &mpsc::UnboundedSender<Message>,
   response_rx: &mut mpsc::Receiver<Response>,
 ) {
-  let (mut outgoing, incoming) = stream.split();
+  let ws_stream = match tokio_tungstenite::accept_async(stream).await {
+    Ok(ws_stream) => ws_stream,
+    Err(err) => {
+      warn!(
+        "Failed to accept WebSocket connection from {}: {}",
+        remote_address, err
+      );
+      return;
+    }
+  };
+
+  info!("WebSocket connection established from {}", remote_address);
+
+  let (mut outgoing, incoming) = ws_stream.split();
 
   tokio::select! {
     result = forward_incoming(incoming, message_tx) => match result {
       Ok(()) => info!(
-        "Incoming websocket handler task for {} stopped",
+        "Incoming WebSocket handler task for {} stopped",
         remote_address,
       ),
       Err(err) => error!(
-        "Error in websocket handler task for {}: {:#}",
+        "Error in WebSocket handler task for {}: {:#}",
         remote_address, err,
       ),
     },
     result = forward_outgoing(response_rx, &mut outgoing) => match result {
       Ok(()) => info!(
-        "Outgoing websocket handler for {} stopped",
+        "Outgoing WebSocket handler for {} stopped",
         remote_address,
       ),
       Err(err) => warn!(
-        "Error in outgoing websocket handler for {}: {:#}",
+        "Error in outgoing WebSocket handler for {}: {:#}",
         remote_address, err,
       ),
     },
   };
 
   match outgoing.close().await {
-    Ok(()) => info!("Closed websocket for {}", remote_address),
+    Ok(()) => info!("Closed WebSocket for {}", remote_address),
     Err(err) => {
-      error!("Error closing websocket for {}: {}", remote_address, err,)
+      error!("Error closing WebSocket for {}: {}", remote_address, err,)
     }
   };
 }
@@ -125,7 +138,7 @@ impl Listener {
     info!("Accepting control connections on {}", self.address());
 
     loop {
-      let (raw_stream, remote_address) = tokio::select! {
+      let (stream, remote_address) = tokio::select! {
         result = self.inner.accept() => {
           result.context("accepting connection")?
         },
@@ -146,11 +159,7 @@ impl Listener {
 
       info!("Accepted control connection from {}", remote_address);
 
-      // TODO: Continue iterating in case of error.
-      let ws_stream = tokio_tungstenite::accept_async(raw_stream).await?;
-      info!("WebSocket connection established from {}", remote_address);
-
-      handle(ws_stream, &remote_address, &message_tx, &mut response_rx).await
+      handle(stream, &remote_address, &message_tx, &mut response_rx).await
     }
 
     Ok(())