etcdv3 · davidli2010 · May 29, 2024 · Sep 26, 2023 · Oct 3, 2023 · Mar 30, 2024
diff --git a/src/auth.rs b/src/auth.rs
@@ -1,19 +1,22 @@
 //! Authentication service.
 
+use crate::error::Error;
+use crate::AuthClient;
+
 use http::{header::AUTHORIZATION, HeaderValue, Request};
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use std::task::{Context, Poll};
 use tower_service::Service;
 
 #[derive(Debug, Clone)]
 pub struct AuthService<S> {
     inner: S,
-    token: Option<Arc<HeaderValue>>,
+    token: Arc<Mutex<Option<HeaderValue>>>,
 }
 
 impl<S> AuthService<S> {
     #[inline]
-    pub fn new(inner: S, token: Option<Arc<HeaderValue>>) -> Self {
+    pub fn new(inner: S, token: Arc<Mutex<Option<HeaderValue>>>) -> Self {
         Self { inner, token }
     }
 }
@@ -33,12 +36,37 @@ where
 
     #[inline]
     fn call(&mut self, mut request: Request<Body>) -> Self::Future {
-        if let Some(token) = &self.token {
-            request
-                .headers_mut()
-                .insert(AUTHORIZATION, token.as_ref().clone());
+        if let Some(token) = self.token.lock().unwrap().as_ref() {
+            request.headers_mut().insert(AUTHORIZATION, token.clone());
         }
 
         self.inner.call(request)
     }
 }
+
+#[derive(Clone)]
+pub struct AuthHandle {
+    token: Arc<Mutex<Option<HeaderValue>>>,
+    cli: AuthClient,
+}
+
+impl AuthHandle {
+    #[inline]
+    pub(crate) fn new(token: Arc<Mutex<Option<HeaderValue>>>, cli: AuthClient) -> Self {
+        Self { token, cli }
+    }
+
+    /// Updates client authentication.
+    pub async fn update_auth(&mut self, name: String, password: String) -> Result<(), Error> {
+        let resp = self.cli.authenticate(name, password).await?;
+
+        self.token.lock().unwrap().replace(resp.token().parse()?);
+
+        Ok(())
+    }
+
+    /// Removes client authentication.
+    pub fn remove_auth(&mut self) {
+        self.token.lock().unwrap().take();
+    }
+}
diff --git a/src/client.rs b/src/client.rs
@@ -34,14 +34,16 @@ use crate::rpc::maintenance::{
     HashResponse, MaintenanceClient, MoveLeaderResponse, SnapshotStreaming, StatusResponse,
 };
 use crate::rpc::watch::{WatchClient, WatchOptions, WatchStream, Watcher};
+use crate::AuthHandle;
 #[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;
+use std::sync::{Arc, Mutex};
 use std::time::Duration;
 use tokio::sync::mpsc::Sender;
 
@@ -64,6 +66,7 @@ pub struct Client {
     cluster: ClusterClient,
     election: ElectionClient,
     options: Option<ConnectOptions>,
+    auth_handle: AuthHandle,
     tx: Sender<Change<Uri, Endpoint>>,
 }
 
@@ -104,7 +107,10 @@ impl Client {
         }
 
         let mut options = options;
-        let auth_token = Self::auth(channel.clone(), &mut options).await?;
+
+        let auth_token = Arc::new(Mutex::new(None));
+        Self::auth(channel.clone(), &mut options, &auth_token).await?;
+
         Ok(Self::build_client(channel, tx, auth_token, options))
     }
 
