From 8688bf9c072be9514aa260081f73e595587f2f90 Mon Sep 17 00:00:00 2001
From: Mikhail Katychev <mkatych@gmail.com>
Date: Mon, 19 Aug 2024 14:44:24 -0500
Subject: [PATCH] Simplified `Endpoint::new` initialization

---
 tonic/src/transport/channel/endpoint.rs | 18 +++---------------
 tonic/src/transport/channel/tls.rs      | 10 ++++++++++
 2 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/tonic/src/transport/channel/endpoint.rs b/tonic/src/transport/channel/endpoint.rs
index 488160a01..7306a65de 100644
--- a/tonic/src/transport/channel/endpoint.rs
+++ b/tonic/src/transport/channel/endpoint.rs
@@ -6,6 +6,7 @@ use super::Channel;
 use super::ClientTlsConfig;
 use crate::transport::Error;
 use bytes::Bytes;
+use http::uri::Scheme;
 use http::{uri::Uri, HeaderValue};
 use hyper::rt;
 use hyper_util::client::legacy::connect::HttpConnector;
@@ -50,21 +51,8 @@ impl Endpoint {
     {
         let me = dst.try_into().map_err(|e| Error::from_source(e.into()))?;
         #[cfg(feature = "tls")]
-        if let Some(tls_config) = me
-            .uri
-            .scheme()
-            .map(|s| s.as_str() == http::uri::Scheme::HTTPS.as_str())
-            .unwrap_or(false)
-            .then(|| {
-                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
-            })
-        {
-            return me.tls_config(tls_config);
+        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,