code review adjustment: streamline code and reduce complexity to avoid arcs/mutexes/clones when building ClientConfigs

Author: ton-anywhere

bitreq/src/client.rs

@@ -13,7 +13,7 @@ use std::sync::{Arc, Mutex};
     all(feature = "native-tls", feature = "tokio-native-tls"),
     all(feature = "rustls", feature = "tokio-rustls")
 ))]
-use crate::connection::certificates::Certificates;
+use crate::connection::certificates::{Certificates, CertificatesBuilder};
 use crate::connection::AsyncConnection;
 use crate::request::{OwnedConnectionParams as ConnectionKey, ParsedRequest};
 use crate::{Error, Request, Response};
@@ -27,23 +27,6 @@ pub(crate) struct ClientConfig {
     pub(crate) tls: Option<TlsConfig>,
 }
 
-impl ClientConfig {
-    #[cfg(any(
-        all(feature = "native-tls", feature = "tokio-native-tls"),
-        all(feature = "rustls", feature = "tokio-rustls")
-    ))]
-    pub fn build(self) -> Result<Self, Error> {
-        let tls = self.tls.map(|tls| tls.build()).transpose()?;
-        Ok(Self { tls })
-    }
-
-    #[cfg(not(any(
-        all(feature = "native-tls", feature = "tokio-native-tls"),
-        all(feature = "rustls", feature = "tokio-rustls")
-    )))]
-    pub fn build(self) -> Result<Self, Error> { Ok(Self {}) }
-}
-
 #[cfg(any(
     all(feature = "native-tls", feature = "tokio-native-tls"),
     all(feature = "rustls", feature = "tokio-rustls")
@@ -58,34 +41,16 @@ pub(crate) struct TlsConfig {
     all(feature = "rustls", feature = "tokio-rustls")
 ))]
 impl TlsConfig {
-    fn new(cert_der: Vec<u8>) -> Result<Self, Error> {
-        let certificates = Certificates::new(Some(cert_der))?;
-
-        Ok(Self { certificates })
-    }
-
-    #[cfg(all(feature = "rustls", feature = "tokio-rustls"))]
-    fn build(mut self) -> Result<Self, Error> {
-        if self.certificates.disable_default {
-            return Ok(self);
-        }
-
-        self.certificates = self.certificates.with_root_certificates();
-        Ok(self)
-    }
-
-    #[cfg(all(feature = "native-tls", not(feature = "rustls"), feature = "tokio-native-tls"))]
-    fn build(mut self) -> Result<Self, Error> {
-        let certificates = self.certificates.build()?;
-
-        self.certificates = certificates;
-        Ok(self)
-    }
+    fn new(certificates: Certificates) -> Self { Self { certificates } }
 }
 
 pub struct ClientBuilder {
     capacity: usize,
-    client_config: Option<ClientConfig>,
+    #[cfg(any(
+        all(feature = "native-tls", feature = "tokio-native-tls"),
+        all(feature = "rustls", feature = "tokio-rustls")
+    ))]
+    certificates: Option<CertificatesBuilder>,
 }
 
 /// Builder for configuring a `Client` with custom settings.
@@ -106,22 +71,39 @@ pub struct ClientBuilder {
 /// ```
 impl ClientBuilder {
     /// Creates a new `ClientBuilder` with a default pool capacity of 10.
-    pub fn new() -> Self { Self { capacity: 10, client_config: None } }
+    #[cfg(any(
+        all(feature = "native-tls", feature = "tokio-native-tls"),
+        all(feature = "rustls", feature = "tokio-rustls")
+    ))]
+    pub fn new() -> Self { Self { capacity: 10, certificates: None } }
+
+    /// Creates a new `ClientBuilder` with a default pool capacity of 10.
+    #[cfg(not(any(
+        all(feature = "native-tls", feature = "tokio-native-tls"),
+        all(feature = "rustls", feature = "tokio-rustls")
+    )))]
+    pub fn new() -> Self { Self { capacity: 10 } }
 
     /// Sets the maximum number of connections to keep in the pool.
     pub fn with_capacity(mut self, capacity: usize) -> Self {
         self.capacity = capacity;
         self
     }
 
