diff --git a/tests/integration_tests/proto/test.proto b/tests/integration_tests/proto/test.proto index 53df288..653ceba 100644 --- a/tests/integration_tests/proto/test.proto +++ b/tests/integration_tests/proto/test.proto @@ -11,6 +11,8 @@ message Output {} service Test1 { rpc UnaryCall(Input1) returns (Output1); + + rpc StreamCall(Input1) returns (stream Output1); } message Input1 { diff --git a/tests/integration_tests/tests/max_message_size.rs b/tests/integration_tests/tests/max_message_size.rs index 14dbec3..f79a416 100644 --- a/tests/integration_tests/tests/max_message_size.rs +++ b/tests/integration_tests/tests/max_message_size.rs @@ -1,3 +1,6 @@ +use std::pin::Pin; + +use futures::{stream, Stream}; use integration_tests::{ pb::{test1_client, test1_server, Input1, Output1}, trace_init, @@ -110,6 +113,81 @@ fn max_message_send_size() { }); } +#[tokio::test] +async fn response_stream_limit() { + let client_blob = vec![0; 1]; + + let (client, server) = tokio::io::duplex(1024); + + struct Svc; + + #[tonic::async_trait] + impl test1_server::Test1 for Svc { + async fn unary_call(&self, _req: Request) -> Result, Status> { + unimplemented!() + } + + type StreamCallStream = + Pin> + Send + 'static>>; + + async fn stream_call( + &self, + _req: Request, + ) -> Result, Status> { + let blob = Output1 { + buf: vec![0; 6877902], + }; + let stream = stream::iter(vec![Ok(blob.clone()), Ok(blob.clone())]); + + Ok(Response::new(Box::pin(stream))) + } + } + + let svc = test1_server::Test1Server::new(Svc); + + tokio::spawn(async move { + Server::builder() + .add_service(svc) + .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)])) + .await + .unwrap(); + }); + + // Move client to an option so we can _move_ the inner value + // on the first attempt to connect. All other attempts will fail. + let mut client = Some(client); + let channel = Endpoint::try_from("http://[::]:50051") + .unwrap() + .connect_with_connector(tower::service_fn(move |_| { + let client = client.take(); + + async move { + if let Some(client) = client { + Ok(client) + } else { + Err(std::io::Error::new( + std::io::ErrorKind::Other, + "Client already taken", + )) + } + } + })) + .await + .unwrap(); + + let client = test1_client::Test1Client::new(channel); + + let mut client = client.max_decoding_message_size(6877902 + 5); + + let req = Request::new(Input1 { + buf: client_blob.clone(), + }); + + let mut stream = client.stream_call(req).await.unwrap().into_inner(); + + while let Some(_b) = stream.message().await.unwrap() {} +} + // Track caller doesn't work on async fn so we extract the async part // into a sync version and assert the response there using track track_caller // so that when this does panic it tells us which line in the test failed not @@ -210,6 +288,16 @@ async fn max_message_run(case: &TestCase) -> Result<(), Status> { buf: self.0.clone(), })) } + + type StreamCallStream = + Pin> + Send + 'static>>; + + async fn stream_call( + &self, + _req: Request, + ) -> Result, Status> { + unimplemented!() + } } let svc = test1_server::Test1Server::new(Svc(server_blob));