check final response in poll_informational

src/proto/streams/recv.rs

@@ -377,6 +377,13 @@ impl Recv {
         // If it's not, we put it back
         if let Some(event) = stream.pending_recv.pop_front(&mut self.buffer) {
             match event {
+                Event::Headers(Client(response)) => {
+                    // Final response
+                    stream
+                        .pending_recv
+                        .push_front(&mut self.buffer, Event::Headers(Client(response)));
+                    return Poll::Ready(None);
+                }
                 Event::InformationalHeaders(Client(response)) => {
                     // Found an informational response, return it
                     return Poll::Ready(Some(Ok(response)));

tests/h2-tests/tests/informational_responses.rs

@@ -233,6 +233,80 @@ async fn invalid_informational_status_returns_error() {
     join(client, srv).await;
 }
 
+#[tokio::test]
+async fn client_poll_informational_responses_none() {
+    h2_support::trace_init!();
+    let (io, mut srv) = mock::new();
+
+    let (sync_sender, sync_receiver) = tokio::sync::oneshot::channel::<()>();
+
+    let srv = async move {
+        let recv_settings = srv.assert_client_handshake().await;
+        assert_default_settings!(recv_settings);
+
+        srv.recv_frame(
+            frames::headers(1)
+                .request("GET", "https://example.com/")
+                .eos(),
+        )
+        .await;
+
+        // Send final response directly
+        srv.send_frame(frames::headers(1).response(StatusCode::OK))
+            .await;
+
+        // The server may not close the stream immediately.
+        // Let's simulate this by waiting from client.
+        // Continue after the client received the response headers
+        tokio::time::timeout(Duration::from_secs(4), sync_receiver)
+            .await
+            .expect("Client blocked on informational headers")
+            .unwrap();
+        srv.send_frame(frames::data(1, b"request body").eos()).await;
+    };
+
+    let client = async move {
+        let (client, connection) = client::handshake(io).await.expect("handshake");
+
+        let request = Request::builder()
+            .method("GET")
+            .uri("https://example.com/")
+            .body(())
+            .unwrap();
+
+        let (mut response_future, _) = client
+            .ready()
+            .await
+            .unwrap()
+            .send_request(request, true)
+            .unwrap();
+
+        tokio::spawn(async move {
+            connection.await.expect("connection error");
+        });
+
+        // Poll for informational responses
+        loop {
+            match poll_fn(|cx| response_future.poll_informational(cx)).await {
+                Some(Ok(rsp)) => panic!("Unexpected informational response {:?}", rsp),
+                Some(Err(e)) => panic!("Error polling informational: {:?}", e),
+                None => break,
+            }
+        }
+        // Let the server continue sending responses
+        sync_sender.send(()).unwrap();
+
+        // Get the final response
+        let response = response_future.await.expect("response error");
+        assert_eq!(response.status(), StatusCode::OK);
+        let (_hdr, mut recv_stream) = response.into_parts();
+        let data = recv_stream.data().await.unwrap().unwrap();
+        assert_eq!("request body", data);
+    };
+
+    join(srv, client).await;
+}
+
 #[tokio::test]
 async fn client_poll_informational_responses() {
     h2_support::trace_init!();