+    #[cfg(any(
+        all(feature = "native-tls", feature = "tokio-native-tls"),
+        all(feature = "rustls", feature = "tokio-rustls")
+    ))]
     /// Builds the `Client` with the configured settings.
     pub fn build(self) -> Result<Client, Error> {
-        let build = self.client_config.map(|c| c.build());
-        let client_config = match build {
-            Some(Ok(config)) => Some(Arc::new(config)),
-            Some(Err(e)) => return Err(e),
-            None => None,
+        let build_config = if let Some(builder) = self.certificates {
+            let certificates = builder.build()?;
+            let tls_config = TlsConfig::new(certificates);
+            Some(ClientConfig { tls: Some(tls_config) })
+        } else {
+            None
         };
+        let client_config = build_config.map(Arc::new);
 
         Ok(Client {
             r#async: Arc::new(Mutex::new(ClientImpl {
@@ -132,6 +114,23 @@ impl ClientBuilder {
             })),
         })
     }
+
+    /// Builds the `Client` with the configured settings.
+    #[cfg(not(any(
+        all(feature = "native-tls", feature = "tokio-native-tls"),
+        all(feature = "rustls", feature = "tokio-rustls")
+    )))]
+    pub fn build(self) -> Result<Client, Error> {
+        Ok(Client {
+            r#async: Arc::new(Mutex::new(ClientImpl {
+                connections: HashMap::new(),
+                lru_order: VecDeque::new(),
+                capacity: self.capacity,
+                client_config: None,
+            })),
+        })
+    }
+
     /// Adds a custom DER-encoded root certificate for TLS verification.
     /// The certificate must be provided in DER format. This method accepts any type
     /// that can be converted into a `Vec<u8>`.
@@ -156,18 +155,13 @@ impl ClientBuilder {
     ))]
     pub fn with_root_certificate<T: Into<Vec<u8>>>(mut self, cert_der: T) -> Result<Self, Error> {
         let cert_der = cert_der.into();
+        if let Some(ref mut certificates) = self.certificates {
+            certificates.append_certificate(cert_der)?;
 
-        if let Some(ref mut client_config) = self.client_config {
-            if let Some(ref mut tls_config) = client_config.tls {
-                let certificates = tls_config.certificates.clone().append_certificate(cert_der)?;
-                tls_config.certificates = certificates;
-
-                return Ok(self);
-            }
+            return Ok(self);
         }
 
-        let tls_config = TlsConfig::new(cert_der)?;
-        self.client_config = Some(ClientConfig { tls: Some(tls_config) });
+        self.certificates = Some(CertificatesBuilder::new(Some(cert_der))?);
         Ok(self)
     }
 
@@ -178,9 +172,11 @@ impl ClientBuilder {
         all(feature = "rustls", feature = "tokio-rustls")
     ))]
     pub fn disable_default_certificates(mut self) -> Result<Self, Error> {
-        let client_config = self.client_config.as_mut().ok_or(Error::InvalidTlsConfig)?;
-        let tls_config = client_config.tls.as_mut().ok_or(Error::InvalidTlsConfig)?;
-        tls_config.certificates = tls_config.certificates.clone().disable_default()?;
+        match self.certificates {
+            Some(ref mut certificates) => certificates.disable_default()?,
+            None => return Err(Error::InvalidTlsConfig),
+        };
+
         Ok(self)
     }
 }

bitreq/src/connection.rs

