From 25d39c13fc8937e86373821093d063202484ac00 Mon Sep 17 00:00:00 2001
From: ByteBaker <42913098+ByteBaker@users.noreply.github.com>
Date: Tue, 20 Aug 2024 00:43:09 +0530
Subject: [PATCH] feat: further TLS options on ClientOptions: #5034 (#6148)

* feat: further TLS options on ClientOptions: #5034

* Rename to Certificate and with_root_certificate, add docs

---------

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
---
 object_store/src/client/mod.rs | 64 ++++++++++++++++++++++++++++++++++
 object_store/src/lib.rs        |  4 +--
 2 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/object_store/src/client/mod.rs b/object_store/src/client/mod.rs
index 43fd65892c2c..c45833b89d2b 100644
--- a/object_store/src/client/mod.rs
+++ b/object_store/src/client/mod.rs
@@ -167,10 +167,60 @@ impl FromStr for ClientConfigKey {
     }
 }
 
+/// Represents a CA certificate provided by the user.
+///
+/// This is used to configure the client to trust a specific certificate. See
+/// [Self::from_pem] for an example
+#[derive(Debug, Clone)]
+pub struct Certificate(reqwest::tls::Certificate);
+
+impl Certificate {
+    /// Create a `Certificate` from a PEM encoded certificate.
+    ///
+    /// # Example from a PEM file
+    ///
+    /// ```no_run
+    /// # use object_store::Certificate;
+    /// # use std::fs::File;
+    /// # use std::io::Read;
+    /// let mut buf = Vec::new();
+    /// File::open("my_cert.pem").unwrap()
+    ///   .read_to_end(&mut buf).unwrap();
+    /// let cert = Certificate::from_pem(&buf).unwrap();
+    ///
+    /// ```
+    pub fn from_pem(pem: &[u8]) -> Result<Self> {
+        Ok(Self(
+            reqwest::tls::Certificate::from_pem(pem).map_err(map_client_error)?,
+        ))
+    }
+
+    /// Create a collection of `Certificate` from a PEM encoded certificate
+    /// bundle.
+    ///
+    /// Files that contain such collections have extensions such as `.crt`,
+    /// `.cer` and `.pem` files.
+    pub fn from_pem_bundle(pem_bundle: &[u8]) -> Result<Vec<Self>> {
+        Ok(reqwest::tls::Certificate::from_pem_bundle(pem_bundle)
+            .map_err(map_client_error)?
+            .into_iter()
+            .map(Self)
+            .collect())
+    }
+
+    /// Create a `Certificate` from a binary DER encoded certificate.
+    pub fn from_der(der: &[u8]) -> Result<Self> {
+        Ok(Self(
+            reqwest::tls::Certificate::from_der(der).map_err(map_client_error)?,
+        ))
+    }
+}
+
 /// HTTP client configuration for remote object stores
 #[derive(Debug, Clone)]
 pub struct ClientOptions {
     user_agent: Option<ConfigValue<HeaderValue>>,
+    root_certificates: Vec<Certificate>,
     content_type_map: HashMap<String, String>,
     default_content_type: Option<String>,
     default_headers: Option<HeaderMap>,
@@ -201,6 +251,7 @@ impl Default for ClientOptions {
         // we opt for a slightly higher default timeout of 30 seconds
         Self {
             user_agent: None,
+            root_certificates: Default::default(),
             content_type_map: Default::default(),
             default_content_type: None,
             default_headers: None,
@@ -310,6 +361,15 @@ impl ClientOptions {
         self
     }
 
+    /// Add a custom root certificate.
+    ///
+    /// This can be used to connect to a server that has a self-signed
+    /// certificate for example.
+    pub fn with_root_certificate(mut self, certificate: Certificate) -> Self {
+        self.root_certificates.push(certificate);
+        self
+    }
+
     /// Set the default CONTENT_TYPE for uploads
     pub fn with_default_content_type(mut self, mime: impl Into<String>) -> Self {
         self.default_content_type = Some(mime.into());
@@ -541,6 +601,10 @@ impl ClientOptions {
             builder = builder.proxy(proxy);
         }
 
+        for certificate in &self.root_certificates {
+            builder = builder.add_root_certificate(certificate.0.clone());
+        }
+
         if let Some(timeout) = &self.timeout {
             builder = builder.timeout(timeout.get()?)
         }
diff --git a/object_store/src/lib.rs b/object_store/src/lib.rs
index 4184d58a0a33..4b43f0cdebae 100644
--- a/object_store/src/lib.rs
+++ b/object_store/src/lib.rs
@@ -526,8 +526,8 @@ mod client;
 
 #[cfg(feature = "cloud")]
 pub use client::{
-    backoff::BackoffConfig, retry::RetryConfig, ClientConfigKey, ClientOptions, CredentialProvider,
-    StaticCredentialProvider,
+    backoff::BackoffConfig, retry::RetryConfig, Certificate, ClientConfigKey, ClientOptions,
+    CredentialProvider, StaticCredentialProvider,
 };
 
 #[cfg(feature = "cloud")]