channel_provider: tls cfg (#77)

Author: alissa-tung

.cargo/config.toml

@@ -1,4 +1,4 @@
 [env]
 RUST_LOG = "hstreamdb"
 
-TEST_SERVER_ADDR = "http://127.0.0.1:6570"
+TEST_SERVER_ADDR = "hstream://127.0.0.1:6570"

Cargo.lock

@@ -27,7 +27,7 @@ prost-types = "0.11.1"
 prost = "0.11.0"
 tokio = { version = "1.21.0", features = ["rt-multi-thread", "parking_lot"] }
 tokio-stream = "0.1.9"
-tonic = "0.8.0"
+tonic = { version = "0.8.2", features = ["tls"] }
 url = "2.2.2"
 workspace-hack = { version = "0.1", path = "../utils/workspace-hack" }
 

src/hstreamdb/Cargo.toml

@@ -1,10 +1,11 @@
 use std::collections::HashMap;
+use std::default::default;
 use std::iter::FromIterator;
 
 use hstreamdb_pb::h_stream_api_client::HStreamApiClient;
 use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender};
 use tokio::sync::oneshot;
-use tonic::transport::{Channel, Endpoint};
+use tonic::transport::{Channel, ClientTlsConfig, Endpoint};
 
 use crate::client::get_available_node_addrs;
 use crate::common;
@@ -67,8 +68,39 @@ impl Channels {
     }
 }
 
+#[derive(Debug, Default)]
 pub struct ChannelProviderSettings {
-    pub concurrency_limit: Option<usize>,
+    concurrency_limit: Option<usize>,
+    pub(crate) client_tls_config: Option<ClientTlsConfig>,
+}
+
+pub struct ChannelProviderSettingsBuilder(ChannelProviderSettings);
+
+impl ChannelProviderSettings {
+    pub fn builder() -> ChannelProviderSettingsBuilder {
+        ChannelProviderSettingsBuilder(default())
+    }
+}
+
+impl ChannelProviderSettingsBuilder {
+    pub fn build(self) -> ChannelProviderSettings {
+        let ChannelProviderSettingsBuilder(channel_provider_settings) = self;
+        channel_provider_settings
+    }
+
+    pub fn set_concurrency_limit(self, concurrency_limit: usize) -> Self {
+        Self(ChannelProviderSettings {
+            concurrency_limit: Some(concurrency_limit),
+            ..self.0
+        })
+    }
+
+    pub fn set_client_tls_config(self, client_tls_config: ClientTlsConfig) -> Self {
+        Self(ChannelProviderSettings {
+            client_tls_config: Some(client_tls_config),
+            ..self.0
+        })
+    }
 }
 
 impl ChannelProvider {
@@ -91,6 +123,9 @@ impl ChannelProvider {
                     if let Some(concurrency_limit) = settings.concurrency_limit {
                         endpoint = endpoint.concurrency_limit(concurrency_limit)
                     }
+                    if let Some(client_tls_config) = settings.client_tls_config.clone() {
+                        endpoint = endpoint.tls_config(client_tls_config)?
+                    }
                     match endpoint.connect().await {
                         Err(err) => {
                             log::warn!("connect to endpoint error: uri = {uri}, {err}");

src/hstreamdb/src/channel_provider.rs

@@ -26,10 +26,26 @@ impl Client {
     where
         Destination: std::convert::Into<String>,
     {
+        const HSTREAM_PREFIX: &str = "hstream";
         let server_url = server_url.into();
-        let url = Url::parse(&server_url)?;
-        let mut hstream_api_client = HStreamApiClient::connect(server_url).await?;
-        let url_scheme = url.scheme().to_string();
+        let (url_scheme, url) = {
+            let url = Url::parse(&server_url)?;
+            if url.scheme() == HSTREAM_PREFIX {
+                let url_scheme = if channel_provider_settings.client_tls_config.is_none() {
+                    "http"
+                } else {
+                    "https"
+                };
+                let server_url = &server_url[7..];
+                (
+                    url_scheme.to_string(),
+                    Url::parse(format!("{url_scheme}{server_url}").as_str())?,
+                )
+            } else {
+                (url.scheme().to_string(), url)
+            }
+        };
+        let mut hstream_api_client = HStreamApiClient::connect(String::from(url)).await?;
         let channels = new_channel_provider(
             &url_scheme,
             &mut hstream_api_client,
@@ -242,14 +258,9 @@ mod tests {
     #[tokio::test(flavor = "multi_thread")]
     async fn test_stream_cld() {
         let addr = env::var("TEST_SERVER_ADDR").unwrap();
-        let mut client = Client::new(
-            addr,
-            ChannelProviderSettings {
-                concurrency_limit: None,
-            },
-        )
-        .await
-        .unwrap();
+        let mut client = Client::new(addr, ChannelProviderSettings::builder().build())
+            .await
+            .unwrap();
 
         let make_stream = |stream_name| Stream {
             stream_name,
@@ -282,14 +293,9 @@ mod tests {
     #[tokio::test(flavor = "multi_thread")]
     async fn test_subscription_cld() {
         let addr = env::var("TEST_SERVER_ADDR").unwrap();
-        let mut client = Client::new(
-            addr,
-            ChannelProviderSettings {
-                concurrency_limit: None,
-            },
-        )
-        .await
-        .unwrap();
+        let mut client = Client::new(addr, ChannelProviderSettings::builder().build())
+            .await
+            .unwrap();
 
         let make_stream = |stream_name| Stream {
             stream_name,

src/hstreamdb/src/client.rs

@@ -118,14 +118,9 @@ mod tests {
         env_logger::init();
 
         let addr = env::var("TEST_SERVER_ADDR").unwrap();
-        let mut client = Client::new(
-            addr,
-            ChannelProviderSettings {
-                concurrency_limit: None,
-            },
-        )
-        .await
-        .unwrap();
+        let mut client = Client::new(addr, ChannelProviderSettings::builder().build())
+            .await
+            .unwrap();
         let stream_name = format!("stream-{}", rand_alphanumeric(10));
         client
             .create_stream(Stream {
@@ -145,9 +140,7 @@ mod tests {
                     .set_max_batch_len(100)
                     .set_batch_deadline(1000)
                     .build(),
-                ChannelProviderSettings {
-                    concurrency_limit: None,
-                },
+                ChannelProviderSettings::builder().build(),
                 None,
             )
             .await

src/hstreamdb/src/flow_controller.rs

@@ -133,14 +133,9 @@ mod tests {
     #[tokio::test(flavor = "multi_thread")]
     async fn test_partition_key_to_shard_id() {
         let addr = env::var("TEST_SERVER_ADDR").unwrap();
-        let mut client = Client::new(
-            addr,
-            ChannelProviderSettings {
-                concurrency_limit: None,
-            },
-        )
-        .await
-        .unwrap();
+        let mut client = Client::new(addr, ChannelProviderSettings::builder().build())
+            .await
+            .unwrap();
 
         let stream_name = rand_alphanumeric(20);