From fd2e80cd1b1e52b2e09e12e2985890424411754f Mon Sep 17 00:00:00 2001
From: Anthony Ramine <nox@nox.paris>
Date: Mon, 19 Jul 2021 14:45:10 +0200
Subject: [PATCH] fix(h2): improve I/O errors emitted by H2Upgraded

---
 src/proto/h2/mod.rs | 43 ++++++++++++++++++++++++++++++-------------
 1 file changed, 30 insertions(+), 13 deletions(-)

diff --git a/src/proto/h2/mod.rs b/src/proto/h2/mod.rs
index ed8fbef74a..1900a6a31c 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,34 @@ 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)),
-        })
+
+        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>> {