@@ -206,28 +212,29 @@ impl Client {
     async fn auth(
         channel: Channel,
         options: &mut Option<ConnectOptions>,
-    ) -> Result<Option<Arc<http::HeaderValue>>> {
+        auth_token: &Arc<Mutex<Option<HeaderValue>>>,
+    ) -> Result<()> {
         let user = match options {
-            None => return Ok(None),
+            None => return Ok(()),
             Some(opt) => {
                 // Take away the user, the password should not be stored in client.
                 opt.user.take()
             }
         };
 
         if let Some((name, password)) = user {
-            let mut tmp_auth = AuthClient::new(channel, None);
+            let mut tmp_auth = AuthClient::new(channel, auth_token.clone());
             let resp = tmp_auth.authenticate(name, password).await?;
-            Ok(Some(Arc::new(resp.token().parse()?)))
-        } else {
-            Ok(None)
+            auth_token.lock().unwrap().replace(resp.token().parse()?);
         }
+
+        Ok(())
     }
 
     fn build_client(
         channel: Channel,
         tx: Sender<Change<Uri, Endpoint>>,
-        auth_token: Option<Arc<http::HeaderValue>>,
+        auth_token: Arc<Mutex<Option<HeaderValue>>>,
         options: Option<ConnectOptions>,
     ) -> Self {
         let kv = KvClient::new(channel.clone(), auth_token.clone());
@@ -237,7 +244,9 @@ impl Client {
         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);
+        let election = ElectionClient::new(channel, auth_token.clone());
+
+        let auth_handle = AuthHandle::new(auth_token, auth.clone());
 
         Self {
             kv,
@@ -249,6 +258,7 @@ impl Client {
             cluster,
             election,
             options,
+            auth_handle,
             tx,
         }
     }
@@ -334,6 +344,12 @@ impl Client {
         self.election.clone()
     }
 
+    /// Gets a auth handle.
+    #[inline]
+    pub fn auth_handle(&self) -> AuthHandle {
+        self.auth_handle.clone()
+    }
+
     /// Put the given key into the key-value store.
     /// A put request increments the revision of the key-value store
     /// and generates one event in the event history.
@@ -726,6 +742,18 @@ impl Client {
     pub async fn resign(&mut self, option: Option<ResignOptions>) -> Result<ResignResponse> {
         self.election.resign(option).await
     }
+
+    /// Updates client authentication.
+    #[inline]
+    pub async fn update_auth(&mut self, name: String, password: String) -> Result<()> {
+        self.auth_handle.update_auth(name, password).await
+    }
+
+    /// Removes client authentication.
+    #[inline]
+    pub fn remove_auth(&mut self) {
+        self.auth_handle.remove_auth();
+    }
 }
 
 /// Options for `Connect` operation.

diff --git a/src/lib.rs b/src/lib.rs
@@ -64,6 +64,7 @@ mod error;
 mod openssl_tls;
 mod rpc;
 
+pub use crate::auth::AuthHandle;
 pub use crate::client::{Client, ConnectOptions};
 pub use crate::error::Error;
 pub use crate::rpc::auth::{

diff --git a/src/rpc/auth.rs b/src/rpc/auth.rs
@@ -35,6 +35,7 @@ use crate::rpc::pb::etcdserverpb::{
 use crate::rpc::ResponseHeader;
 use crate::rpc::{get_prefix, KeyRange};
 use http::HeaderValue;
+use std::sync::Mutex;
 use std::{string::String, sync::Arc};
 use tonic::{IntoRequest, Request};
 
@@ -48,7 +49,7 @@ pub struct AuthClient {
 impl AuthClient {
     /// Creates an auth client.
     #[inline]
-    pub(crate) fn new(channel: Channel, auth_token: Option<Arc<HeaderValue>>) -> Self {
+    pub(crate) fn new(channel: Channel, auth_token: Arc<Mutex<Option<HeaderValue>>>) -> Self {
         let inner = PbAuthClient::new(AuthService::new(channel, auth_token));
         Self { inner }
     }

diff --git a/src/rpc/cluster.rs b/src/rpc/cluster.rs
@@ -14,6 +14,7 @@ use crate::rpc::pb::etcdserverpb::{
 };
 use crate::rpc::ResponseHeader;
 use http::HeaderValue;
+use std::sync::Mutex;
 use std::{string::String, sync::Arc};
 use tonic::{IntoRequest, Request};
 
@@ -27,7 +28,7 @@ pub struct ClusterClient {
 impl ClusterClient {
     /// Creates an Cluster client.
     #[inline]
-    pub(crate) fn new(channel: Channel, auth_token: Option<Arc<HeaderValue>>) -> Self {
+    pub(crate) fn new(channel: Channel, auth_token: Arc<Mutex<Option<HeaderValue>>>) -> Self {
         let inner = PbClusterClient::new(AuthService::new(channel, auth_token));
         Self { inner }
     }

diff --git a/src/rpc/election.rs b/src/rpc/election.rs
@@ -12,6 +12,7 @@ use crate::rpc::pb::v3electionpb::{
 };
 use crate::rpc::{KeyValue, ResponseHeader};
 use http::HeaderValue;
+use std::sync::Mutex;
 use std::task::{Context, Poll};
 use std::{pin::Pin, sync::Arc};
 use tokio_stream::Stream;
@@ -479,7 +480,7 @@ impl From<&PbLeaderKey> for &LeaderKey {
 impl ElectionClient {
     /// Creates a election
     #[inline]
-    pub(crate) fn new(channel: Channel, auth_token: Option<Arc<HeaderValue>>) -> Self {
+    pub(crate) fn new(channel: Channel, auth_token: Arc<Mutex<Option<HeaderValue>>>) -> Self {
         let inner = PbElectionClient::new(AuthService::new(channel, auth_token));
         Self { inner }
     }

diff --git a/src/rpc/kv.rs b/src/rpc/kv.rs
@@ -20,7 +20,7 @@ use crate::rpc::pb::etcdserverpb::{
 };
 use crate::rpc::{get_prefix, KeyRange, KeyValue, ResponseHeader};
 use http::HeaderValue;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use tonic::{IntoRequest, Request};
 
 /// Client for KV operations.
@@ -33,7 +33,7 @@ pub struct KvClient {
 impl KvClient {
     /// Creates a kv client.
     #[inline]
-    pub(crate) fn new(channel: Channel, auth_token: Option<Arc<HeaderValue>>) -> Self {
+    pub(crate) fn new(channel: Channel, auth_token: Arc<Mutex<Option<HeaderValue>>>) -> Self {
         let inner = PbKvClient::new(AuthService::new(channel, auth_token));
         Self { inner }
     }

diff --git a/src/rpc/lease.rs b/src/rpc/lease.rs
@@ -17,7 +17,7 @@ use crate::rpc::ResponseHeader;
 use crate::Error;
 use http::HeaderValue;
 use std::pin::Pin;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use std::task::{Context, Poll};
 use tokio::sync::mpsc::{channel, Sender};
 use tokio_stream::wrappers::ReceiverStream;
@@ -34,7 +34,7 @@ pub struct LeaseClient {
 impl LeaseClient {
     /// Creates a `LeaseClient`.
     #[inline]
-    pub(crate) fn new(channel: Channel, auth_token: Option<Arc<HeaderValue>>) -> Self {
+    pub(crate) fn new(channel: Channel, auth_token: Arc<Mutex<Option<HeaderValue>>>) -> Self {
         let inner = PbLeaseClient::new(AuthService::new(channel, auth_token));
         Self { inner }
     }

diff --git a/src/rpc/lock.rs b/src/rpc/lock.rs
@@ -6,7 +6,7 @@ use crate::channel::Channel;
 use crate::error::Result;
 use crate::rpc::ResponseHeader;
 use http::HeaderValue;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use tonic::{IntoRequest, Request};
 use v3lockpb::lock_client::LockClient as PbLockClient;
 use v3lockpb::{
@@ -24,7 +24,7 @@ pub struct LockClient {
 impl LockClient {
     /// Creates a lock client.
     #[inline]
-    pub(crate) fn new(channel: Channel, auth_token: Option<Arc<HeaderValue>>) -> Self {
+    pub(crate) fn new(channel: Channel, auth_token: Arc<Mutex<Option<HeaderValue>>>) -> Self {
         let inner = PbLockClient::new(AuthService::new(channel, auth_token));
         Self { inner }
     }

diff --git a/src/rpc/maintenance.rs b/src/rpc/maintenance.rs
@@ -20,7 +20,7 @@ use crate::rpc::ResponseHeader;
 use etcdserverpb::maintenance_client::MaintenanceClient as PbMaintenanceClient;
 use etcdserverpb::AlarmMember as PbAlarmMember;
 use http::HeaderValue;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use tonic::codec::Streaming as PbStreaming;
 use tonic::{IntoRequest, Request};
 
@@ -556,7 +556,7 @@ impl MoveLeaderResponse {
 impl MaintenanceClient {
     /// Creates a maintenance client.
     #[inline]
-    pub(crate) fn new(channel: Channel, auth_token: Option<Arc<HeaderValue>>) -> Self {
+    pub(crate) fn new(channel: Channel, auth_token: Arc<Mutex<Option<HeaderValue>>>) -> Self {
         let inner = PbMaintenanceClient::new(AuthService::new(channel, auth_token));
         Self { inner }
     }

diff --git a/src/rpc/watch.rs b/src/rpc/watch.rs
@@ -15,7 +15,7 @@ use crate::rpc::pb::mvccpb::Event as PbEvent;
 use crate::rpc::{KeyRange, KeyValue, ResponseHeader};
 use http::HeaderValue;
 use std::pin::Pin;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use std::task::{Context, Poll};
 use tokio::sync::mpsc::{channel, Sender};
 use tokio_stream::{wrappers::ReceiverStream, Stream};
@@ -31,7 +31,7 @@ pub struct WatchClient {
 impl WatchClient {
     /// Creates a watch client.
     #[inline]
-    pub(crate) fn new(channel: Channel, auth_token: Option<Arc<HeaderValue>>) -> Self {
+    pub(crate) fn new(channel: Channel, auth_token: Arc<Mutex<Option<HeaderValue>>>) -> Self {
         let inner = PbWatchClient::new(AuthService::new(channel, auth_token));
         Self { inner }
     }