implement client builder for rustls

Author: ton-anywhere

bitreq/src/client.rs

@@ -9,10 +9,138 @@
 use std::collections::{hash_map, HashMap, VecDeque};
 use std::sync::{Arc, Mutex};
 
+#[cfg(feature = "tokio-rustls")]
+use crate::connection::tls_config::{TlsConfig, TlsConfigBuilder};
 use crate::connection::AsyncConnection;
 use crate::request::{OwnedConnectionParams as ConnectionKey, ParsedRequest};
 use crate::{Error, Request, Response};
 
+#[derive(Clone)]
+pub(crate) struct ClientConfig {
+    #[cfg(feature = "tokio-rustls")]
+    pub(crate) tls: Option<TlsConfig>,
+}
+
+pub struct ClientBuilder {
+    capacity: usize,
+    #[cfg(feature = "tokio-rustls")]
+    tls_config: Option<TlsConfigBuilder>,
+}
+
+/// Builder for configuring a `Client` with custom settings.
+///
+/// # Example
+///
+/// ```no_run
+/// # async fn example() -> Result<(), bitreq::Error> {
+/// use bitreq::{Client, RequestExt};
+///
+/// let client = Client::builder().with_capacity(20).build()?;
+///
+/// let response = bitreq::get("https://example.com")
+///     .send_async_with_client(&client)
+///     .await?;
+/// # Ok(())
+/// # }
+/// ```
+impl ClientBuilder {
+    /// Creates a new `ClientBuilder` with a default pool capacity of 10.
+    pub fn new() -> Self {
+        Self {
+            capacity: 10,
+            #[cfg(feature = "tokio-rustls")]
+            tls_config: None,
+        }
+    }
+
+    /// 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(feature = "tokio-rustls")]
+    /// Builds the `Client` with the configured settings.
+    pub fn build(self) -> Result<Client, Error> {
+        let build_config = if let Some(builder) = self.tls_config {
+            let tls_config = builder.build()?;
+            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 {
+                connections: HashMap::new(),
+                lru_order: VecDeque::new(),
+                capacity: self.capacity,
+                client_config,
+            })),
+        })
+    }
+
+    /// Builds the `Client` with the configured settings.
+    #[cfg(not(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>`.
+    /// The certificate is appended to the default trust store rather than replacing it.
+    /// The trust store used depends on the TLS backend: system certificates for native-tls,
+    /// Mozilla's root certificates(rustls-webpki) and/or system certificates(rustls-native-certs) for rustls.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// # use bitreq::Client;
+    /// # async fn example() -> Result<(), bitreq::Error> {
+    /// let client = Client::builder()
+    ///     .with_root_certificate(include_bytes!("../tests/test_cert.der"))?
+    ///     .build()?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    #[cfg(feature = "tokio-rustls")]
+    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 tls_config) = self.tls_config {
+            tls_config.append_certificate(cert_der)?;
+
+            return Ok(self);
+        }
+
+        self.tls_config = Some(TlsConfigBuilder::new(Some(cert_der))?);
+        Ok(self)
+    }
+
+    /// Disables default root certificates for TLS connections.
+    /// Returns [`Error::InvalidTlsConfig`] if TLS has not been configured.
+    #[cfg(feature = "tokio-rustls")]
+    pub fn disable_default_certificates(mut self) -> Result<Self, Error> {
+        match self.tls_config {
+            Some(ref mut tls_config) => tls_config.disable_default_certificates()?,
+            None => return Err(Error::InvalidTlsConfig),
+        };
+
+        Ok(self)
+    }
+}
+
+impl Default for ClientBuilder {
+    fn default() -> Self { Self::new() }
+}
+
 /// A client that caches connections for reuse.
 ///
 /// The client maintains a pool of up to `capacity` connections, evicting
@@ -39,10 +167,11 @@ struct ClientImpl<T> {
     connections: HashMap<ConnectionKey, Arc<T>>,
     lru_order: VecDeque<ConnectionKey>,
     capacity: usize,
+    client_config: Option<Arc<ClientConfig>>,
 }
 
 impl Client {
-    /// Creates a new `Client` with the specified connection cache capacity.
+    /// Creates a new `Client` with the specified connection pool capacity.
     ///
     /// # Arguments
     ///
@@ -54,10 +183,14 @@ impl Client {
                 connections: HashMap::new(),
                 lru_order: VecDeque::new(),
                 capacity,
+                client_config: None,
             })),
         }
     }
 
