use crate::channel::Channel;
use crate::error::{Error, Result};
#[cfg(feature = "tls-openssl")]
use crate::openssl_tls::{self, OpenSslClientConfig, OpenSslConnector};
use crate::rpc::auth::Permission;
use crate::rpc::auth::{AuthClient, AuthDisableResponse, AuthEnableResponse};
use crate::rpc::auth::{
RoleAddResponse, RoleDeleteResponse, RoleGetResponse, RoleGrantPermissionResponse,
RoleListResponse, RoleRevokePermissionOptions, RoleRevokePermissionResponse, UserAddOptions,
UserAddResponse, UserChangePasswordResponse, UserDeleteResponse, UserGetResponse,
UserGrantRoleResponse, UserListResponse, UserRevokeRoleResponse,
};
use crate::rpc::cluster::{
ClusterClient, MemberAddOptions, MemberAddResponse, MemberListResponse, MemberPromoteResponse,
MemberRemoveResponse, MemberUpdateResponse,
};
use crate::rpc::election::{
CampaignResponse, ElectionClient, LeaderResponse, ObserveStream, ProclaimOptions,
ProclaimResponse, ResignOptions, ResignResponse,
};
use crate::rpc::kv::{
CompactionOptions, CompactionResponse, DeleteOptions, DeleteResponse, GetOptions, GetResponse,
KvClient, PutOptions, PutResponse, Txn, TxnResponse,
};
use crate::rpc::lease::{
LeaseClient, LeaseGrantOptions, LeaseGrantResponse, LeaseKeepAliveStream, LeaseKeeper,
LeaseLeasesResponse, LeaseRevokeResponse, LeaseTimeToLiveOptions, LeaseTimeToLiveResponse,
};
use crate::rpc::lock::{LockClient, LockOptions, LockResponse, UnlockResponse};
use crate::rpc::maintenance::{
AlarmAction, AlarmOptions, AlarmResponse, AlarmType, DefragmentResponse, HashKvResponse,
HashResponse, MaintenanceClient, MoveLeaderResponse, SnapshotStreaming, StatusResponse,
};
use crate::rpc::watch::{WatchClient, WatchOptions, WatchStream, Watcher};
#[cfg(feature = "tls-openssl")]
use crate::OpenSslResult;
#[cfg(feature = "tls")]
use crate::TlsOptions;
use http::uri::Uri;
use http::HeaderValue;
use std::str::FromStr;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use tokio::sync::mpsc::Sender;
use tonic::transport::Endpoint;
use tower::discover::Change;
const HTTP_PREFIX: &str = "http://";
const HTTPS_PREFIX: &str = "https://";
#[derive(Clone)]
pub struct Client {
kv: KvClient,
watch: WatchClient,
lease: LeaseClient,
lock: LockClient,
auth: AuthClient,
maintenance: MaintenanceClient,
cluster: ClusterClient,
election: ElectionClient,
options: Option<ConnectOptions>,
tx: Sender<Change<Uri, Endpoint>>,
}
impl Client {
pub async fn connect<E: AsRef<str>, S: AsRef<[E]>>(
endpoints: S,
options: Option<ConnectOptions>,
) -> Result<Self> {
let endpoints = {
let mut eps = Vec::new();
for e in endpoints.as_ref() {
let channel = Self::build_endpoint(e.as_ref(), &options)?;
eps.push(channel);
}
eps
};
if endpoints.is_empty() {
return Err(Error::InvalidArgs(String::from("empty endpoints")));
}
#[cfg(not(feature = "tls-openssl"))]
let (channel, tx) = Channel::balance_channel(64);
#[cfg(feature = "tls-openssl")]
let (channel, tx) = openssl_tls::balanced_channel(
options
.clone()
.and_then(|o| o.otls)
.unwrap_or_else(OpenSslConnector::create_default)?,
)?;
for endpoint in endpoints {
tx.send(Change::Insert(endpoint.uri().clone(), endpoint))
.await
.unwrap();
}
let mut options = options;
let auth_token = Arc::new(RwLock::new(None));
Self::auth(channel.clone(), &mut options, &auth_token).await?;
Ok(Self::build_client(channel, tx, auth_token, options))
}
fn build_endpoint(url: &str, options: &Option<ConnectOptions>) -> Result<Endpoint> {
#[cfg(feature = "tls-openssl")]
use tonic::transport::Channel;
let mut endpoint = if url.starts_with(HTTP_PREFIX) {
#[cfg(feature = "tls")]
if let Some(connect_options) = options {
if connect_options.tls.is_some() {
return Err(Error::InvalidArgs(String::from(
"TLS options are only supported with HTTPS URLs",
)));
}
}
Channel::builder(url.parse()?)
} else if url.starts_with(HTTPS_PREFIX) {
#[cfg(not(any(feature = "tls", feature = "tls-openssl")))]
return Err(Error::InvalidArgs(String::from(
"HTTPS URLs are only supported with the feature \"tls\"",
)));
#[cfg(all(feature = "tls-openssl", not(feature = "tls")))]
{
Channel::builder(url.parse()?)
}
#[cfg(feature = "tls")]
{
let tls = if let Some(connect_options) = options {
connect_options.tls.clone()
} else {
None
}
.unwrap_or_else(TlsOptions::new);
Channel::builder(url.parse()?).tls_config(tls)?
}
} else {
#[cfg(feature = "tls")]
{
let tls = if let Some(connect_options) = options {
connect_options.tls.clone()
} else {
None
};
match tls {
Some(tls) => {
let e = HTTPS_PREFIX.to_owned() + url;
Channel::builder(e.parse()?).tls_config(tls)?
}
None => {
let e = HTTP_PREFIX.to_owned() + url;
Channel::builder(e.parse()?)
}
}
}
#[cfg(all(feature = "tls-openssl", not(feature = "tls")))]
{
let pfx = if options.as_ref().and_then(|o| o.otls.as_ref()).is_some() {
HTTPS_PREFIX
} else {
HTTP_PREFIX
};
let e = pfx.to_owned() + url;
Channel::builder(e.parse()?)
}
#[cfg(all(not(feature = "tls"), not(feature = "tls-openssl")))]
{
let e = HTTP_PREFIX.to_owned() + url;
Channel::builder(e.parse()?)
}
};
if let Some(opts) = options {
if let Some((interval, timeout)) = opts.keep_alive {
endpoint = endpoint
.keep_alive_while_idle(opts.keep_alive_while_idle)
.http2_keep_alive_interval(interval)
.keep_alive_timeout(timeout);
}
if let Some(timeout) = opts.timeout {
endpoint = endpoint.timeout(timeout);
}
if let Some(timeout) = opts.connect_timeout {
endpoint = endpoint.connect_timeout(timeout);
}
if let Some(tcp_keepalive) = opts.tcp_keepalive {
endpoint = endpoint.tcp_keepalive(Some(tcp_keepalive));
}
}
Ok(endpoint)
}
async fn auth(
channel: Channel,
options: &mut Option<ConnectOptions>,
auth_token: &Arc<RwLock<Option<HeaderValue>>>,
) -> Result<()> {
let user = match options {
None => return Ok(()),
Some(opt) => {
opt.user.take()
}
};
if let Some((name, password)) = user {
let mut tmp_auth = AuthClient::new(channel, auth_token.clone());
let resp = tmp_auth.authenticate(name, password).await?;
auth_token.write().unwrap().replace(resp.token().parse()?);
}
Ok(())
}
fn build_client(
channel: Channel,
tx: Sender<Change<Uri, Endpoint>>,
auth_token: Arc<RwLock<Option<HeaderValue>>>,
options: Option<ConnectOptions>,
) -> Self {
let kv = KvClient::new(channel.clone(), auth_token.clone());
let watch = WatchClient::new(channel.clone(), auth_token.clone());
let lease = LeaseClient::new(channel.clone(), auth_token.clone());
let lock = LockClient::new(channel.clone(), auth_token.clone());
let auth = AuthClient::new(channel.clone(), auth_token.clone());
let cluster = ClusterClient::new(channel.clone(), auth_token.clone());
let maintenance = MaintenanceClient::new(channel.clone(), auth_token.clone());
let election = ElectionClient::new(channel, auth_token);
Self {
kv,
watch,
lease,
lock,
auth,
maintenance,
cluster,
election,
options,
tx,
}
}
#[inline]
pub async fn add_endpoint<E: AsRef<str>>(&self, endpoint: E) -> Result<()> {
let endpoint = Self::build_endpoint(endpoint.as_ref(), &self.options)?;
let tx = &self.tx;
tx.send(Change::Insert(endpoint.uri().clone(), endpoint))
.await
.map_err(|e| Error::EndpointError(format!("failed to add endpoint because of {}", e)))
}
#[inline]
pub async fn remove_endpoint<E: AsRef<str>>(&self, endpoint: E) -> Result<()> {
let uri = http::Uri::from_str(endpoint.as_ref())?;
let tx = &self.tx;
tx.send(Change::Remove(uri)).await.map_err(|e| {
Error::EndpointError(format!("failed to remove endpoint because of {}", e))
})
}
#[inline]
pub fn kv_client(&self) -> KvClient {
self.kv.clone()
}
#[inline]
pub fn watch_client(&self) -> WatchClient {
self.watch.clone()
}
#[inline]
pub fn lease_client(&self) -> LeaseClient {
self.lease.clone()
}
#[inline]
pub fn auth_client(&self) -> AuthClient {
self.auth.clone()
}
#[inline]
pub fn maintenance_client(&self) -> MaintenanceClient {
self.maintenance.clone()
}
#[inline]
pub fn cluster_client(&self) -> ClusterClient {
self.cluster.clone()
}
#[inline]
pub fn lock_client(&self) -> LockClient {
self.lock.clone()
}
#[inline]
pub fn election_client(&self) -> ElectionClient {
self.election.clone()
}
#[inline]
pub async fn put(
&mut self,
key: impl Into<Vec<u8>>,
value: impl Into<Vec<u8>>,
options: Option<PutOptions>,
) -> Result<PutResponse> {
self.kv.put(key, value, options).await
}
#[inline]
pub async fn get(
&mut self,
key: impl Into<Vec<u8>>,
options: Option<GetOptions>,
) -> Result<GetResponse> {
self.kv.get(key, options).await
}
#[inline]
pub async fn delete(
&mut self,
key: impl Into<Vec<u8>>,
options: Option<DeleteOptions>,
) -> Result<DeleteResponse> {
self.kv.delete(key, options).await
}
#[inline]
pub async fn compact(
&mut self,
revision: i64,
options: Option<CompactionOptions>,
) -> Result<CompactionResponse> {
self.kv.compact(revision, options).await
}
#[inline]
pub async fn txn(&mut self, txn: Txn) -> Result<TxnResponse> {
self.kv.txn(txn).await
}
#[inline]
pub async fn watch(
&mut self,
key: impl Into<Vec<u8>>,
options: Option<WatchOptions>,
) -> Result<(Watcher, WatchStream)> {
self.watch.watch(key, options).await
}
#[inline]
pub async fn lease_grant(
&mut self,
ttl: i64,
options: Option<LeaseGrantOptions>,
) -> Result<LeaseGrantResponse> {
self.lease.grant(ttl, options).await
}
#[inline]
pub async fn lease_revoke(&mut self, id: i64) -> Result<LeaseRevokeResponse> {
self.lease.revoke(id).await
}
#[inline]
pub async fn lease_keep_alive(
&mut self,
id: i64,
) -> Result<(LeaseKeeper, LeaseKeepAliveStream)> {
self.lease.keep_alive(id).await
}
#[inline]
pub async fn lease_time_to_live(
&mut self,
id: i64,
options: Option<LeaseTimeToLiveOptions>,
) -> Result<LeaseTimeToLiveResponse> {
self.lease.time_to_live(id, options).await
}
#[inline]
pub async fn leases(&mut self) -> Result<LeaseLeasesResponse> {
self.lease.leases().await
}
#[inline]
pub async fn lock(
&mut self,
name: impl Into<Vec<u8>>,
options: Option<LockOptions>,
) -> Result<LockResponse> {
self.lock.lock(name, options).await
}
#[inline]
pub async fn unlock(&mut self, key: impl Into<Vec<u8>>) -> Result<UnlockResponse> {
self.lock.unlock(key).await
}
#[inline]
pub async fn auth_enable(&mut self) -> Result<AuthEnableResponse> {
self.auth.auth_enable().await
}
#[inline]
pub async fn auth_disable(&mut self) -> Result<AuthDisableResponse> {
self.auth.auth_disable().await
}
#[inline]
pub async fn role_add(&mut self, name: impl Into<String>) -> Result<RoleAddResponse> {
self.auth.role_add(name).await
}
#[inline]
pub async fn role_delete(&mut self, name: impl Into<String>) -> Result<RoleDeleteResponse> {
self.auth.role_delete(name).await
}
#[inline]
pub async fn role_get(&mut self, name: impl Into<String>) -> Result<RoleGetResponse> {
self.auth.role_get(name).await
}
#[inline]
pub async fn role_list(&mut self) -> Result<RoleListResponse> {
self.auth.role_list().await
}
#[inline]
pub async fn role_grant_permission(
&mut self,
name: impl Into<String>,
perm: Permission,
) -> Result<RoleGrantPermissionResponse> {
self.auth.role_grant_permission(name, perm).await
}
#[inline]
pub async fn role_revoke_permission(
&mut self,
name: impl Into<String>,
key: impl Into<Vec<u8>>,
options: Option<RoleRevokePermissionOptions>,
) -> Result<RoleRevokePermissionResponse> {
self.auth.role_revoke_permission(name, key, options).await
}
#[inline]
pub async fn user_add(
&mut self,
name: impl Into<String>,
password: impl Into<String>,
options: Option<UserAddOptions>,
) -> Result<UserAddResponse> {
self.auth.user_add(name, password, options).await
}
#[inline]
pub async fn user_get(&mut self, name: impl Into<String>) -> Result<UserGetResponse> {
self.auth.user_get(name).await
}
#[inline]
pub async fn user_list(&mut self) -> Result<UserListResponse> {
self.auth.user_list().await
}
#[inline]
pub async fn user_delete(&mut self, name: impl Into<String>) -> Result<UserDeleteResponse> {
self.auth.user_delete(name).await
}
#[inline]
pub async fn user_change_password(
&mut self,
name: impl Into<String>,
password: impl Into<String>,
) -> Result<UserChangePasswordResponse> {
self.auth.user_change_password(name, password).await
}
#[inline]
pub async fn user_grant_role(
&mut self,
user: impl Into<String>,
role: impl Into<String>,
) -> Result<UserGrantRoleResponse> {
self.auth.user_grant_role(user, role).await
}
#[inline]
pub async fn user_revoke_role(
&mut self,
user: impl Into<String>,
role: impl Into<String>,
) -> Result<UserRevokeRoleResponse> {
self.auth.user_revoke_role(user, role).await
}
#[inline]
pub async fn alarm(
&mut self,
alarm_action: AlarmAction,
alarm_type: AlarmType,
options: Option<AlarmOptions>,
) -> Result<AlarmResponse> {
self.maintenance
.alarm(alarm_action, alarm_type, options)
.await
}
#[inline]
pub async fn status(&mut self) -> Result<StatusResponse> {
self.maintenance.status().await
}
#[inline]
pub async fn defragment(&mut self) -> Result<DefragmentResponse> {
self.maintenance.defragment().await
}
#[inline]
pub async fn hash(&mut self) -> Result<HashResponse> {
self.maintenance.hash().await
}
#[inline]
pub async fn hash_kv(&mut self, revision: i64) -> Result<HashKvResponse> {
self.maintenance.hash_kv(revision).await
}
#[inline]
pub async fn snapshot(&mut self) -> Result<SnapshotStreaming> {
self.maintenance.snapshot().await
}
#[inline]
pub async fn member_add<E: AsRef<str>, S: AsRef<[E]>>(
&mut self,
urls: S,
options: Option<MemberAddOptions>,
) -> Result<MemberAddResponse> {
let mut eps = Vec::new();
for e in urls.as_ref() {
let e = e.as_ref();
let url = if e.starts_with(HTTP_PREFIX) || e.starts_with(HTTPS_PREFIX) {
e.to_string()
} else {
HTTP_PREFIX.to_owned() + e
};
eps.push(url);
}
self.cluster.member_add(eps, options).await
}
#[inline]
pub async fn member_remove(&mut self, id: u64) -> Result<MemberRemoveResponse> {
self.cluster.member_remove(id).await
}
#[inline]
pub async fn member_update(
&mut self,
id: u64,
url: impl Into<Vec<String>>,
) -> Result<MemberUpdateResponse> {
self.cluster.member_update(id, url).await
}
#[inline]
pub async fn member_promote(&mut self, id: u64) -> Result<MemberPromoteResponse> {
self.cluster.member_promote(id).await
}
#[inline]
pub async fn member_list(&mut self) -> Result<MemberListResponse> {
self.cluster.member_list().await
}
#[inline]
pub async fn move_leader(&mut self, target_id: u64) -> Result<MoveLeaderResponse> {
self.maintenance.move_leader(target_id).await
}
#[inline]
pub async fn campaign(
&mut self,
name: impl Into<Vec<u8>>,
value: impl Into<Vec<u8>>,
lease: i64,
) -> Result<CampaignResponse> {
self.election.campaign(name, value, lease).await
}
#[inline]
pub async fn proclaim(
&mut self,
value: impl Into<Vec<u8>>,
options: Option<ProclaimOptions>,
) -> Result<ProclaimResponse> {
self.election.proclaim(value, options).await
}
#[inline]
pub async fn leader(&mut self, name: impl Into<Vec<u8>>) -> Result<LeaderResponse> {
self.election.leader(name).await
}
#[inline]
pub async fn observe(&mut self, name: impl Into<Vec<u8>>) -> Result<ObserveStream> {
self.election.observe(name).await
}
#[inline]
pub async fn resign(&mut self, option: Option<ResignOptions>) -> Result<ResignResponse> {
self.election.resign(option).await
}
pub async fn set_client_auth(&mut self, name: String, password: String) -> Result<()> {
self.auth.set_client_auth(name, password).await
}
pub fn remove_client_auth(&mut self) {
self.auth.remove_client_auth();
}
}
#[derive(Debug, Default, Clone)]
pub struct ConnectOptions {
user: Option<(String, String)>,
keep_alive: Option<(Duration, Duration)>,
keep_alive_while_idle: bool,
timeout: Option<Duration>,
connect_timeout: Option<Duration>,
tcp_keepalive: Option<Duration>,
#[cfg(feature = "tls")]
tls: Option<TlsOptions>,
#[cfg(feature = "tls-openssl")]
otls: Option<OpenSslResult<OpenSslConnector>>,
}
impl ConnectOptions {
#[inline]
pub fn with_user(mut self, name: impl Into<String>, password: impl Into<String>) -> Self {
self.user = Some((name.into(), password.into()));
self
}
#[cfg_attr(docsrs, doc(cfg(feature = "tls")))]
#[cfg(feature = "tls")]
#[inline]
pub fn with_tls(mut self, tls: TlsOptions) -> Self {
self.tls = Some(tls);
self
}
#[cfg_attr(docsrs, doc(cfg(feature = "tls-openssl")))]
#[cfg(feature = "tls-openssl")]
#[inline]
pub fn with_openssl_tls(mut self, otls: OpenSslClientConfig) -> Self {
self.otls = Some(otls.build());
self
}
#[inline]
pub fn with_keep_alive(mut self, interval: Duration, timeout: Duration) -> Self {
self.keep_alive = Some((interval, timeout));
self
}
#[inline]
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.timeout = Some(timeout);
self
}
#[inline]
pub fn with_connect_timeout(mut self, timeout: Duration) -> Self {
self.connect_timeout = Some(timeout);
self
}
#[inline]
pub fn with_tcp_keepalive(mut self, tcp_keepalive: Duration) -> Self {
self.tcp_keepalive = Some(tcp_keepalive);
self
}
#[inline]
pub fn with_keep_alive_while_idle(mut self, enabled: bool) -> Self {
self.keep_alive_while_idle = enabled;
self
}
#[inline]
pub const fn new() -> Self {
ConnectOptions {
user: None,
keep_alive: None,
keep_alive_while_idle: true,
timeout: None,
connect_timeout: None,
tcp_keepalive: None,
#[cfg(feature = "tls")]
tls: None,
#[cfg(feature = "tls-openssl")]
otls: None,
}
}
}