@@ -315,7 +315,9 @@ impl AsyncConnection {
         client_config: Option<Arc<ClientConfig>>,
     ) -> Result<AsyncHttpStream, Error> {
         if let Some(client_config) = client_config {
-            rustls_stream::wrap_async_stream_with_configs(socket, host, client_config).await
+            let tls_config = client_config.tls.as_ref().unwrap();
+            let certificates = tls_config.certificates.clone();
+            rustls_stream::wrap_async_stream_with_configs(socket, host, certificates).await
         } else {
             rustls_stream::wrap_async_stream(socket, host).await
         }

bitreq/src/connection/certificates.rs

@@ -1,125 +1,101 @@
-#[cfg(feature = "rustls")]
+#[cfg(any(feature = "rustls", feature = "native-tls"))]
 use std::sync::Arc;
-#[cfg(all(feature = "native-tls", not(feature = "rustls"), feature = "tokio-native-tls"))]
-use std::sync::{Arc, Mutex};
 
 #[cfg(all(feature = "native-tls", not(feature = "rustls")))]
 use native_tls::{Certificate, TlsConnector, TlsConnectorBuilder};
 #[cfg(feature = "rustls")]
 use rustls::RootCertStore;
+#[cfg(all(feature = "native-tls", not(feature = "rustls"), feature = "tokio-native-tls"))]
+use tokio_native_tls::TlsConnector as AsyncTlsConnector;
 #[cfg(feature = "rustls-webpki")]
 use webpki_roots::TLS_SERVER_ROOTS;
 
 use crate::Error;
 
-#[derive(Clone)]
 #[cfg(feature = "rustls")]
-pub(crate) struct Certificates {
-    pub(crate) inner: Arc<RootCertStore>,
+pub(crate) struct CertificatesBuilder {
+    pub(crate) inner: RootCertStore,
     pub(crate) disable_default: bool,
 }
 
-#[derive(Clone)]
 #[cfg(all(feature = "native-tls", not(feature = "rustls"), feature = "tokio-native-tls"))]
-pub(crate) struct Certificates {
-    pub(crate) inner: CertificatesInner,
+pub(crate) struct CertificatesBuilder {
+    pub(crate) inner: TlsConnectorBuilder,
 }
 
-#[derive(Clone)]
-#[cfg(all(feature = "native-tls", not(feature = "rustls"), feature = "tokio-native-tls"))]
-pub(crate) enum CertificatesInner {
-    Builder(Arc<Mutex<TlsConnectorBuilder>>),
-    Built(TlsConnector),
-}
-
-impl Certificates {
+impl CertificatesBuilder {
     #[cfg(feature = "rustls")]
     pub(crate) fn new(cert_der: Option<Vec<u8>>) -> Result<Self, Error> {
-        let certificates = Self { inner: Arc::new(RootCertStore::empty()), disable_default: false };
+        let mut certificates = Self { inner: RootCertStore::empty(), disable_default: false };
 
         if let Some(cert_der) = cert_der {
-            certificates.append_certificate(cert_der)
-        } else {
-            Ok(certificates)
+            certificates.append_certificate(cert_der)?;
         }
+
+        Ok(certificates)
     }
 
     #[cfg(all(feature = "native-tls", not(feature = "rustls"), feature = "tokio-native-tls"))]
     pub(crate) fn new(cert_der: Option<Vec<u8>>) -> Result<Self, Error> {
         let builder = TlsConnector::builder();
-        let inner = CertificatesInner::Builder(Arc::new(Mutex::new(builder)));
-        let certificates = Self { inner: inner };
+        let mut certificates = Self { inner: builder };
 
         if let Some(cert_der) = cert_der {
-            certificates.append_certificate(cert_der)
-        } else {
-            Ok(certificates)
+            certificates.append_certificate(cert_der)?;
         }
+
+        Ok(certificates)
     }
 
     #[cfg(feature = "rustls")]
-    pub(crate) fn append_certificate(mut self, cert_der: Vec<u8>) -> Result<Self, Error> {
-        let certificates = Arc::make_mut(&mut self.inner);
-        certificates.add(&rustls::Certificate(cert_der)).map_err(Error::RustlsAppendCert)?;
+    pub(crate) fn append_certificate(&mut self, cert_der: Vec<u8>) -> Result<&mut Self, Error> {
+        self.inner.add(&rustls::Certificate(cert_der)).map_err(Error::RustlsAppendCert)?;
 
         Ok(self)
     }
 
     #[cfg(all(feature = "native-tls", not(feature = "rustls"), feature = "tokio-native-tls"))]
-    pub(crate) fn append_certificate(mut self, cert_der: Vec<u8>) -> Result<Self, Error> {
-        let new_inner = match self.inner {
-            CertificatesInner::Builder(builder_mutex) => {
-                let certificate = Certificate::from_der(&cert_der)?;
-
-                {
-                    let mut builder_guard = builder_mutex.lock().unwrap();
-                    builder_guard.add_root_certificate(certificate);
-                }
+    pub(crate) fn append_certificate(&mut self, cert_der: Vec<u8>) -> Result<&mut Self, Error> {
+        let certificate = Certificate::from_der(&cert_der)?;
+        self.inner.add_root_certificate(certificate);
 
-                CertificatesInner::Builder(builder_mutex)
-            }
-            CertificatesInner::Built(_) => return Err(Error::NativeTlsAppendCert),
-        };
-
-        self.inner = new_inner;
         Ok(self)
     }
 
     #[cfg(all(feature = "native-tls", not(feature = "rustls"), feature = "tokio-native-tls"))]
-    pub(crate) fn build(mut self) -> Result<Self, Error> {
-        let new_inner = match self.inner {
-            CertificatesInner::Builder(builder_mutex) => {
-                let mut builder_guard = builder_mutex.lock().unwrap();
-                let connector = builder_guard.build()?;
-
-                CertificatesInner::Built(connector)
-            }
-            CertificatesInner::Built(_) => return Ok(self),
-        };
+    pub(crate) fn build(self) -> Result<Certificates, Error> {
+        let connector = self.inner.build()?;
+        let async_connector = AsyncTlsConnector::from(connector);
 
-        self.inner = new_inner;
-        Ok(self)
+        Ok(Certificates(Arc::new(async_connector)))
     }
 
     #[cfg(feature = "rustls")]
-    pub(crate) fn with_root_certificates(mut self) -> Self {
-        let root_certificates = Arc::make_mut(&mut self.inner);
+    pub(crate) fn build(mut self) -> Result<Certificates, Error> {
+        if !self.disable_default {
+            self.with_root_certificates();
+        }
+
+        Ok(Certificates(Arc::new(self.inner)))
+    }
 
+    #[cfg(feature = "rustls")]
+    fn with_root_certificates(&mut self) -> &mut Self {
         // Try to load native certs
         #[cfg(feature = "https-rustls-probe")]
         if let Ok(os_roots) = rustls_native_certs::load_native_certs() {
             for root_cert in os_roots {
                 // Ignore erroneous OS certificates, there's nothing
                 // to do differently in that situation anyways.
-                let _ = root_certificates.add(&rustls::Certificate(root_cert.0));
+                let _ = self.inner.add(&rustls::Certificate(root_cert.0));
             }
         }
 
         #[cfg(feature = "rustls-webpki")]
         {
             #[allow(deprecated)]
             // Need to use add_server_trust_anchors to compile with rustls 0.21.1
-            root_certificates.add_server_trust_anchors(TLS_SERVER_ROOTS.iter().map(|ta| {
+            self.inner.add_server_trust_anchors(TLS_SERVER_ROOTS.iter().map(|ta| {
                 rustls::OwnedTrustAnchor::from_subject_spki_name_constraints(
                     ta.subject,
                     ta.spki,
@@ -131,19 +107,22 @@ impl Certificates {
     }
 
     #[cfg(feature = "rustls")]
-    pub(crate) fn disable_default(mut self) -> Result<Self, Error> {
+    pub(crate) fn disable_default(&mut self) -> Result<&mut Self, Error> {
         self.disable_default = true;
         Ok(self)
     }
 
     #[cfg(all(feature = "native-tls", not(feature = "rustls"), feature = "tokio-native-tls"))]
-    pub(crate) fn disable_default(self) -> Result<Self, Error> {
-        match self.inner {
-            CertificatesInner::Builder(ref builder_mutex) => {
-                builder_mutex.lock().unwrap().disable_built_in_roots(true);
-                Ok(self)
-            }
-            CertificatesInner::Built(_) => return Err(Error::InvalidTlsConfig),
-        }
+    pub(crate) fn disable_default(&mut self) -> Result<&mut Self, Error> {
+        self.inner.disable_built_in_roots(true);
+        Ok(self)
     }
 }
+
+#[derive(Clone)]
+#[cfg(feature = "rustls")]
+pub(crate) struct Certificates(pub(crate) Arc<RootCertStore>);
+
+#[derive(Clone)]
+#[cfg(all(feature = "native-tls", not(feature = "rustls"), feature = "tokio-native-tls"))]
+pub(crate) struct Certificates(pub(crate) Arc<AsyncTlsConnector>);