[bitreq] Bitreq client builder

bitreq/src/client.rs

@@ -39,6 +39,50 @@ struct ClientImpl<T> {
     connections: HashMap<ConnectionKey, Arc<T>>,
     lru_order: VecDeque<ConnectionKey>,
     capacity: usize,
+    client_config: Option<ClientConfig>,
+}
+
+pub struct ClientBuilder {
+    capacity: usize,
+    client_config: Option<ClientConfig>,
+}
+
+#[derive(Clone)]
+pub struct ClientConfig {
+    pub tls: Option<TlsConfig>,
+}
+
+#[derive(Clone)]
+pub struct TlsConfig {
+    pub extra_root_cert: Vec<u8>, // DER-encoded certs
+}
+
+impl ClientBuilder {
+    pub fn new() -> Self {
+        Self { capacity: 1, client_config: None }
+    }
+
+    pub fn with_root_certificate<T: Into<Vec<u8>>>(mut self, cert_der: T) -> Self {
+        let tls_config = TlsConfig { extra_root_cert: cert_der.into() };
+        self.client_config = Some(ClientConfig { tls: Some(tls_config) });
+        self
+    }
+
+    pub fn with_capacity(mut self, capacity: usize) -> Self {
+        self.capacity = capacity;
+        self
+    }
+
+    pub fn build(self) -> Client {
+        Client {
+            r#async: Arc::new(Mutex::new(ClientImpl {
+                connections: HashMap::new(),
+                lru_order: VecDeque::new(),
+                capacity: self.capacity,
+                client_config: self.client_config,
+            })),
+        }
+    }
 }
 
 impl Client {
@@ -54,10 +98,16 @@ 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 +127,16 @@ 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.clone()
+            };
+
+            let connection = if let Some(client_config) = client_config {
+                AsyncConnection::new_with_configs(key, parsed_request.timeout_at, client_config).await?
+            } else {
+                AsyncConnection::new(key, parsed_request.timeout_at).await?
+            };
             let connection = Arc::new(connection);
 
             let mut state = self.r#async.lock().unwrap();

bitreq/src/connection.rs

@@ -25,6 +25,8 @@ use tokio::sync::Mutex as AsyncMutex;
 use crate::request::{ConnectionParams, OwnedConnectionParams, ParsedRequest};
 #[cfg(feature = "async")]
 use crate::Response;
+#[cfg(feature = "async")]
+use crate::client::ClientConfig;
 use crate::{Error, Method, ResponseLazy};
 
 type UnsecuredStream = TcpStream;
@@ -298,6 +300,42 @@ impl AsyncConnection {
         }))))
     }
 
