fix: preserve received END_STREAM on connection error and EOF

src/proto/streams/state.rs

@@ -76,7 +76,8 @@ enum Peer {
 enum Cause {
     EndStream,
     Error(Error),
-    /// The stream was reset after the receive half had already reached EOS.
+    /// The stream was reset, or the connection failed, after the receive half
+    /// had already reached EOS.
     ErrorAfterEndStream(Error),
 
     /// This indicates to the connection that a reset frame must be sent out
@@ -297,27 +298,36 @@ impl State {
 
     /// Handle a connection-level error.
     pub fn handle_error(&mut self, err: &proto::Error) {
+        let recv_end_stream = self.is_recv_end_stream();
         match self.inner {
             Closed(..) => {}
             _ => {
                 tracing::trace!("handle_error; err={:?}", err);
-                self.inner = Closed(Cause::Error(err.clone()));
+                self.inner = Closed(if recv_end_stream {
+                    Cause::ErrorAfterEndStream(err.clone())
+                } else {
+                    Cause::Error(err.clone())
+                });
             }
         }
     }
 
     pub fn recv_eof(&mut self) {
+        let recv_end_stream = self.is_recv_end_stream();
         match self.inner {
             Closed(..) => {}
             ref state => {
                 tracing::trace!("recv_eof; state={:?}", state);
-                self.inner = Closed(Cause::Error(
-                    io::Error::new(
-                        io::ErrorKind::BrokenPipe,
-                        "stream closed because of a broken pipe",
-                    )
-                    .into(),
-                ));
+                let error = io::Error::new(
+                    io::ErrorKind::BrokenPipe,
+                    "stream closed because of a broken pipe",
+                )
+                .into();
+                self.inner = Closed(if recv_end_stream {
+                    Cause::ErrorAfterEndStream(error)
+                } else {
+                    Cause::Error(error)
+                });
             }
         }
     }
@@ -519,4 +529,41 @@ mod tests {
             Some(Reason::NO_ERROR)
         );
     }
+
+    fn half_closed_remote() -> State {
+        let mut state = State::default();
+        state.send_open(false).unwrap();
+
+        let mut headers =
+            frame::Headers::new(StreamId::from(1), Default::default(), HeaderMap::new());
+        headers.set_end_stream();
+        state.recv_open(&headers).unwrap();
+        assert!(state.is_recv_end_stream());
+        state
+    }
+
+    #[test]
+    fn handle_error_preserves_received_end_stream() {
+        let mut state = half_closed_remote();
+
+        state.handle_error(&Error::library_go_away(Reason::INTERNAL_ERROR));
+
+        assert!(state.is_recv_end_stream());
+        assert_eq!(state.ensure_recv_open().unwrap(), false);
+        assert_eq!(
+            state.ensure_reason(PollReset::Streaming).unwrap(),
+            Some(Reason::INTERNAL_ERROR)
+        );
+    }
+
+    #[test]
+    fn recv_eof_preserves_received_end_stream() {
+        let mut state = half_closed_remote();
+
+        state.recv_eof();
+
+        assert!(state.is_recv_end_stream());
+        assert_eq!(state.ensure_recv_open().unwrap(), false);
+        assert!(state.ensure_reason(PollReset::Streaming).is_err());
+    }
 }

tests/h2-tests/tests/stream_states.rs

@@ -1231,6 +1231,86 @@ async fn assert_recv_end_stream_survives_reset(reset_reason: Reason) {
     join(srv, client).await;
 }
 
+#[tokio::test]
+async fn recv_end_stream_survives_go_away() {
+    h2_support::trace_init!();
+    let (io, mut srv) = mock::new();
+
+    let srv = async move {
+        let settings = srv.assert_client_handshake().await;
+        assert_default_settings!(settings);
+        srv.recv_frame(frames::headers(1).request("POST", "https://example.com/"))
+            .await;
+        srv.send_frame(frames::headers(1).response(200).eos()).await;
+        // last_stream_id 0 makes the GOAWAY close stream 1.
+        srv.send_frame(frames::go_away(0).reason(Reason::INTERNAL_ERROR))
+            .await;
+    };
+
+    let client = async move {
+        let (mut client, mut conn) = client::handshake(io).await.expect("handshake");
+        let request = Request::builder()
+            .method(Method::POST)
+            .uri("https://example.com/")
+            .body(())
+            .unwrap();
+        let (response, mut send_stream) = client.send_request(request, false).unwrap();
+
+        // Process the GOAWAY before polling the response future.
+        let reason = conn
+            .drive(poll_fn(move |cx| send_stream.poll_reset(cx)))
+            .await
+            .unwrap();
+        assert_eq!(reason, Reason::INTERNAL_ERROR);
+
+        let response = response.await.unwrap();
+        assert!(response.body().is_end_stream());
+
+        let mut body = response.into_body();
+        assert!(body.data().await.is_none());
+        assert!(body.trailers().await.unwrap().is_none());
+    };
+
+    join(srv, client).await;
+}
+
+#[tokio::test]
+async fn recv_end_stream_survives_eof() {
+    h2_support::trace_init!();
+    let (io, mut srv) = mock::new();
+
+    let srv = async move {
+        let settings = srv.assert_client_handshake().await;
+        assert_default_settings!(settings);
+        srv.recv_frame(frames::headers(1).request("POST", "https://example.com/"))
+            .await;
+        srv.send_frame(frames::headers(1).response(200).eos()).await;
+        // Dropping the mock closes the connection.
+    };
+
+    let client = async move {
+        let (mut client, conn) = client::handshake(io).await.expect("handshake");
+        let request = Request::builder()
+            .method(Method::POST)
+            .uri("https://example.com/")
+            .body(())
+            .unwrap();
+        let (response, _send_stream) = client.send_request(request, false).unwrap();
+
+        // Run the connection until it sees EOF.
+        conn.await.unwrap();
+
+        let response = response.await.unwrap();
+        assert!(response.body().is_end_stream());
+
+        let mut body = response.into_body();
+        assert!(body.data().await.is_none());
+        assert!(body.trailers().await.unwrap().is_none());
+    };
+
+    join(srv, client).await;
+}
+
 #[tokio::test]
 async fn rst_while_closing() {
     // Test to reproduce panic in issue #246 --- receipt of a RST_STREAM frame