+    /// Create a builder for a client
+    pub fn builder() -> ClientBuilder { ClientBuilder::new() }
+
     /// Sends a request asynchronously using a cached connection if available.
     pub async fn send_async(&self, request: Request) -> Result<Response, Error> {
         let parsed_request = ParsedRequest::new(request)?;
@@ -77,7 +210,13 @@ impl Client {
         let conn = if let Some(conn) = conn_opt {
             conn
         } else {
-            let connection = AsyncConnection::new(key, parsed_request.timeout_at).await?;
+            let client_config = {
+                let state = self.r#async.lock().unwrap();
+                state.client_config.as_ref().map(Arc::clone)
+            };
+
+            let connection =
+                AsyncConnection::new(key, parsed_request.timeout_at, client_config).await?;
             let connection = Arc::new(connection);
 
             let mut state = self.r#async.lock().unwrap();

bitreq/src/connection.rs

@@ -22,6 +22,8 @@ use tokio::net::TcpStream as AsyncTcpStream;
 #[cfg(feature = "async")]
 use tokio::sync::Mutex as AsyncMutex;
 
+#[cfg(feature = "async")]
+use crate::client::ClientConfig;
 use crate::request::{ConnectionParams, OwnedConnectionParams, ParsedRequest};
 #[cfg(feature = "async")]
 use crate::Response;
@@ -31,6 +33,8 @@ type UnsecuredStream = TcpStream;
 
 #[cfg(any(feature = "rustls", feature = "native-tls"))]
 mod rustls_stream;
+#[cfg(feature = "tokio-rustls")]
+pub(crate) mod tls_config;
 #[cfg(any(feature = "rustls", feature = "native-tls"))]
 type SecuredStream = rustls_stream::SecuredStream;
 
@@ -238,6 +242,7 @@ struct AsyncConnectionState {
     /// Defaults to 60 seconds after open to align with nginx's default timeout of 75 seconds, but
     /// can be overridden by the `Keep-Alive` header.
     socket_new_requests_timeout: Mutex<Instant>,
+    client_config: Option<Arc<ClientConfig>>,
 }
 
 #[cfg(feature = "async")]
@@ -266,13 +271,15 @@ impl AsyncConnection {
     pub(crate) async fn new(
         params: ConnectionParams<'_>,
         timeout_at: Option<Instant>,
+        client_config: Option<Arc<ClientConfig>>,
     ) -> Result<AsyncConnection, Error> {
+        let client_config_ref = &client_config;
+
         let future = async move {
             let socket = Self::connect(params).await?;
 
             if params.https {
-                // temp call
-                Self::wrap_async_stream(socket, params.host).await
+                Self::wrap_async_stream(socket, params.host, client_config_ref).await
             } else {
                 Ok(AsyncHttpStream::Unsecured(socket))
             }
@@ -293,26 +300,47 @@ impl AsyncConnection {
             readable_request_id: AtomicUsize::new(0),
             min_dropped_reader_id: AtomicUsize::new(usize::MAX),
             socket_new_requests_timeout: Mutex::new(Instant::now() + Duration::from_secs(60)),
+            client_config,
         }))))
     }
 
-    /// Temp Method implementation
-    #[cfg(any(feature = "tokio-rustls", feature = "tokio-native-tls"))]
+    // =======
+    /// Temp method. Required to compile
+    #[cfg(all(feature = "tokio-native-tls", not(feature = "tokio-rustls")))]
     async fn wrap_async_stream(
         socket: AsyncTcpStream,
         host: &str,
+        _client_config: &Option<Arc<ClientConfig>>,
     ) -> Result<AsyncHttpStream, Error> {
         rustls_stream::wrap_async_stream(socket, host).await
     }
+    // =======
+
+    /// Call the correct wrapper function depending on whether client_configs are present
+    #[cfg(feature = "tokio-rustls")]
+    async fn wrap_async_stream(
+        socket: AsyncTcpStream,
+        host: &str,
+        client_config: &Option<Arc<ClientConfig>>,
+    ) -> Result<AsyncHttpStream, Error> {
+        if let Some(client_config) = client_config {
+            let tls_config = client_config.tls.as_ref().unwrap().clone();
+            rustls_stream::wrap_async_stream_with_configs(socket, host, tls_config).await
+        } else {
+            rustls_stream::wrap_async_stream(socket, host).await
+        }
+    }
 
-    /// Temp Method implementation
+    /// Error treatment function, should not be called under normal circustances
     #[cfg(not(any(feature = "tokio-rustls", feature = "tokio-native-tls")))]
     async fn wrap_async_stream(
         _socket: AsyncTcpStream,
         _host: &str,
+        _client_config: &Option<Arc<ClientConfig>>,
     ) -> Result<AsyncHttpStream, Error> {
         Err(Error::HttpsFeatureNotEnabled)
     }
+
     async fn tcp_connect(host: &str, port: u16) -> Result<AsyncTcpStream, Error> {
         #[cfg(feature = "log")]
         log::trace!("Looking up host {host}");
@@ -461,9 +489,13 @@ impl AsyncConnection {
                     retry_new_connection!(_internal);
                 };
                 (_internal) => {
-                    let new_connection =
-                        AsyncConnection::new(request.connection_params(), request.timeout_at)
-                            .await?;
+                    let config = conn.client_config.as_ref().map(Arc::clone);
+                    let new_connection = AsyncConnection::new(
+                        request.connection_params(),
+                        request.timeout_at,
+                        config,
+                    )
+                    .await?;
                     *self.0.lock().unwrap() = Arc::clone(&*new_connection.0.lock().unwrap());
                     core::mem::drop(read);
                     // Note that this cannot recurse infinitely as we'll always be able to send at
@@ -818,7 +850,8 @@ async fn async_handle_redirects(
             let new_connection;
             if needs_new_connection {
                 new_connection =
-                    AsyncConnection::new(request.connection_params(), request.timeout_at).await?;
+                    AsyncConnection::new(request.connection_params(), request.timeout_at, None)
+                        .await?;
                 connection = &new_connection;
             }
             connection.send(request).await