From 22475d803df668e3d13795fdbd68eed081587f1e Mon Sep 17 00:00:00 2001
From: Mikhail Katychev <mkatych@gmail.com>
Date: Wed, 21 Aug 2024 09:09:12 -0500
Subject: [PATCH] fix(tonic-build,tonic) Add back TLS handling in genereated
 `Client::connect` code (#1866)

* tls feature flag for Endpoint::new

* added unit test

* Simplified `Endpoint::new` initialization
---
 tests/integration_tests/tests/connection.rs |  5 +++++
 tonic/src/transport/channel/endpoint.rs     |  5 +++++
 tonic/src/transport/channel/tls.rs          | 10 ++++++++++
 3 files changed, 20 insertions(+)

diff --git a/tests/integration_tests/tests/connection.rs b/tests/integration_tests/tests/connection.rs
index bb67adf55..841600bcf 100644
--- a/tests/integration_tests/tests/connection.rs
+++ b/tests/integration_tests/tests/connection.rs
@@ -26,6 +26,11 @@ async fn connect_returns_err() {
     assert!(res.is_err());
 }
 
+#[tokio::test]
+async fn connect_handles_tls() {
+    TestClient::connect("https://example.com").await.unwrap();
+}
+
 #[tokio::test]
 async fn connect_returns_err_via_call_after_connected() {
     let (tx, rx) = oneshot::channel();
diff --git a/tonic/src/transport/channel/endpoint.rs b/tonic/src/transport/channel/endpoint.rs
index 508415f63..aaf977346 100644
--- a/tonic/src/transport/channel/endpoint.rs
+++ b/tonic/src/transport/channel/endpoint.rs
@@ -49,6 +49,11 @@ impl Endpoint {
         D::Error: Into<crate::Error>,
     {
         let me = dst.try_into().map_err(|e| Error::from_source(e.into()))?;
+        #[cfg(feature = "tls")]
+        if me.uri.scheme() == Some(&http::uri::Scheme::HTTPS) {
+            return me.tls_config(ClientTlsConfig::new().with_enabled_roots());
+        }
+
         Ok(me)
     }
 
diff --git a/tonic/src/transport/channel/tls.rs b/tonic/src/transport/channel/tls.rs
index 8c845f0ef..0e1ec254c 100644
--- a/tonic/src/transport/channel/tls.rs
+++ b/tonic/src/transport/channel/tls.rs
@@ -81,6 +81,16 @@ impl ClientTlsConfig {
         }
     }
 
+    /// Activates all TLS roots enabled through `tls-*-roots` feature flags
+    pub fn with_enabled_roots(self) -> Self {
+        let config = ClientTlsConfig::new();
+        #[cfg(feature = "tls-native-roots")]
+        let config = config.with_native_roots();
+        #[cfg(feature = "tls-webpki-roots")]
+        let config = config.with_webpki_roots();
+        config
+    }
+
     pub(crate) fn into_tls_connector(self, uri: &Uri) -> Result<TlsConnector, crate::Error> {
         let domain = match &self.domain {
             Some(domain) => domain,