tonic: Introduce a new method on Endpoint to override the origin (#1013)

tonic/src/transport/channel/endpoint.rs

@@ -24,6 +24,7 @@ use tower::make::MakeConnection;
 #[derive(Clone)]
 pub struct Endpoint {
     pub(crate) uri: Uri,
+    pub(crate) origin: Option<Uri>,
     pub(crate) user_agent: Option<HeaderValue>,
     pub(crate) timeout: Option<Duration>,
     pub(crate) concurrency_limit: Option<usize>,
@@ -106,6 +107,25 @@ impl Endpoint {
             .map_err(|_| Error::new_invalid_user_agent())
     }
 
+    /// Set a custom origin.
+    ///
+    /// Override the `origin`, mainly useful when you are reaching a Server/LoadBalancer
+    /// which serves multiple services at the same time.
+    /// It will play the role of SNI (Server Name Indication).
+    ///
+    /// ```
+    /// # use tonic::transport::Endpoint;
+    /// # let mut builder = Endpoint::from_static("https://proxy.com");
+    /// builder.origin("https://example.com".parse().expect("http://example.com must be a valid URI"));
+    /// // origin: "https://example.com"
+    /// ```
+    pub fn origin(self, origin: Uri) -> Self {
+        Endpoint {
+            origin: Some(origin),
+            ..self
+        }
+    }
+
     /// Apply a timeout to each request.
     ///
     /// ```
@@ -395,6 +415,7 @@ impl From<Uri> for Endpoint {
     fn from(uri: Uri) -> Self {
         Self {
             uri,
+            origin: None,
             user_agent: None,
             concurrency_limit: None,
             rate_limit: None,

tonic/src/transport/service/add_origin.rs

@@ -1,4 +1,6 @@
 use futures_core::future::BoxFuture;
+use http::uri::Authority;
+use http::uri::Scheme;
 use http::{Request, Uri};
 use std::task::{Context, Poll};
 use tower_service::Service;
@@ -6,12 +8,21 @@ use tower_service::Service;
 #[derive(Debug)]
 pub(crate) struct AddOrigin<T> {
     inner: T,
-    origin: Uri,
+    scheme: Option<Scheme>,
+    authority: Option<Authority>,
 }
 
 impl<T> AddOrigin<T> {
     pub(crate) fn new(inner: T, origin: Uri) -> Self {
-        Self { inner, origin }
+        let http::uri::Parts {
+            scheme, authority, ..
+        } = origin.into_parts();
+
+        Self {
+            inner,
+            scheme,
+            authority,
+        }
     }
 }
 
@@ -30,24 +41,24 @@ where
     }
 
     fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
-        // Split the request into the head and the body.
-        let (mut head, body) = req.into_parts();
-
-        // Split the request URI into parts.
-        let mut uri: http::uri::Parts = head.uri.into();
-        let set_uri = self.origin.clone().into_parts();
-
-        if set_uri.scheme.is_none() || set_uri.authority.is_none() {
+        if self.scheme.is_none() || self.authority.is_none() {
             let err = crate::transport::Error::new_invalid_uri();
             return Box::pin(async move { Err::<Self::Response, _>(err.into()) });
         }
 
-        // Update the URI parts, setting hte scheme and authority
-        uri.scheme = Some(set_uri.scheme.expect("expected scheme"));
-        uri.authority = Some(set_uri.authority.expect("expected authority"));
+        // Split the request into the head and the body.
+        let (mut head, body) = req.into_parts();
 
         // Update the the request URI
-        head.uri = http::Uri::from_parts(uri).expect("valid uri");
+        head.uri = {
+            // Split the request URI into parts.
+            let mut uri: http::uri::Parts = head.uri.into();
+            // Update the URI parts, setting hte scheme and authority
+            uri.scheme = self.scheme.clone();
+            uri.authority = self.authority.clone();
+
+            http::Uri::from_parts(uri).expect("valid uri")
+        };
 
         let request = Request::from_parts(head, body);
 

tonic/src/transport/service/connection.rs

@@ -55,7 +55,11 @@ impl Connection {
         }
 
         let stack = ServiceBuilder::new()
-            .layer_fn(|s| AddOrigin::new(s, endpoint.uri.clone()))
+            .layer_fn(|s| {
+                let origin = endpoint.origin.as_ref().unwrap_or(&endpoint.uri).clone();
+
+                AddOrigin::new(s, origin)
+            })
             .layer_fn(|s| UserAgent::new(s, endpoint.user_agent.clone()))
             .layer_fn(|s| GrpcTimeout::new(s, endpoint.timeout))
             .option_layer(endpoint.concurrency_limit.map(ConcurrencyLimitLayer::new))