From 931691ffcf73213e1661564bbcdbe15b8b977dc1 Mon Sep 17 00:00:00 2001
From: moiseev-signal <122060238+moiseev-signal@users.noreply.github.com>
Date: Wed, 20 Dec 2023 15:16:19 -0800
Subject: [PATCH] Allow specifying certificates when creating
 CdsiEndpointConnection

---
 rust/net/src/env.rs   | 16 +++++++++++++++-
 rust/net/src/infra.rs | 26 +++++++++++++++-----------
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/rust/net/src/env.rs b/rust/net/src/env.rs
index a49e2eba..3000e953 100644
--- a/rust/net/src/env.rs
+++ b/rust/net/src/env.rs
@@ -62,10 +62,24 @@ impl<T: TransportConnector> CdsiEndpointConnection<SingleRouteThrottlingConnecti
         cdsi: CdsiEndpoint<'static>,
         connect_timeout: Duration,
         transport_connector: T,
+    ) -> Self {
+        Self::with_certs(
+            cdsi,
+            connect_timeout,
+            transport_connector,
+            RootCertificates::Signal,
+        )
+    }
+
+    pub fn with_certs(
+        cdsi: CdsiEndpoint<'static>,
+        connect_timeout: Duration,
+        transport_connector: T,
+        certs: RootCertificates,
     ) -> Self {
         Self {
             connection_manager: SingleRouteThrottlingConnectionManager::new(
-                cdsi.direct_connection(),
+                cdsi.direct_connection().with_certs(certs),
                 connect_timeout,
             ),
             connector: WebSocketClientConnector::new(
diff --git a/rust/net/src/infra.rs b/rust/net/src/infra.rs
index faf4decf..a1451490 100644
--- a/rust/net/src/infra.rs
+++ b/rust/net/src/infra.rs
@@ -96,6 +96,11 @@ impl ConnectionParams {
         decorators.push(decorator);
         self
     }
+
+    pub fn with_certs(mut self, certs: RootCertificates) -> Self {
+        self.certs = certs;
+        self
+    }
 }
 
 impl HttpRequestDecoratorSeq {
@@ -167,7 +172,7 @@ impl TransportConnector for TcpSslTransportConnector {
         )
         .await?;
 
-        let ssl_config = client_ssl_connector_builder(connection_params.certs.clone(), alpn)?
+        let ssl_config = Self::builder(connection_params.certs.clone(), alpn)?
             .build()
             .configure()?;
 
@@ -179,6 +184,15 @@ impl TransportConnector for TcpSslTransportConnector {
     }
 }
 
+impl TcpSslTransportConnector {
+    fn builder(certs: RootCertificates, alpn: &[u8]) -> Result<SslConnectorBuilder, NetError> {
+        let mut ssl = SslConnector::builder(SslMethod::tls_client())?;
+        ssl.set_verify_cert_store(certs.try_into()?)?;
+        ssl.set_alpn_protos(alpn)?;
+        Ok(ssl)
+    }
+}
+
 pub(crate) async fn connect_tcp(
     dns_resolver: &DnsResolver,
     host: &str,
@@ -197,16 +211,6 @@ pub(crate) async fn connect_tcp(
     Err(NetError::TcpConnectionFailed)
 }
 
-pub(crate) fn client_ssl_connector_builder(
-    certs: RootCertificates,
-    alpn: &[u8],
-) -> Result<SslConnectorBuilder, NetError> {
-    let mut ssl = SslConnector::builder(SslMethod::tls_client())?;
-    ssl.set_verify_cert_store(certs.try_into()?)?;
-    ssl.set_alpn_protos(alpn)?;
-    Ok(ssl)
-}
-
 #[cfg(test)]
 pub(crate) mod test {
     use hyper::Request;