From 068a9a69844cdfbcd44c178bfc62a6e24e5d04a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20CORTIER?= <bcortier@proton.me>
Date: Fri, 5 May 2023 01:28:53 -0400
Subject: [PATCH] Fix poll_flush on closed connection

Flushing when there is nothing to send anymore should succeed.

Relevant code in `tungstenite` crate code:
https://github.com/snapview/tungstenite-rs/blob/e5efe537b87a6705467043fe44bb220ddf7c1ce8/src/protocol/mod.rs#L390
---
 src/lib.rs | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index a89fccd..f8ff30f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -374,7 +374,7 @@ where
     fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
         match (*self).with_context(None, |s| s.write_message(item)) {
             Ok(()) => Ok(()),
-            Err(::tungstenite::Error::Io(err)) if err.kind() == std::io::ErrorKind::WouldBlock => {
+            Err(WsError::Io(err)) if err.kind() == std::io::ErrorKind::WouldBlock => {
                 // the message was accepted and queued
                 // isn't an error.
                 Ok(())
@@ -387,7 +387,19 @@ where
     }
 
     fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        (*self).with_context(Some((ContextWaker::Write, cx)), |s| cvt(s.write_pending()))
+        match (*self).with_context(Some((ContextWaker::Write, cx)), |s| s.write_pending()) {
+            Ok(()) => Poll::Ready(Ok(())),
+            Err(WsError::ConnectionClosed) => {
+                // WebSocket is closing and there is nothing to send anymore.
+                // Not an failure, the flush operation is a success.
+                Poll::Ready(Ok(()))
+            }
+            Err(WsError::Io(ref e)) if e.kind() == std::io::ErrorKind::WouldBlock => {
+                trace!("WouldBlock");
+                Poll::Pending
+            }
+            Err(e) => Poll::Ready(Err(e)),
+        }
     }
 
     fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
@@ -400,8 +412,8 @@ where
 
         match res {
             Ok(()) => Poll::Ready(Ok(())),
-            Err(::tungstenite::Error::ConnectionClosed) => Poll::Ready(Ok(())),
-            Err(::tungstenite::Error::Io(err)) if err.kind() == std::io::ErrorKind::WouldBlock => {
+            Err(WsError::ConnectionClosed) => Poll::Ready(Ok(())),
+            Err(WsError::Io(err)) if err.kind() == std::io::ErrorKind::WouldBlock => {
                 trace!("WouldBlock");
                 self.closing = true;
                 Poll::Pending