+    pub(crate) async fn new_with_configs(
+        params: ConnectionParams<'_>,
+        timeout_at: Option<Instant>,
+        client_config: ClientConfig,
+    ) -> Result<AsyncConnection, Error> {
+        let future = async move {
+            let socket = Self::connect(params).await?;
+
+            if params.https {
+                #[cfg(not(feature = "tokio-rustls"))]
+                return Err(Error::HttpsFeatureNotEnabled);
+                #[cfg(feature = "tokio-rustls")]
+                rustls_stream::wrap_async_stream_with_configs(socket, params.host, client_config).await
+            } else {
+                Ok(AsyncHttpStream::Unsecured(socket))
+            }
+        };
+        let stream = if let Some(timeout_at) = timeout_at {
+            tokio::time::timeout_at(timeout_at.into(), future)
+                .await
+                .unwrap_or(Err(Error::IoError(timeout_err())))?
+        } else {
+            future.await?
+        };
+        let (read, write) = tokio::io::split(stream);
+
+        Ok(AsyncConnection(Mutex::new(Arc::new(AsyncConnectionState {
+            read: AsyncMutex::new(read),
+            write: AsyncMutex::new(write),
+            next_request_id: AtomicUsize::new(0),
+            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)),
+        }))))
+    }
+
     async fn tcp_connect(host: &str, port: u16) -> Result<AsyncTcpStream, Error> {
         #[cfg(feature = "log")]
         log::trace!("Looking up host {host}");

bitreq/src/connection/rustls_stream.rs

@@ -9,6 +9,8 @@ use std::io;
 use std::net::TcpStream;
 use std::sync::OnceLock;
 
+#[cfg(feature = "rustls")]
+use crate::client::ClientConfig as CustomClientConfig;
 #[cfg(all(feature = "native-tls", not(feature = "rustls")))]
 use native_tls::{HandshakeError, TlsConnector, TlsStream};
 #[cfg(feature = "rustls")]
@@ -63,6 +65,54 @@ fn build_client_config() -> Arc<ClientConfig> {
     Arc::new(config)
 }
 
+#[cfg(feature = "rustls")]
+fn build_root_certificates() -> RootCertStore {
+    println!("*** build_client_config CALLED ***");
+
+    let mut root_certificates = RootCertStore::empty();
+
+    // 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));
+        }
+    }
+
+    #[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| {
+            rustls::OwnedTrustAnchor::from_subject_spki_name_constraints(
+                ta.subject,
+                ta.spki,
+                ta.name_constraints,
+            )
+        }));
+    }
+    root_certificates
+}
+
+#[cfg(feature = "rustls")]
+fn append_certificate(mut certificates: RootCertStore, certificate: Vec<u8>) -> RootCertStore {
+    match certificates.add(&rustls::Certificate(certificate)) {
+        Ok(_) => println!("Certificate added successfully"),
+        Err(e) => println!("Failed to add certificate: {:?}", e),
+    }
+    certificates
+}
+
+#[cfg(feature = "rustls")]
+fn build_rustls_client_config(certificates: RootCertStore) -> Arc<ClientConfig> {
+    let config = ClientConfig::builder()
+        .with_safe_defaults()
+        .with_root_certificates(certificates)
+        .with_no_client_auth();
+    Arc::new(config)
+}
+
 #[cfg(feature = "rustls")]
 pub(super) fn wrap_stream(tcp: TcpStream, host: &str) -> Result<SecuredStream, Error> {
     #[cfg(feature = "log")]
@@ -106,6 +156,33 @@ pub(super) async fn wrap_async_stream(
     Ok(AsyncHttpStream::Secured(Box::new(tls)))
 }
 
+#[cfg(all(feature = "rustls", feature = "tokio-rustls"))]
+pub(super) async fn wrap_async_stream_with_configs(
+    tcp: AsyncTcpStream,
+    host: &str,
+    custom_client_config: CustomClientConfig,
+) -> Result<AsyncHttpStream, Error> {
+    #[cfg(feature = "log")]
+    log::trace!("Setting up TLS parameters for {host}.");
+    let dns_name = match ServerName::try_from(host) {
+        Ok(result) => result,
+        Err(err) => return Err(Error::IoError(io::Error::new(io::ErrorKind::Other, err))),
+    };
+
+    let certificates = build_root_certificates();
+    let custom_certificate = custom_client_config.tls.unwrap().extra_root_cert;
+    let certificates = append_certificate(certificates, custom_certificate);
+    let client_config = build_rustls_client_config(certificates);
+    let connector = TlsConnector::from(CONFIG.get_or_init(|| client_config).clone());
+
+    #[cfg(feature = "log")]
+    log::trace!("Establishing TLS session to {host}.");
+
+    let tls = connector.connect(dns_name, tcp).await.map_err(Error::IoError)?;
+
+    Ok(AsyncHttpStream::Secured(Box::new(tls)))
+}
+
 #[cfg(all(feature = "native-tls", not(feature = "rustls")))]
 pub type SecuredStream = TlsStream<TcpStream>;
 

bitreq/tests/main.rs

@@ -16,6 +16,37 @@ async fn test_https() {
     assert_eq!(get_status_code(bitreq::get("https://example.com")).await, 200);
 }
 
+#[tokio::test]
+#[cfg(feature = "rustls")]
+async fn test_https_with_client() {
+    setup();
+    let client = bitreq::Client::new(1);
+    let response = client.send_async(bitreq::get("https://example.com")).await.unwrap();
+    assert_eq!(response.status_code, 200);
+}
+
+#[tokio::test]
+#[cfg(feature = "rustls")]
+async fn test_https_with_client_builder() {
+    setup();
+    let client = bitreq::Client::builder()
+    .build();
+    let response = client.send_async(bitreq::get("https://example.com")).await.unwrap();
+    assert_eq!(response.status_code, 200);
+}
+
+#[tokio::test]
+#[cfg(feature = "rustls")]
+async fn test_https_with_client_builder_and_cert() {
+    setup();
+    let cert_der = include_bytes!("test_cert.der");
+    let client = bitreq::Client::builder()
+    .with_root_certificate(cert_der.as_slice())
+    .build();
+    let response = client.send_async(bitreq::get("https://example.com")).await.unwrap();
+    assert_eq!(response.status_code, 200);
+}
+
 #[tokio::test]
 #[cfg(feature = "json-using-serde")]
 async fn test_json_using_serde() {