From ec8dbccaa37d40590778817ac23abf98fdb7b57a Mon Sep 17 00:00:00 2001 From: limbooverlambda Date: Wed, 26 Jun 2024 18:52:33 -0700 Subject: [PATCH] Resolving the TLS url issue (#459) * resolving the url issue Signed-off-by: limbooverlambda * fix formatting Signed-off-by: limbooverlambda * make check fixes Signed-off-by: limbooverlambda --------- Signed-off-by: limbooverlambda --- src/common/security.rs | 47 +++++++++++++++++++++++-------------- src/kv/key.rs | 1 - src/kv/mod.rs | 1 - src/raw/client.rs | 1 - src/transaction/requests.rs | 2 +- 5 files changed, 31 insertions(+), 21 deletions(-) diff --git a/src/common/security.rs b/src/common/security.rs index 483759cf..89e074b3 100644 --- a/src/common/security.rs +++ b/src/common/security.rs @@ -8,10 +8,10 @@ use std::time::Duration; use log::info; use regex::Regex; -use tonic::transport::Certificate; use tonic::transport::Channel; use tonic::transport::ClientTlsConfig; use tonic::transport::Identity; +use tonic::transport::{Certificate, Endpoint}; use crate::internal_err; use crate::Result; @@ -77,27 +77,40 @@ impl SecurityManager { where Factory: FnOnce(Channel) -> Client, { - let addr = "http://".to_string() + &SCHEME_REG.replace(addr, ""); - info!("connect to rpc server at endpoint: {:?}", addr); + let channel = if !self.ca.is_empty() { + self.tls_channel(addr).await? + } else { + self.default_channel(addr).await? + }; + let ch = channel.connect().await?; - let mut builder = Channel::from_shared(addr)? - .tcp_keepalive(Some(Duration::from_secs(10))) - .keep_alive_timeout(Duration::from_secs(3)); + Ok(factory(ch)) + } - if !self.ca.is_empty() { - let tls = ClientTlsConfig::new() - .ca_certificate(Certificate::from_pem(&self.ca)) - .identity(Identity::from_pem( - &self.cert, - load_pem_file("private key", &self.key)?, - )); - builder = builder.tls_config(tls)?; - }; + async fn tls_channel(&self, addr: &str) -> Result { + let addr = "https://".to_string() + &SCHEME_REG.replace(addr, ""); + let builder = self.endpoint(addr.to_string())?; + let tls = ClientTlsConfig::new() + .ca_certificate(Certificate::from_pem(&self.ca)) + .identity(Identity::from_pem( + &self.cert, + load_pem_file("private key", &self.key)?, + )); + let builder = builder.tls_config(tls)?; + Ok(builder) + } - let ch = builder.connect().await?; + async fn default_channel(&self, addr: &str) -> Result { + let addr = "http://".to_string() + &SCHEME_REG.replace(addr, ""); + self.endpoint(addr) + } - Ok(factory(ch)) + fn endpoint(&self, addr: String) -> Result { + let endpoint = Channel::from_shared(addr)? + .tcp_keepalive(Some(Duration::from_secs(10))) + .keep_alive_timeout(Duration::from_secs(3)); + Ok(endpoint) } } diff --git a/src/kv/key.rs b/src/kv/key.rs index 1b4f0606..94fe8a94 100644 --- a/src/kv/key.rs +++ b/src/kv/key.rs @@ -2,7 +2,6 @@ use std::fmt; use std::ops::Bound; -use std::u8; #[allow(unused_imports)] #[cfg(test)] diff --git a/src/kv/mod.rs b/src/kv/mod.rs index 489110e6..d0958ee2 100644 --- a/src/kv/mod.rs +++ b/src/kv/mod.rs @@ -1,6 +1,5 @@ // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0. use std::fmt; -use std::u8; mod bound_range; pub mod codec; diff --git a/src/raw/client.rs b/src/raw/client.rs index 71d40b2a..9a166278 100644 --- a/src/raw/client.rs +++ b/src/raw/client.rs @@ -3,7 +3,6 @@ use core::ops::Range; use std::str::FromStr; use std::sync::Arc; -use std::u32; use futures::StreamExt; use log::debug; diff --git a/src/transaction/requests.rs b/src/transaction/requests.rs index 4f0a6174..231c9e5a 100644 --- a/src/transaction/requests.rs +++ b/src/transaction/requests.rs @@ -252,7 +252,7 @@ pub fn new_prewrite_request( req.start_version = start_version; req.lock_ttl = lock_ttl; // FIXME: Lite resolve lock is currently disabled - req.txn_size = std::u64::MAX; + req.txn_size = u64::MAX; req }