feat: support SASL authentication (#46)

Closes #12.

Author: kezhuw

.github/workflows/ci.yml

@@ -31,8 +31,11 @@ jobs:
       with:
           toolchain: stable
           override: true
+    - name: Install krb5
+      run: |
+        sudo apt install -y libkrb5-dev
     - name: Build code
-      run: cargo build
+      run: cargo build --all-features
   test:
     needs: [build]
     runs-on: ubuntu-latest
@@ -46,8 +49,11 @@ jobs:
       with:
           toolchain: stable
           override: true
+    - name: Install krb5
+      run: |
+        sudo apt install -y libkrb5-dev
     - name: Test code
-      run: cargo test -- --nocapture
+      run: cargo test --all-features -- --nocapture
   coverage:
     if: github.event_name == 'pull_request' || (github.event_name == 'push' && github.ref_type == 'branch' && github.ref_name == 'master')
     needs: [test]
@@ -65,9 +71,12 @@ jobs:
           override: true
           components: llvm-tools-preview
     - run: cargo install grcov
-    - run: cargo build --verbose
+    - name: Install krb5
+      run: |
+        sudo apt install -y libkrb5-dev
+    - run: cargo build --all-features --verbose
     - name: Run tests
-      run: LLVM_PROFILE_FILE="zookeeper-client-%p-%m.profraw" cargo test --verbose -- --nocapture
+      run: LLVM_PROFILE_FILE="zookeeper-client-%p-%m.profraw" cargo test --all-features --verbose -- --nocapture
     - name: Generate coverage report
       run: grcov $(find . -name "zookeeper-*.profraw" -print) --binary-path ./target/debug/ -s . -t lcov --branch --ignore-not-existing --ignore "/*" -o lcov.info
     - name: Upload to codecov.io
@@ -86,15 +95,21 @@ jobs:
           toolchain: stable
           override: true
           components: clippy
+    - name: Install krb5
+      run: |
+        sudo apt install -y libkrb5-dev
     - name: Lint code
-      run: cargo clippy --no-deps -- -D clippy::all
+      run: cargo clippy --all-features --no-deps -- -D clippy::all
   release:
     if: github.event_name == 'push' && github.ref_type == 'tag'
     needs: [build, test, lint]
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v2
     - uses: Swatinem/rust-cache@v1
+    - name: Install krb5
+      run: |
+        sudo apt install -y libkrb5-dev
     - name: publish crate
       env:
         CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}

Cargo.toml

@@ -14,6 +14,12 @@ rust-version = "1.65"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[features]
+default = []
+sasl = ["sasl-gssapi", "sasl-digest-md5"]
+sasl-digest-md5 = ["rsasl/unstable_custom_mechanism", "md5", "linkme", "hex"]
+sasl-gssapi = ["rsasl/gssapi"]
+
 [dependencies]
 bytes = "1.1.0"
 tokio = {version = "1.15.0", features = ["full"]}
@@ -35,6 +41,10 @@ derive-where = "1.2.7"
 tokio-rustls = "0.26.0"
 fastrand = "2.0.2"
 tracing = "0.1.40"
+rsasl = { version = "2.0.1", default-features = false, features = ["provider", "config_builder", "registry_static", "std"], optional = true }
+md5 = { version = "0.7.0", optional = true }
+hex = { version = "0.4.3", optional = true }
+linkme = { version = "0.2", optional = true }
 
 [dev-dependencies]
 test-log = { version = "0.2.15", features = ["log", "trace"] }

Makefile

@@ -9,10 +9,10 @@ check_fmt:
 	cargo +nightly fmt --all -- --check
 
 lint:
-	cargo clippy --no-deps -- -D clippy::all
+	cargo clippy --all-features --no-deps -- -D clippy::all
 
 build:
-	cargo build
+	cargo build --all-features
 
 test:
-	cargo test
+	cargo test --all-features

README.md

