From ffdb47883190a8889cf30b716294383392a763c5 Mon Sep 17 00:00:00 2001
From: Josh Leeb-du Toit <josh.leebdutoit@gmail.com>
Date: Tue, 26 Jun 2018 04:14:35 +1000
Subject: [PATCH] feat(http2): quickly cancel when receiving RST_STREAM

Update Http2 proto to cancel quick when the stream is reset, on an
`RST_STREAM` frame.

Closes: #1549
---
 src/proto/h2/mod.rs    | 14 ++++++++++++++
 src/proto/h2/server.rs | 17 ++++++++++++++++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/src/proto/h2/mod.rs b/src/proto/h2/mod.rs
index 8742d70c1c..de877f0bc7 100644
--- a/src/proto/h2/mod.rs
+++ b/src/proto/h2/mod.rs
@@ -116,6 +116,13 @@ where
                             None => return Err(::Error::new_canceled(None::<::Error>)),
                         }
                     }
+                } else {
+                    if let Async::Ready(reason) =
+                        self.body_tx.poll_reset().map_err(|e| ::Error::new_h2(e))?
+                    {
+                        debug!("stream received RST_STREAM: {:?}", reason);
+                        return Err(::Error::new_h2(reason.into()));
+                    }
                 }
 
                 match try_ready!(self.stream.poll_data().map_err(|e| self.on_err(e))) {
@@ -148,6 +155,13 @@ where
                     }
                 }
             } else {
+                if let Async::Ready(reason) =
+                    self.body_tx.poll_reset().map_err(|e| ::Error::new_h2(e))?
+                {
+                    debug!("stream received RST_STREAM: {:?}", reason);
+                    return Err(::Error::new_h2(reason.into()));
+                }
+
                 match try_ready!(self.stream.poll_trailers().map_err(|e| self.on_err(e))) {
                     Some(trailers) => {
                         self.body_tx
diff --git a/src/proto/h2/server.rs b/src/proto/h2/server.rs
index 6777bfc109..d78d4a254e 100644
--- a/src/proto/h2/server.rs
+++ b/src/proto/h2/server.rs
@@ -174,7 +174,22 @@ where
         loop {
             let next = match self.state {
                 H2StreamState::Service(ref mut h) => {
-                    let res = try_ready!(h.poll().map_err(::Error::new_user_service));
+                    let res = match h.poll() {
+                        Ok(Async::Ready(r)) => r,
+                        Ok(Async::NotReady) => {
+                            // Body is not yet ready, so we want to check if the client has sent a
+                            // RST_STREAM frame which would cancel the current request.
+                            if let Async::Ready(reason) =
+                                self.reply.poll_reset().map_err(|e| ::Error::new_h2(e))?
+                            {
+                                debug!("stream received RST_STREAM: {:?}", reason);
+                                return Err(::Error::new_h2(reason.into()));
+                            }
+                            return Ok(Async::NotReady);
+                        }
+                        Err(e) => return Err(::Error::new_user_service(e)),
+                    };
+
                     let (head, body) = res.into_parts();
                     let mut res = ::http::Response::from_parts(head, ());
                     super::strip_connection_headers(res.headers_mut());