diff --git a/tests/integration_tests/tests/connection.rs b/tests/integration_tests/tests/connection.rs index d7d4cd49b..126ebe03f 100644 --- a/tests/integration_tests/tests/connection.rs +++ b/tests/integration_tests/tests/connection.rs @@ -1,3 +1,4 @@ +use http::{header::HeaderName, HeaderValue}; use integration_tests::pb::{test_client::TestClient, test_server, Input, Output}; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -6,6 +7,7 @@ use tonic::{ transport::{server::TcpIncoming, Endpoint, Server}, Code, Request, Response, Status, }; +use tower_http::set_header::SetRequestHeaderLayer; struct Svc(Arc>>>); @@ -69,6 +71,63 @@ async fn connect_returns_err_via_call_after_connected() { jh.await.unwrap(); } +#[tokio::test] +async fn endpoint_layer_stacks_and_applies_to_requests() { + struct HeaderSvc; + + #[tonic::async_trait] + impl test_server::Test for HeaderSvc { + async fn unary_call(&self, req: Request) -> Result, Status> { + // Both layers must have been applied for this to succeed, proving the + // layers stack rather than the second `layer` call replacing the first. + match ( + req.metadata().get("x-first"), + req.metadata().get("x-second"), + ) { + (Some(_), Some(_)) => Ok(Response::new(Output {})), + _ => Err(Status::internal("a header set by a layer is missing")), + } + } + } + + let (tx, rx) = oneshot::channel(); + let svc = test_server::TestServer::new(HeaderSvc); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let incoming = TcpIncoming::from(listener).with_nodelay(Some(true)); + + let jh = tokio::spawn(async move { + Server::builder() + .add_service(svc) + .serve_with_incoming_shutdown(incoming, async { drop(rx.await) }) + .await + .unwrap(); + }); + + // The layers are configured directly on the `Endpoint`, so no type + // parameters leak onto the resulting `Channel`. + let channel = Endpoint::from_shared(format!("http://{addr}")) + .unwrap() + .layer(SetRequestHeaderLayer::overriding( + HeaderName::from_static("x-first"), + HeaderValue::from_static("first"), + )) + .layer(SetRequestHeaderLayer::overriding( + HeaderName::from_static("x-second"), + HeaderValue::from_static("second"), + )) + .connect_lazy(); + + let mut client = TestClient::new(channel); + + tokio::time::sleep(Duration::from_millis(100)).await; + client.unary_call(Request::new(Input {})).await.unwrap(); + + tx.send(()).unwrap(); + jh.await.unwrap(); +} + #[tokio::test] async fn connect_lazy_reconnects_after_first_failure() { let (tx, rx) = oneshot::channel(); diff --git a/tonic/Cargo.toml b/tonic/Cargo.toml index 69aef608c..c2f8c3c01 100644 --- a/tonic/Cargo.toml +++ b/tonic/Cargo.toml @@ -133,6 +133,7 @@ allowed_external_types = [ "tower_layer::Layer", "tower_layer::stack::Stack", "tower_layer::identity::Identity", + "tower::util::boxed::sync::BoxService", ] [[bench]] diff --git a/tonic/src/transport/channel/endpoint.rs b/tonic/src/transport/channel/endpoint.rs index 4a63746ce..4c03f9a77 100644 --- a/tonic/src/transport/channel/endpoint.rs +++ b/tonic/src/transport/channel/endpoint.rs @@ -5,11 +5,12 @@ use super::ClientTlsConfig; use super::service::TlsConnector; use super::service::{self, Executor, SharedExec}; use super::uds_connector::UdsConnector; +use crate::body::Body; use crate::transport::Error; #[cfg(feature = "_tls-any")] use crate::transport::error; use bytes::Bytes; -use http::{HeaderValue, uri::Uri}; +use http::{HeaderValue, Request, Response, uri::Uri}; use hyper::rt; use hyper_util::client::legacy::connect::HttpConnector; #[cfg(feature = "_tls-any")] @@ -17,8 +18,22 @@ use std::sync::Arc; use std::{fmt, future::Future, net::IpAddr, pin::Pin, str, str::FromStr, time::Duration}; #[cfg(feature = "_tls-any")] use tokio_rustls::rustls::client::danger::ServerCertVerifier; +use tower::layer::Layer; +use tower::layer::util::Stack; +use tower::util::{BoxLayer, BoxService}; use tower_service::Service; +/// A boxed [`Layer`] applied to the [`Connection`](super::service::Connection) service. +/// +/// The layer wraps the boxed connection service, allowing arbitrary `tower` +/// middleware to be added to an [`Endpoint`] without leaking type parameters. +pub(crate) type BoxedLayer = BoxLayer< + BoxService, Response, crate::BoxError>, + Request, + Response, + crate::BoxError, +>; + #[derive(Clone, PartialEq, Eq, Hash)] pub(crate) enum EndpointType { Uri(Uri), @@ -56,6 +71,7 @@ pub struct Endpoint { pub(crate) http2_adaptive_window: Option, pub(crate) local_address: Option, pub(crate) executor: SharedExec, + pub(crate) layer: Option, } impl Endpoint { @@ -105,6 +121,7 @@ impl Endpoint { http2_adaptive_window: None, executor: SharedExec::tokio(), local_address: None, + layer: None, } } @@ -136,6 +153,7 @@ impl Endpoint { http2_adaptive_window: None, executor: SharedExec::tokio(), local_address: None, + layer: None, } } @@ -519,6 +537,29 @@ impl Endpoint { self } + /// Add a [`tower::Layer`] to the stack of middleware applied to each + /// connection. + /// + /// Calling this method multiple times stacks the layers, with the most + /// recently added layer wrapping the previously added ones. + pub fn layer(mut self, layer: L) -> Self + where + L: Layer, Response, crate::BoxError>> + + Send + + Sync + + 'static, + L::Service: Service, Response = Response, Error = crate::BoxError> + + Send + + 'static, + >>::Future: Send + 'static, + { + self.layer = Some(match self.layer.take() { + Some(existing) => BoxLayer::new(Stack::new(existing, layer)), + None => BoxLayer::new(layer), + }); + self + } + pub(crate) fn connector(&self, c: C) -> service::Connector { service::Connector::new( c, diff --git a/tonic/src/transport/channel/service/connection.rs b/tonic/src/transport/channel/service/connection.rs index 50ac77b23..81be04629 100644 --- a/tonic/src/transport/channel/service/connection.rs +++ b/tonic/src/transport/channel/service/connection.rs @@ -80,9 +80,13 @@ impl Connection { let conn = Reconnect::new(make_service, endpoint.uri().clone(), is_lazy); - Self { - inner: BoxService::new(stack.layer(conn)), - } + let inner = BoxService::new(stack.layer(conn)); + let inner = match &endpoint.layer { + Some(layer) => layer.layer(inner), + None => inner, + }; + + Self { inner } } pub(crate) async fn connect(