From 1d2cc257dc8598f63a1a9303e2cc37222fe91d24 Mon Sep 17 00:00:00 2001
From: Sergio Benitez <sb@sergio.bz>
Date: Thu, 30 Mar 2023 15:24:18 -0700
Subject: [PATCH] Preprocess tungstenite 'ConnectionClosed' errors.

Tungstenite abuses `Err(ConnectionClosed)` to indicate the non-error
condition of a websocket closing. This commit changes error processing
such that the error is caught and converted into a successful
termination of websocket handling.
---
 examples/upgrade/src/ws.rs | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/examples/upgrade/src/ws.rs b/examples/upgrade/src/ws.rs
index ace8c5fa..48353186 100644
--- a/examples/upgrade/src/ws.rs
+++ b/examples/upgrade/src/ws.rs
@@ -87,14 +87,23 @@ impl<'r, 'o: 'r, S> Responder<'r, 'o> for MessageStream<'o, S>
     }
 }
 
+/// Returns `Ok(true)` if processing should continue, `Ok(false)` if processing
+/// has terminated without error, and `Err(e)` if an error has occurred.
+fn handle_result(result: Result<()>) -> io::Result<bool> {
+    match result {
+        Ok(_) => Ok(true),
+        Err(Error::ConnectionClosed) => Ok(false),
+        Err(Error::Io(e)) => Err(e),
+        Err(e) => Err(io::Error::new(io::ErrorKind::Other, e))
+    }
+}
+
 #[rocket::async_trait]
 impl IoHandler for Channel<'_> {
     async fn io(&mut self, io: IoStream) -> io::Result<()> {
         let stream = WebSocketStream::from_raw_socket(io, Role::Server, None).await;
-        (self.handler)(stream).await.map_err(|e| match e {
-            Error::Io(e) => e,
-            other => io::Error::new(io::ErrorKind::Other, other)
-        })
+        let result = (self.handler)(stream).await;
+        handle_result(result).map(|_| ())
     }
 }
 
@@ -112,10 +121,9 @@ impl<'r, S> IoHandler for MessageStream<'r, S>
                 Err(e) => Err(e)
             };
 
-            result.map_err(|e| match e {
-                Error::Io(e) => e,
-                other => io::Error::new(io::ErrorKind::Other, other)
-            })?;
+            if !handle_result(result)? {
+                return Ok(());
+            }
         }
 
         Ok(())