@@ -79,9 +79,6 @@ latch.create("/app/data", b"data", &zk::CreateMode::Ephemeral.with_acls(zk::Acls
 
 For more examples, see [zookeeper.rs](tests/zookeeper.rs).
 
-## TODO
-* [ ] Sasl authentication
-
 ## License
 The MIT License (MIT). See [LICENSE](LICENSE) for the full license text.
 

src/client/mod.rs

@@ -45,6 +45,8 @@ use crate::proto::{
 };
 pub use crate::proto::{EnsembleUpdate, Stat};
 use crate::record::{self, Record, StaticRecord};
+#[cfg(any(feature = "sasl-digest-md5", feature = "sasl-gssapi"))]
+use crate::sasl::SaslOptions;
 use crate::session::StateReceiver;
 pub use crate::session::{EventType, SessionId, SessionInfo, SessionState, WatchedEvent};
 use crate::tls::TlsOptions;
@@ -1538,6 +1540,8 @@ pub(crate) struct Version(u32, u32, u32);
 #[derive(Clone, Debug)]
 pub struct Connector {
     tls: Option<TlsOptions>,
+    #[cfg(any(feature = "sasl-digest-md5", feature = "sasl-gssapi"))]
+    sasl: Option<SaslOptions>,
     authes: Vec<AuthPacket>,
     session: Option<SessionInfo>,
     readonly: bool,
@@ -1553,6 +1557,8 @@ impl Connector {
     fn new() -> Self {
         Self {
             tls: None,
+            #[cfg(any(feature = "sasl-digest-md5", feature = "sasl-gssapi"))]
+            sasl: None,
             authes: Default::default(),
             session: None,
             readonly: false,
@@ -1624,6 +1630,13 @@ impl Connector {
         self
     }
 
+    /// Specifies SASL options.
+    #[cfg(any(feature = "sasl-digest-md5", feature = "sasl-gssapi"))]
+    pub fn sasl(&mut self, options: impl Into<SaslOptions>) -> &mut Self {
+        self.sasl = Some(options.into());
+        self
+    }
+
     /// Fail session establishment eagerly with [Error::NoHosts] when all hosts has been tried.
     ///
     /// This permits fail-fast without wait up to [Self::session_timeout] in [Self::connect]. This
@@ -1657,6 +1670,8 @@ impl Connector {
             self.readonly,
             self.detached,
             tls_config,
+            #[cfg(any(feature = "sasl-digest-md5", feature = "sasl-gssapi"))]
+            self.sasl.take(),
             self.session_timeout,
             self.connection_timeout,
         );

src/lib.rs

@@ -6,6 +6,8 @@ mod endpoint;
 mod error;
 mod proto;
 mod record;
+#[cfg(any(feature = "sasl-digest-md5", feature = "sasl-gssapi"))]
+mod sasl;
 mod session;
 mod tls;
 mod util;
@@ -14,3 +16,9 @@ pub use self::acl::{Acl, Acls, AuthId, AuthUser, Permission};
 pub use self::error::Error;
 pub use self::tls::TlsOptions;
 pub use crate::client::*;
+#[cfg(feature = "sasl-digest-md5")]
+pub use crate::sasl::DigestMd5SaslOptions;
+#[cfg(feature = "sasl-gssapi")]
+pub use crate::sasl::GssapiSaslOptions;
+#[cfg(any(feature = "sasl-digest-md5", feature = "sasl-gssapi"))]
+pub use crate::sasl::SaslOptions;

src/record.rs

@@ -422,16 +422,44 @@ impl<'a> DeserializableRecord<'a> for &'a [u8] {
     type Error = InsufficientBuf;
 
     fn deserialize(buf: &mut ReadingBuf<'a>) -> Result<&'a [u8], Self::Error> {
+        Option::<&[u8]>::deserialize(buf).map(|opt| opt.unwrap_or_default())
+    }
+}
+
+impl SerializableRecord for Option<&[u8]> {
+    fn serialize(&self, buf: &mut dyn BufMut) {
+        match self {
+            None => buf.put_i32(-1),
+            Some(bytes) => bytes.serialize(buf),
+        }
+    }
+}
+
+impl DynamicRecord for Option<&[u8]> {
+    fn serialized_len(&self) -> usize {
+        match self {
+            None => 4,
+            Some(bytes) => 4 + bytes.len(),
+        }
+    }
+}
+
+impl<'a> DeserializableRecord<'a> for Option<&'a [u8]> {
+    type Error = InsufficientBuf;
+
+    fn deserialize(buf: &mut ReadingBuf<'a>) -> Result<Option<&'a [u8]>, Self::Error> {
         let n = i32::deserialize(buf)?;
-        if n <= 0 {
-            return Ok(Default::default());
+        if n < 0 {
+            return Ok(None);
+        } else if n == 0 {
+            return Ok(Some(Default::default()));
         } else if n > buf.len() as i32 {
             return Err(InsufficientBuf);
         }
         let n = n as usize;
         let bytes = unsafe { buf.get_unchecked(..n) };
         unsafe { *buf = buf.get_unchecked(n..) };
-        Ok(bytes)
+        Ok(Some(bytes))
     }
 }
 

src/sasl/digest_md5.rs

@@ -0,0 +1,70 @@
+use std::borrow::Cow;
+
+use rsasl::callback::{Context, Request, SessionCallback, SessionData};
+use rsasl::prelude::*;
+use rsasl::property::{AuthId, Password, Realm};
+
+use super::{Result, SaslInitiator, SaslInnerOptions, SaslOptions, SaslSession};
+
+#[derive(Clone, Debug)]
+pub struct DigestMd5SaslOptions {
+    realm: Option<Cow<'static, str>>,
+    username: Cow<'static, str>,
+    password: Cow<'static, str>,
+}
+
+impl DigestMd5SaslOptions {
+    fn realm(&self) -> Option<&str> {
+        self.realm.as_ref().map(|s| s.as_ref())
+    }
+
+    pub(crate) fn new(username: impl Into<Cow<'static, str>>, password: impl Into<Cow<'static, str>>) -> Self {
+        Self { realm: None, username: username.into(), password: password.into() }
+    }
+
+    /// Specifies the client chosen realm.
+    #[cfg(test)]
+    pub fn with_realm(self, realm: impl Into<Cow<'static, str>>) -> Self {
+        Self { realm: Some(realm.into()), ..self }
+    }
+}
+
+impl From<DigestMd5SaslOptions> for SaslOptions {
+    fn from(options: DigestMd5SaslOptions) -> Self {
+        Self(SaslInnerOptions::DigestMd5(options))
+    }
+}
+
+struct DigestSessionCallback {
+    options: DigestMd5SaslOptions,
+}
+
+impl SessionCallback for DigestSessionCallback {
+    fn callback(
+        &self,
+        _session_data: &SessionData,
+        _context: &Context,
+        request: &mut Request<'_>,
+    ) -> Result<(), SessionError> {
+        if request.is::<Realm>() {
+            if let Some(realm) = self.options.realm() {
+                request.satisfy::<Realm>(realm)?;
+            }
+        } else if request.is::<AuthId>() {
+            request.satisfy::<AuthId>(&self.options.username)?;
+        } else if request.is::<Password>() {
+            request.satisfy::<Password>(self.options.password.as_bytes())?;
+        }
+        Ok(())
+    }
+}
+
+impl SaslInitiator for DigestMd5SaslOptions {
+    fn new_session(&self, _hostname: &str) -> Result<SaslSession> {
+        let callback = DigestSessionCallback { options: self.clone() };
+        let config = SASLConfig::builder().with_defaults().with_callback(callback).unwrap();
+        let client = SASLClient::new(config);
+        let session = client.start_suggested(&[Mechname::parse(b"DIGEST-MD5").unwrap()]).unwrap();
+        SaslSession::new(session)
+    }
+}

src/sasl/gssapi.rs

@@ -0,0 +1,80 @@
+use std::borrow::Cow;
+
+use rsasl::callback::{Context, Request, SessionCallback, SessionData};
+use rsasl::mechanisms::gssapi::properties::GssService;
+use rsasl::prelude::*;
+use rsasl::property::Hostname;
+
+use super::{Result, SaslInitiator, SaslInnerOptions, SaslOptions, SaslSession};
+
+impl From<GssapiSaslOptions> for SaslOptions {
+    fn from(options: GssapiSaslOptions) -> Self {
+        Self(SaslInnerOptions::Gssapi(options))
+    }
+}
+
+/// GSSAPI SASL options.
+#[derive(Clone, Debug)]
+pub struct GssapiSaslOptions {
+    username: Cow<'static, str>,
+    hostname: Option<Cow<'static, str>>,
+}
+
+impl GssapiSaslOptions {
+    pub(crate) fn new() -> Self {
+        Self { username: Cow::from("zookeeper"), hostname: None }
+    }
+
+    /// Specifies the primary part of Kerberos principal.
+    ///
+    /// It is `zookeeper.sasl.client.username` in Java client, but the word "client" is misleading
+    /// as it is the username of targeting server.
+    ///
+    /// Defaults to "zookeeper".
+    pub fn with_username(self, username: impl Into<Cow<'static, str>>) -> Self {
+        Self { username: username.into(), ..self }
+    }
+
+    /// Specifies the instance part of Kerberos principal.
+    ///
+    /// Defaults to hostname or ip of targeting server in connecting string.
+    pub fn with_hostname(self, hostname: impl Into<Cow<'static, str>>) -> Self {
+        Self { hostname: Some(hostname.into()), ..self }
+    }
+
+    fn hostname_or(&self, hostname: &str) -> Cow<'static, str> {
+        match self.hostname.as_ref() {
+            None => Cow::Owned(hostname.to_string()),
+            Some(hostname) => hostname.clone(),
+        }
+    }
+}
+
+impl SaslInitiator for GssapiSaslOptions {
+    fn new_session(&self, hostname: &str) -> Result<SaslSession> {
+        struct GssapiOptionsProvider {
+            username: Cow<'static, str>,
+            hostname: Cow<'static, str>,
+        }
+        impl SessionCallback for GssapiOptionsProvider {
+            fn callback(
+                &self,
+                _session_data: &SessionData,
+                _context: &Context,
+                request: &mut Request<'_>,
+            ) -> Result<(), SessionError> {
+                if request.is::<Hostname>() {
+                    request.satisfy::<Hostname>(&self.hostname)?;
+                } else if request.is::<GssService>() {
+                    request.satisfy::<GssService>(&self.username)?;
+                }
+                Ok(())
+            }
+        }
+        let provider = GssapiOptionsProvider { username: self.username.clone(), hostname: self.hostname_or(hostname) };
+        let config = SASLConfig::builder().with_defaults().with_callback(provider).unwrap();
+        let client = SASLClient::new(config);
+        let session = client.start_suggested(&[Mechname::parse(b"GSSAPI").unwrap()]).unwrap();
+        SaslSession::new(session)
+    }
+}