tls: bug fix and enhancement (#91)

Author: alissa-tung

Cargo.lock

@@ -1,7 +1,7 @@
 [package]
 edition = "2021"
 name = "hstreamdb"
-version = "0.2.0"
+version = "0.2.1"
 
 license = "BSD-3-Clause"
 description = "Rust client library for HStreamDB"

src/hstreamdb/Cargo.toml

@@ -5,10 +5,11 @@ 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, ClientTlsConfig, Endpoint};
+use tonic::transport::{Channel, Endpoint};
 
 use crate::client::get_available_node_addrs;
 use crate::common;
+use crate::tls::ClientTlsConfig;
 
 #[derive(Debug)]
 pub(crate) struct Request(

src/hstreamdb/src/channel_provider.rs

@@ -5,14 +5,15 @@ use hstreamdb_pb::{
     GetSubscriptionRequest, ListConsumersRequest, ListStreamsRequest, ListSubscriptionsRequest,
     LookupSubscriptionRequest, NodeState,
 };
-use tonic::transport::{Channel, ClientTlsConfig};
+use tonic::transport::{Channel, Endpoint};
 use tonic::Request;
 use url::Url;
 
 use crate::appender::Appender;
 use crate::channel_provider::{new_channel_provider, ChannelProviderSettings, Channels};
 use crate::common::Error::PBUnwrapError;
 use crate::producer::{FlushCallback, FlushSettings, Producer};
+use crate::tls::ClientTlsConfig;
 use crate::{common, flow_controller, format_url, producer};
 
 pub struct Client {
@@ -29,41 +30,29 @@ impl Client {
     where
         Destination: std::convert::Into<String>,
     {
-        const HSTREAM_PREFIX: &str = "hstream";
-        let server_url = server_url.into();
-        let (url_scheme, url) = {
-            let url = {
-                let mut url = Url::parse(&server_url)?;
-                if url.port().is_none() {
-                    url.set_port(Some(6570))
-                        .map_err(|()| common::Error::SetPortError(server_url.to_string()))?;
-                }
-                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 if url.scheme() == "hstreams" {
-                let url_scheme = "https";
-                (
-                    url_scheme.to_string(),
-                    Url::parse(format!("{url_scheme}{server_url}").as_str())?,
-                )
-            } else {
-                (url.scheme().to_string(), url)
+        let server_url: String = server_url.into();
+        Url::parse(&server_url)?;
+        let server_url = set_scheme(&server_url).ok_or(common::Error::InvalidUrl(server_url))?;
+        let (url_scheme, server_url) = {
+            let mut server_url = Url::parse(&server_url)?;
+            let port = server_url.port();
+            if port.is_none() {
+                server_url
+                    .set_port(Some(6570))
+                    .map_err(|()| common::Error::InvalidUrl(server_url.to_string()))?;
             }
+            (server_url.scheme().to_string(), server_url)
         };
-        let mut hstream_api_client = HStreamApiClient::connect(String::from(url)).await?;
+
+        log::debug!("client init connect: scheme = {url_scheme}, url = {server_url}");
         let tls_config = channel_provider_settings.client_tls_config.clone();
+        let mut hstream_api_client = HStreamApiClient::new({
+            let mut endpoint = Endpoint::new(server_url.to_string())?;
+            if let Some(tls_config) = tls_config.clone() {
+                endpoint = endpoint.tls_config(tls_config)?;
+            }
+            endpoint.connect().await?
+        });
         let channels = new_channel_provider(
             &url_scheme,
             &mut hstream_api_client,
@@ -78,6 +67,13 @@ impl Client {
     }
 }
 
+fn set_scheme(url: &str) -> Option<String> {
+    Some(
+        url.replace("hstream://", "http://")
+            .replace("hstreams://", "https://"),
+    )
+}
+
 impl Client {
     async fn new_channel_provider(
         &self,

src/hstreamdb/src/client.rs

@@ -95,7 +95,7 @@ pub enum Error {
     #[error(transparent)]
     AppenderSendError(producer::SendError),
     #[error("the URL {0} is cannot-be-a-base, or does not have a host, or has the file scheme")]
-    SetPortError(String),
+    InvalidUrl(String),
 }
 
 #[derive(Debug, thiserror::Error)]

src/hstreamdb/src/common.rs

@@ -8,6 +8,7 @@ pub mod consumer;
 mod flow_controller;
 pub mod producer;
 pub mod reader;
+pub mod tls;
 pub mod utils;
 
 pub use channel_provider::ChannelProviderSettings;

src/hstreamdb/src/lib.rs

@@ -0,0 +1,2 @@
+pub use tonic::transport::channel::ClientTlsConfig;
+pub use tonic::transport::{Certificate, Identity};

src/hstreamdb/src/tls.rs

@@ -0,0 +1,32 @@
+// use std::{env, fs};
+
+// use hstreamdb::tls::{Certificate, ClientTlsConfig, Identity};
+// use hstreamdb::{ChannelProviderSettings, Client};
+
+// async fn _test_tls_impl() {
+//     env::set_var("RUST_LOG", "DEBUG");
+//     env_logger::init();
+
+//     let server_url: &str = todo!();
+//     let tls_dir: &str = todo!();
+
+//     let ca_certificate =
+// Certificate::from_pem(fs::read(format!("{tls_dir}/root_ca.crt")).unwrap());
+//     let cert = fs::read(format!("{tls_dir}/client.crt")).unwrap();
+//     let key = fs::read(format!("{tls_dir}/client.key")).unwrap();
+
+//     let client = Client::new(
+//         server_url,
+//         ChannelProviderSettings::builder()
+//             .set_tls_config(
+//                 ClientTlsConfig::new()
+//                     .ca_certificate(ca_certificate)
+//                     .identity(Identity::from_pem(cert, key)),
+//             )
+//             .build(),
+//     )
+//     .await
+//     .unwrap();
+
+//     log::info!("{:?}", client.list_streams().await.unwrap());
+// }