Handle stream error correctly (#2199)

Suficio · LucioFranco · web-flow · commit 8be906daae10 · 2025-03-25T18:49:19.000Z
Co-authored-by: Lucio Franco &lt;luciofranco14@gmail.com&gt;
diff --git a/tests/integration_tests/tests/status.rs b/tests/integration_tests/tests/status.rs
@@ -6,9 +6,12 @@ use integration_tests::pb::{
     test_client, test_server, test_stream_client, test_stream_server, Input, InputStream, Output,
     OutputStream,
 };
+use integration_tests::BoxFuture;
 use std::error::Error;
+use std::task::{Context, Poll};
 use std::time::Duration;
 use tokio::{net::TcpListener, sync::oneshot};
+use tonic::body::Body;
 use tonic::metadata::{MetadataMap, MetadataValue};
 use tonic::{
     transport::{server::TcpIncoming, Endpoint, Server},
@@ -209,6 +212,93 @@ async fn status_from_server_stream_with_source() {
     source.downcast_ref::<tonic::transport::Error>().unwrap();
 }
 
+#[tokio::test]
+async fn status_from_server_stream_with_inferred_status() {
+    integration_tests::trace_init();
+
+    struct Svc;
+
+    #[tonic::async_trait]
+    impl test_stream_server::TestStream for Svc {
+        type StreamCallStream = Stream<OutputStream>;
+
+        async fn stream_call(
+            &self,
+            _: Request<InputStream>,
+        ) -> Result<Response<Self::StreamCallStream>, Status> {
+            let s = tokio_stream::once(Ok(OutputStream {}));
+            Ok(Response::new(Box::pin(s) as Self::StreamCallStream))
+        }
+    }
+
+    #[derive(Clone)]
+    struct TestLayer;
+
+    impl<S> tower::Layer<S> for TestLayer {
+        type Service = TestService;
+
+        fn layer(&self, _: S) -> Self::Service {
+            TestService
+        }
+    }
+
+    #[derive(Clone)]
+    struct TestService;
+
+    impl tower::Service<http::Request<Body>> for TestService {
+        type Response = http::Response<Body>;
+        type Error = std::convert::Infallible;
+        type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn call(&mut self, _: http::Request<Body>) -> Self::Future {
+            Box::pin(async {
+                Ok(http::Response::builder()
+                    .status(http::StatusCode::BAD_GATEWAY)
+                    .body(Body::empty())
+                    .unwrap())
+            })
+        }
+    }
+
+    let svc = test_stream_server::TestStreamServer::new(Svc);
+
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let incoming: TcpIncoming = TcpIncoming::from(listener).with_nodelay(Some(true));
+
+    tokio::spawn(async move {
+        Server::builder()
+            .layer(TestLayer)
+            .add_service(svc)
+            .serve_with_incoming(incoming)
+            .await
+            .unwrap();
+    });
+
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
+    let mut client = test_stream_client::TestStreamClient::connect(format!("http://{addr}"))
+        .await
+        .unwrap();
+
+    let mut stream = client
+        .stream_call(InputStream {})
+        .await
+        .unwrap()
+        .into_inner();
+
+    assert_eq!(
+        stream.message().await.unwrap_err().code(),
+        Code::Unavailable
+    );
+
+    assert_eq!(stream.message().await.unwrap(), None);
+}
+
 #[tokio::test]
 async fn message_and_then_status_from_server_stream() {
     integration_tests::trace_init();
diff --git a/tonic/src/codec/decode.rs b/tonic/src/codec/decode.rs
@@ -400,14 +400,12 @@ impl<T> Stream for Streaming<T> {
             }
 
             if ready!(self.inner.poll_frame(cx))?.is_none() {
-                break;
+                match self.inner.response() {
+                    Ok(()) => return Poll::Ready(None),
+                    Err(err) => self.inner.state = State::Error(Some(err)),
+                }
             }
         }
-
-        Poll::Ready(match self.inner.response() {
-            Ok(()) => None,
-            Err(err) => Some(Err(err)),
-        })
     }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -400,14 +400,12 @@ impl<T> Stream for Streaming<T> {`
`400`	`400`	`}`
`401`	`401`
`402`	`402`	`if ready!(self.inner.poll_frame(cx))?.is_none() {`
`403`		`- break;`
	`403`	`+ match self.inner.response() {`
	`404`	`+ Ok(()) => return Poll::Ready(None),`
	`405`	`+ Err(err) => self.inner.state = State::Error(Some(err)),`
	`406`	`+ }`
`404`	`407`	`}`
`405`	`408`	`}`
`406`		`-`
`407`		`- Poll::Ready(match self.inner.response() {`
`408`		`- Ok(()) => None,`
`409`		`- Err(err) => Some(Err(err)),`
`410`		`- })`
`411`	`409`	`}`
`412`	`410`	`}`
`413`	`411`