From f51c677dec9debf60cb336dc938bae103adf17a0 Mon Sep 17 00:00:00 2001
From: Anthony Ramine <123095+nox@users.noreply.github.com>
Date: Wed, 21 Jul 2021 17:17:05 +0200
Subject: [PATCH] fix(http2): improve I/O errors emitted by H2Upgraded (#2598)

When a `CONNECT` over HTTP2 has been established, and the user tries to write data right when the peer closes the stream, it will no longer return as a "user error". The reset code is checked, and converted into an appropriate `io::ErrorKind`.
---
 src/proto/h2/mod.rs | 45 ++++++++++++++++++++++++++++++++-------------
 1 file changed, 32 insertions(+), 13 deletions(-)

diff --git a/src/proto/h2/mod.rs b/src/proto/h2/mod.rs
index ed8fbef74a..b8312aff64 100644
--- a/src/proto/h2/mod.rs
+++ b/src/proto/h2/mod.rs
@@ -1,5 +1,5 @@
 use bytes::{Buf, Bytes};
-use h2::{RecvStream, SendStream};
+use h2::{Reason, RecvStream, SendStream};
 use http::header::{HeaderName, CONNECTION, TE, TRAILER, TRANSFER_ENCODING, UPGRADE};
 use http::HeaderMap;
 use pin_project_lite::pin_project;
@@ -313,7 +313,11 @@ where
                         break buf;
                     }
                     Some(Err(e)) => {
-                        return Poll::Ready(Err(h2_to_io_error(e)));
+                        return Poll::Ready(match e.reason() {
+                            Some(Reason::NO_ERROR) | Some(Reason::CANCEL) => Ok(()),
+                            Some(Reason::STREAM_CLOSED) => Err(io::ErrorKind::BrokenPipe.into()),
+                            _ => Err(h2_to_io_error(e)),
+                        })
                     }
                 }
             };
@@ -335,21 +339,36 @@ where
         cx: &mut Context<'_>,
         buf: &[u8],
     ) -> Poll<Result<usize, io::Error>> {
-        if let Poll::Ready(reset) = self.send_stream.poll_reset(cx) {
-            return Poll::Ready(Err(h2_to_io_error(match reset {
-                Ok(reason) => reason.into(),
-                Err(e) => e,
-            })));
-        }
         if buf.is_empty() {
             return Poll::Ready(Ok(0));
         }
         self.send_stream.reserve_capacity(buf.len());
-        Poll::Ready(match ready!(self.send_stream.poll_capacity(cx)) {
-            None => Ok(0),
-            Some(Ok(cnt)) => self.send_stream.write(&buf[..cnt], false).map(|()| cnt),
-            Some(Err(e)) => Err(h2_to_io_error(e)),
-        })
+
+        // We ignore all errors returned by `poll_capacity` and `write`, as we
+        // will get the correct from `poll_reset` anyway.
+        let cnt = match ready!(self.send_stream.poll_capacity(cx)) {
+            None => Some(0),
+            Some(Ok(cnt)) => self
+                .send_stream
+                .write(&buf[..cnt], false)
+                .ok()
+                .map(|()| cnt),
+            Some(Err(_)) => None,
+        };
+
+        if let Some(cnt) = cnt {
+            return Poll::Ready(Ok(cnt));
+        }
+
+        Poll::Ready(Err(h2_to_io_error(
+            match ready!(self.send_stream.poll_reset(cx)) {
+                Ok(Reason::NO_ERROR) | Ok(Reason::CANCEL) | Ok(Reason::STREAM_CLOSED) => {
+                    return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
+                }
+                Ok(reason) => reason.into(),
+                Err(e) => e,
+            },
+        )))
     }
 
     fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {