From 112110e1060b832725d02dd72a11387c3ab31a54 Mon Sep 17 00:00:00 2001
From: Lucio Franco <luciofranco14@gmail.com>
Date: Wed, 15 Jun 2022 12:02:34 -0400
Subject: [PATCH] feat: Add `Grpc::with_origin` for clients

---
 examples/src/tls/client_rustls.rs | 26 +++-----------------------
 tonic-build/src/client.rs         |  6 ++++++
 tonic/src/client/grpc.rs          | 23 +++++++++++++++++++++++
 3 files changed, 32 insertions(+), 23 deletions(-)

diff --git a/examples/src/tls/client_rustls.rs b/examples/src/tls/client_rustls.rs
index f5a7e4d09..ba0e0e4ee 100644
--- a/examples/src/tls/client_rustls.rs
+++ b/examples/src/tls/client_rustls.rs
@@ -48,30 +48,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let client = hyper::Client::builder().build(connector);
 
-    // Hyper expects an absolute `Uri` to allow it to know which server to connect too.
-    // Currently, tonic's generated code only sets the `path_and_query` section so we
-    // are going to write a custom tower layer in front of the hyper client to add the
-    // scheme and authority.
-    //
-    // Again, this Uri is `example.com` because our tls certs is signed with this SNI but above
-    // we actually map this back to `[::1]:50051` before the `Uri` is passed to hyper's `HttpConnector`
-    // to allow it to correctly establish the tcp connection to the local `tls-server`.
+    // Using `with_origin` will let the codegenerated client set the `scheme` and
+    // `authority` from the porvided `Uri`.
     let uri = Uri::from_static("https://example.com");
-    let svc = tower::ServiceBuilder::new()
-        .map_request(move |mut req: http::Request<tonic::body::BoxBody>| {
-            let uri = Uri::builder()
-                .scheme(uri.scheme().unwrap().clone())
-                .authority(uri.authority().unwrap().clone())
-                .path_and_query(req.uri().path_and_query().unwrap().clone())
-                .build()
-                .unwrap();
-
-            *req.uri_mut() = uri;
-            req
-        })
-        .service(client);
-
-    let mut client = EchoClient::new(svc);
+    let mut client = EchoClient::with_origin(client, uri);
 
     let request = tonic::Request::new(EchoRequest {
         message: "hello".into(),
diff --git a/tonic-build/src/client.rs b/tonic-build/src/client.rs
index fa89724eb..48d455154 100644
--- a/tonic-build/src/client.rs
+++ b/tonic-build/src/client.rs
@@ -44,6 +44,7 @@ pub fn generate<T: Service>(
                 clippy::let_unit_value,
             )]
             use tonic::codegen::*;
+            use tonic::codegen::http::Uri;
 
             #service_doc
             #(#struct_attributes)*
@@ -66,6 +67,11 @@ pub fn generate<T: Service>(
                     Self { inner }
                 }
 
+                pub fn with_origin(inner: T, origin: Uri) -> Self {
+                    let inner = tonic::client::Grpc::with_origin(inner, origin);
+                    Self { inner }
+                }
+
                 pub fn with_interceptor<F>(inner: T, interceptor: F) -> #service_ident<InterceptedService<T, F>>
                 where
                     F: tonic::service::Interceptor,
diff --git a/tonic/src/client/grpc.rs b/tonic/src/client/grpc.rs
index 9ad0bf1da..a1a3675ef 100644
--- a/tonic/src/client/grpc.rs
+++ b/tonic/src/client/grpc.rs
@@ -30,6 +30,7 @@ use std::fmt;
 /// [gRPC protocol definition]: https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests
 pub struct Grpc<T> {
     inner: T,
+    origin: Uri,
     /// Which compression encodings does the client accept?
     accept_compression_encodings: EnabledCompressionEncodings,
     /// The compression encoding that will be applied to requests.
@@ -41,6 +42,20 @@ impl<T> Grpc<T> {
     pub fn new(inner: T) -> Self {
         Self {
             inner,
+            origin: Uri::default(),
+            send_compression_encodings: None,
+            accept_compression_encodings: EnabledCompressionEncodings::default(),
+        }
+    }
+
+    /// Creates a new gRPC client with the provided [`GrpcService`] and `Uri`.
+    ///
+    /// The provided Uri will use only the scheme and authority parts as the
+    /// path_and_query portion will be set for each method.
+    pub fn with_origin(inner: T, origin: Uri) -> Self {
+        Self {
+            inner,
+            origin,
             send_compression_encodings: None,
             accept_compression_encodings: EnabledCompressionEncodings::default(),
         }
@@ -211,8 +226,13 @@ impl<T> Grpc<T> {
         M1: Send + Sync + 'static,
         M2: Send + Sync + 'static,
     {
+        let scheme = self.origin.scheme().cloned();
+        let authority = self.origin.authority().cloned();
+
         let mut parts = Parts::default();
         parts.path_and_query = Some(path);
+        parts.scheme = scheme;
+        parts.authority = authority;
 
         let uri = Uri::from_parts(parts).expect("path_and_query only is valid Uri");
 
@@ -296,6 +316,7 @@ impl<T: Clone> Clone for Grpc<T> {
     fn clone(&self) -> Self {
         Self {
             inner: self.inner.clone(),
+            origin: self.origin.clone(),
             send_compression_encodings: self.send_compression_encodings,
             accept_compression_encodings: self.accept_compression_encodings,
         }
@@ -308,6 +329,8 @@ impl<T: fmt::Debug> fmt::Debug for Grpc<T> {
 
         f.field("inner", &self.inner);
 
+        f.field("origin", &self.origin);
+
         f.field("compression_encoding", &self.send_compression_encodings);
 
         f.field(