refactor!: add SessionInfo to box session id, password and readonly

We also use it to reject readonly session reestablish request as
readonly sessions are local to connected node.

Author: kezhuw

src/client/mod.rs

@@ -13,7 +13,7 @@ use thiserror::Error;
 use tokio::sync::{mpsc, watch};
 
 pub use self::watcher::{OneshotWatcher, PersistentWatcher, StateWatcher};
-use super::session::{Depot, MarshalledRequest, Session, SessionOperation, WatchReceiver, PASSWORD_LEN};
+use super::session::{Depot, MarshalledRequest, Session, SessionOperation, WatchReceiver};
 use crate::acl::{Acl, Acls, AuthUser};
 use crate::chroot::{Chroot, ChrootPath, OwnedChroot};
 use crate::endpoint::{self, IterableEndpoints};
@@ -44,7 +44,7 @@ use crate::proto::{
 pub use crate::proto::{EnsembleUpdate, Stat};
 use crate::record::{self, Record, StaticRecord};
 use crate::session::StateReceiver;
-pub use crate::session::{EventType, SessionId, SessionState, WatchedEvent};
+pub use crate::session::{EventType, SessionId, SessionInfo, SessionState, WatchedEvent};
 use crate::tls::TlsOptions;
 use crate::util;
 
@@ -215,7 +215,7 @@ impl CreateSequence {
 pub struct Client {
     chroot: OwnedChroot,
     version: Version,
-    session: (SessionId, Vec<u8>),
+    session: SessionInfo,
     session_timeout: Duration,
     requester: mpsc::UnboundedSender<SessionOperation>,
     state_watcher: StateWatcher,
@@ -243,7 +243,7 @@ impl Client {
     pub(crate) fn new(
         chroot: OwnedChroot,
         version: Version,
-        session: (SessionId, Vec<u8>),
+        session: SessionInfo,
         timeout: Duration,
         requester: mpsc::UnboundedSender<SessionOperation>,
         state_receiver: watch::Receiver<SessionState>,
@@ -265,18 +265,18 @@ impl Client {
         self.chroot.path()
     }
 
-    /// ZooKeeper session id.
-    pub fn session_id(&self) -> SessionId {
-        self.session.0
+    /// ZooKeeper session info.
+    pub fn session(&self) -> &SessionInfo {
+        &self.session
     }
 
-    /// Session password.
-    pub fn session_password(&self) -> &[u8] {
-        self.session.1.as_slice()
+    /// ZooKeeper session id.
+    pub fn session_id(&self) -> SessionId {
+        self.session().id()
     }
 
     /// Consumes this instance into session info.
-    pub fn into_session(self) -> (SessionId, Vec<u8>) {
+    pub fn into_session(self) -> SessionInfo {
         self.session
     }
 
@@ -1538,7 +1538,7 @@ pub(crate) struct Version(u32, u32, u32);
 pub struct Connector {
     tls: Option<TlsOptions>,
     authes: Vec<AuthPacket>,
-    session: Option<(SessionId, Vec<u8>)>,
+    session: Option<SessionInfo>,
     readonly: bool,
     detached: bool,
     server_version: Version,
@@ -1590,8 +1590,8 @@ impl Connector {
     }
 
     /// Specifies session to reestablish.
-    pub fn session(&mut self, id: SessionId, password: Vec<u8>) -> &mut Self {
-        self.session = Some((id, password));
+    pub fn session(&mut self, session: SessionInfo) -> &mut Self {
+        self.session = Some(session);
         self
     }
 
@@ -1623,14 +1623,12 @@ impl Connector {
 
     async fn connect_internally(&mut self, secure: bool, cluster: &str) -> Result<Client> {
         let (endpoints, chroot) = endpoint::parse_connect_string(cluster, secure)?;
-        if let Some((id, password)) = &self.session {
-            if id.0 == 0 {
-                return Err(Error::BadArguments(&"session id must not be 0"));
-            } else if password.is_empty() {
-                return Err(Error::BadArguments(&formatcp!(
-                    "session password is empty, it should have length of {}",
-                    PASSWORD_LEN
-                )));
+        if let Some(session) = self.session.as_ref() {
+            if session.is_readonly() {
+                return Err(Error::new_other(
+                    format!("can't reestablish readonly and hence local session {}", session.id()),
+                    None,
+                ));
             }
         }
         if self.session_timeout < Duration::ZERO {
@@ -1653,7 +1651,7 @@ impl Connector {
         let mut connecting_depot = Depot::for_connecting();
         let conn = session.start(&mut endpoints, &mut buf, &mut connecting_depot).await?;
         let (sender, receiver) = mpsc::unbounded_channel();
-        let session_info = (session.session_id, session.session_password.clone());
+        let session_info = session.session.clone();
         let session_timeout = session.session_timeout;
         tokio::spawn(async move {
             session.serve(endpoints, conn, buf, connecting_depot, receiver).await;
@@ -1729,12 +1727,6 @@ impl ClientBuilder {
         self
     }
 
-    /// Specifies session to reestablish.
-    pub fn with_session(&mut self, id: SessionId, password: Vec<u8>) -> &mut Self {
-        self.connector.session(id, password);
-        self
-    }
-
     /// Specifies client assumed server version of ZooKeeper cluster.
     ///
     /// Client will issue server compatible protocol to avoid [Error::Unimplemented] for some

src/error.rs

@@ -127,12 +127,19 @@ impl Error {
         }
     }
 
+    pub(crate) fn new_other(
+        message: impl Into<Arc<String>>,
+        source: Option<Arc<dyn std::error::Error + Send + Sync + 'static>>,
+    ) -> Self {
+        Self::Other(OtherError { message: message.into(), source })
+    }
+
     pub(crate) fn other(message: impl Into<String>, source: impl std::error::Error + Send + Sync + 'static) -> Self {
-        Self::Other(OtherError { message: Arc::new(message.into()), source: Some(Arc::new(source)) })
+        Self::new_other(message.into(), Some(Arc::new(source)))
     }
 
     pub(crate) fn other_from(source: impl std::error::Error + Send + Sync + 'static) -> Self {
-        Self::Other(OtherError { message: Arc::new(source.to_string()), source: Some(Arc::new(source)) })
+        Self::new_other(source.to_string(), Some(Arc::new(source)))
     }
 }
 

src/session/mod.rs

@@ -27,7 +27,7 @@ pub use self::request::{
     StateReceiver,
     StateResponser,
 };
-pub use self::types::{EventType, SessionId, SessionState, WatchedEvent};
+pub use self::types::{EventType, SessionId, SessionInfo, SessionState, WatchedEvent};
 pub use self::watch::{OneshotReceiver, PersistentReceiver, WatchReceiver};
 use self::watch::{WatchManager, WatcherId};
 use crate::deadline::Deadline;
@@ -71,11 +71,9 @@ pub struct Session {
     ping_timeout: Duration,
     session_expired_timeout: Duration,
 
-    pub session_id: SessionId,
+    pub session: SessionInfo,
     session_state: SessionState,
     pub session_timeout: Duration,
-    pub session_password: Vec<u8>,
-    session_readonly: bool,
 
     pub authes: Vec<MarshalledRequest>,
     state_sender: tokio::sync::watch::Sender<SessionState>,
@@ -86,16 +84,15 @@ pub struct Session {
 
 impl Session {
     pub fn new(
-        session: Option<(SessionId, Vec<u8>)>,
+        session: Option<SessionInfo>,
         authes: &[AuthPacket],
         readonly: bool,
         detached: bool,
         tls_config: ClientConfig,
         session_timeout: Duration,
         connection_timeout: Duration,
     ) -> (Session, tokio::sync::watch::Receiver<SessionState>) {
-        let (session_id, session_password) =
-            session.unwrap_or_else(|| (SessionId(0), Vec::with_capacity(PASSWORD_LEN)));
+        let session = session.unwrap_or_else(|| SessionInfo::new(SessionId(0), Vec::with_capacity(PASSWORD_LEN)));
         let (state_sender, state_receiver) = tokio::sync::watch::channel(SessionState::Disconnected);
         let now = Instant::now();
         let (watch_manager, unwatch_receiver) = WatchManager::new();
@@ -114,11 +111,9 @@ impl Session {
             session_expired_timeout: Duration::ZERO,
             connector: Connector::new(tls_config),
 
-            session_id,
+            session,
             session_timeout,
             session_state: SessionState::Disconnected,
-            session_password,
-            session_readonly: session_id.0 == 0,
 
             authes: authes.iter().map(|auth| MarshalledRequest::new(OpCode::Auth, auth)).collect(),
             state_sender,
@@ -132,7 +127,7 @@ impl Session {
 
     fn is_readonly_allowed(&self) -> bool {
         // Session downgrade is not allowed as partitioned session will expired finally by quorum.
-        self.readonly && self.session_readonly
+        self.readonly && self.session.readonly
     }
 
     async fn close_requester<T: RequestOperation>(mut requester: mpsc::UnboundedReceiver<T>, err: &Error) {
@@ -215,7 +210,7 @@ impl Session {
     ) {
         if let Err(err) = self.serve_session(endpoints, &mut conn, buf, depot, requester, unwatch_requester).await {
             self.resolve_serve_error(&err);
-            log::info!("ZooKeeper session {} state {} error {}", self.session_id, self.session_state, err);
+            log::info!("ZooKeeper session {} state {} error {}", self.session.id, self.session_state, err);
             depot.error(&err);
         } else {
             self.change_state(SessionState::Disconnected);
@@ -301,7 +296,7 @@ impl Session {
             depot.pop_ping()?;
             if let Some(last_ping) = self.last_ping.take() {
                 let elapsed = Instant::now() - last_ping;
-                log::debug!("ZooKeeper session {} got ping response after {}ms", self.session_id, elapsed.as_millis());
+                log::debug!("ZooKeeper session {} got ping response after {}ms", self.session.id, elapsed.as_millis());
             }
             return Ok(());
         }
@@ -349,7 +344,7 @@ impl Session {
     }
 
     fn complete_connect(&mut self) {
-        let state = if self.session_readonly { SessionState::ConnectedReadOnly } else { SessionState::SyncConnected };
+        let state = if self.session.readonly { SessionState::ConnectedReadOnly } else { SessionState::SyncConnected };
         self.change_state(state);
     }
 
@@ -361,11 +356,11 @@ impl Session {
         } else if !self.is_readonly_allowed() && response.readonly {
             return Err(Error::ConnectionLoss);
         }
-        self.session_id = SessionId(response.session_id);
+        self.session.id = SessionId(response.session_id);
+        self.session.password.clear();
+        self.session.password.extend_from_slice(response.password);
+        self.session.readonly = response.readonly;
         self.reset_timeout(Duration::from_millis(response.session_timeout as u64));
-        self.session_password.clear();
-        self.session_password.extend_from_slice(response.password);
-        self.session_readonly = response.readonly;
         self.complete_connect();
         Ok(())
     }
@@ -450,7 +445,7 @@ impl Session {
         unwatch_requester: &mut mpsc::UnboundedReceiver<(WatcherId, StateResponser)>,
     ) -> Result<(), Error> {
         let mut seek_for_writable =
-            if self.session_readonly { Some(self.connector.clone().seek_for_writable(endpoints)) } else { None };
+            if self.session.readonly { Some(self.connector.clone().seek_for_writable(endpoints)) } else { None };
         let mut tick = time::interval(self.tick_timeout);
         tick.set_missed_tick_behavior(time::MissedTickBehavior::Skip);
         let mut channel_closed = false;
@@ -519,8 +514,8 @@ impl Session {
             protocol_version: 0,
             last_zxid_seen: 0,
             timeout: self.session_timeout.as_millis() as i32,
-            session_id: if self.session_readonly { 0 } else { self.session_id.0 },
-            password: self.session_password.as_slice(),
+            session_id: if self.session.readonly { 0 } else { self.session.id.0 },
+            password: self.session.password.as_slice(),
             readonly: self.is_readonly_allowed(),
         };
         log::trace!("Sending connect request: {request:?}");
@@ -570,7 +565,7 @@ impl Session {
                 Err(err)
             },
             _ => {
-                log::info!("ZooKeeper succeeds to establish session({}) to {}", self.session_id, endpoint);
+                log::info!("ZooKeeper succeeds to establish session({}) to {}", self.session.id, endpoint);
                 Ok(conn)
             },
         }
@@ -582,7 +577,7 @@ impl Session {
         buf: &mut Vec<u8>,
         depot: &mut Depot,
     ) -> Result<Connection, Error> {
-        let session_timeout = if self.session_id.0 == 0 { self.session_timeout } else { self.session_expired_timeout };
+        let session_timeout = if self.session.id.0 == 0 { self.session_timeout } else { self.session_expired_timeout };
         let mut deadline = Deadline::until(self.last_recv + session_timeout);
         let mut last_error = match self.start_once(endpoints, &mut deadline, buf, depot).await {
             Err(err) => err,

src/session/types.rs

@@ -1,3 +1,4 @@
+use derive_where::derive_where;
 use num_enum::{IntoPrimitive, TryFromPrimitive};
 use strum::EnumIter;
 
@@ -6,7 +7,7 @@ use crate::proto::AddWatchMode;
 use crate::util;
 
 /// Thin wrapper for zookeeper session id. It prints in hex format headed with 0x.
-#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+#[derive(Copy, Clone, PartialEq, Eq)]
 pub struct SessionId(pub i64);
 
 impl std::fmt::Display for SessionId {
@@ -15,6 +16,40 @@ impl std::fmt::Display for SessionId {
     }
 }
 
+impl std::fmt::Debug for SessionId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        std::fmt::Display::fmt(self, f)
+    }
+}
+
+/// ZooKeeper session info.
+#[derive(Clone)]
+#[derive_where(Debug)]
+pub struct SessionInfo {
+    pub(crate) id: SessionId,
+    #[derive_where(skip(Debug))]
+    pub(crate) password: Vec<u8>,
+    pub(crate) readonly: bool,
+}
+
+impl SessionInfo {
+    pub(crate) fn new(id: SessionId, password: Vec<u8>) -> Self {
+        Self { id, password, readonly: id.0 == 0 }
+    }
+
+    /// Session id.
+    pub fn id(&self) -> SessionId {
+        self.id
+    }
+
+    /// Is this an readonly session ?
+    ///
+    /// Readonly sessions are local to connected server thus not eligible for session reestablishment.
+    pub fn is_readonly(&self) -> bool {
+        self.readonly
+    }
+}
+
 /// ZooKeeper session states.
 #[derive(Copy, Clone, Debug, PartialEq, Eq, strum::Display)]
 pub enum SessionState {

tests/zookeeper.rs

@@ -164,11 +164,11 @@ async fn test_connect_session_expired() {
     let cluster = Cluster::new().await;
     let client = cluster.custom_client(None, |connector| connector.detached()).await.unwrap();
     let timeout = client.session_timeout();
-    let (id, password) = client.into_session();
+    let session = client.into_session();
 
     tokio::time::sleep(timeout * 2).await;
 
-    assert_that!(cluster.custom_client(None, |connector| connector.session(id, password)).await.unwrap_err())
+    assert_that!(cluster.custom_client(None, |connector| connector.session(session)).await.unwrap_err())
         .is_equal_to(zk::Error::SessionExpired);
 }
 
@@ -1701,10 +1701,10 @@ async fn test_client_drop() {
 
     let mut state_watcher = client.state_watcher();
     tokio::time::sleep(Duration::from_secs(20)).await;
-    let (id, password) = client.into_session();
+    let session = client.into_session();
     assert_eq!(zk::SessionState::Closed, state_watcher.changed().await);
 
-    cluster.custom_client(None, |connector| connector.session(id, password)).await.unwrap_err();
+    cluster.custom_client(None, |connector| connector.session(session)).await.unwrap_err();
 }
 
 #[test_log::test(tokio::test)]
@@ -1713,10 +1713,10 @@ async fn test_client_detach() {
     let client = cluster.custom_client(None, |connector| connector.detached()).await.unwrap();
 
     let mut state_watcher = client.state_watcher();
-    let (id, password) = client.into_session();
+    let session = client.into_session();
     assert_eq!(zk::SessionState::Closed, state_watcher.changed().await);
 
-    cluster.custom_client(None, |connector| connector.session(id, password)).await.unwrap();
+    cluster.custom_client(None, |connector| connector.session(session)).await.unwrap();
 }
 
 fn generate_ca_cert() -> (Certificate, String) {
@@ -1821,6 +1821,11 @@ async fn test_readonly(tls: bool) {
     };
     assert_that!(client.create("/y", b"", PERSISTENT_OPEN).await.unwrap_err()).is_equal_to(zk::Error::NotReadOnly);
 
+    let session = client.session().clone();
+    assert_eq!(session.is_readonly(), true);
+    assert_that!(zk::Client::connector().session(session).connect("localhost:4001").await.unwrap_err().to_string())
+        .contains("readonly");
+
     let mut state_watcher = client.state_watcher();
     assert_eq!(state_watcher.state(), zk::SessionState::ConnectedReadOnly);