From 1dd08bedaa997855c4c28dfc03ca23c650f6b2c6 Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@abcum.com>
Date: Tue, 20 Sep 2022 23:04:40 +0100
Subject: [PATCH] Handle errors properly when WebSocket client disconnects
 improperly

Closes #181
Closes #195
Closes #140
---
 src/net/rpc.rs | 28 ++++++++++++++++++++++++----
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/net/rpc.rs b/src/net/rpc.rs
index bda2f1d8..35c48e11 100644
--- a/src/net/rpc.rs
+++ b/src/net/rpc.rs
@@ -3,6 +3,7 @@ use crate::cnf::MAX_CONCURRENT_CALLS;
 use crate::dbs::DB;
 use crate::err::Error;
 use crate::net::session;
+use crate::net::LOG;
 use crate::rpc::args::Take;
 use crate::rpc::paths::{ID, METHOD, PARAMS};
 use crate::rpc::res::Failure;
@@ -60,15 +61,34 @@ impl Rpc {
 		let (mut wtx, mut wrx) = ws.split();
 		// Send messages to the client
 		tokio::task::spawn(async move {
+			// Wait for the next message to send
 			while let Some(res) = rcv.next().await {
-				let _ = wtx.send(res).await;
+				// Send the message to the client
+				if let Err(err) = wtx.send(res).await {
+					// Output the WebSocket error to the logs
+					trace!(target: LOG, "WebSocket error: {:?}", err);
+					// It's already failed, so ignore error
+					let _ = wtx.close().await;
+					// Exit out of the loop
+					break;
+				}
 			}
 		});
 		// Get messages from the client
 		while let Some(msg) = wrx.next().await {
-			if let Ok(msg) = msg {
-				if msg.is_text() {
-					tokio::task::spawn(Rpc::call(rpc.clone(), msg, chn.clone()));
+			match msg {
+				// We've received a message from the client
+				Ok(msg) => {
+					if msg.is_text() {
+						tokio::task::spawn(Rpc::call(rpc.clone(), msg, chn.clone()));
+					}
+				}
+				// There was an error receiving the message
+				Err(err) => {
+					// Output the WebSocket error to the logs
+					trace!(target: LOG, "WebSocket error: {:?}", err);
+					// Exit out of the loop
+					break;
 				}
 			}
 		}