Fix sasl auth

.github/workflows/rust.yml

@@ -10,7 +10,7 @@ on:
 
 env:
   CARGO_TERM_COLOR: always
-  RUST_VERSION: 1.80.0
+  RUST_VERSION: 1.88.0
 
 jobs:
   static_analysis:

src/error.rs

@@ -32,11 +32,12 @@ pub enum Error {
     MissingBrokerConfigOptions,
     IncorrectConnectionUsage,
     InvalidSaslMechanism,
+    SaslAuthFailed(String),
 }
 
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{:?}", self)
+        write!(f, "{self:?}")
     }
 }
 

src/lib.rs

@@ -673,7 +673,7 @@ pub mod prelude {
     pub use crate::error::{Error, KafkaCode, Result};
     pub use crate::metadata::ClusterMetadata;
     pub use crate::network::{
-        sasl::{do_sasl, SaslConfig},
+        sasl::{do_sasl_v2, SaslConfig},
         tcp::{SaslTcpConfig, SaslTcpConnection, TcpConnection},
         tls::{SaslTlsConfig, SaslTlsConnection, TlsConnection, TlsConnectionOptions},
         BrokerAddress, BrokerConnection,

src/metadata.rs

@@ -87,7 +87,7 @@ impl<'a, T: BrokerConnection + Clone + Debug> ClusterMetadata<T> {
         Some(leader.node_id)
     }
 
-    #[instrument(name = "metadata-sync")]
+    #[instrument(name = "metadata-sync", level = "debug")]
     pub async fn sync(&mut self) -> Result<()> {
         tracing::debug!("Syncing metadata");
         // let mut set = JoinSet::new();
@@ -187,10 +187,7 @@ impl<'a, T: BrokerConnection + Clone + Debug> ClusterMetadata<T> {
                 // Do we have this topic already?
                 if let Some(existing_partitions) = broker_ownership.get_mut(&new_topic_name) {
                     // Don't push the partition on more than once
-                    if !existing_partitions
-                        .iter()
-                        .any(|existing_partition| *existing_partition == *new_partition)
-                    {
+                    if !existing_partitions.contains(new_partition) {
                         existing_partitions.push(*new_partition);
                     }
                 } else {

src/network/sasl.rs

@@ -10,7 +10,7 @@ use crate::{
 };
 use bytes::Bytes;
 use rsasl::prelude::*;
-use std::io::Cursor;
+use std::{collections::HashSet, future::Future, io::Cursor};
 
 /// SASL Credentials
 #[derive(Clone, Debug)]
@@ -62,63 +62,198 @@ pub async fn sasl_authentication(
     SaslAuthenticationResponse::try_from(authentication_response.freeze())
 }
 
-pub async fn do_sasl(
+pub async fn start_sasl_session(
     broker_conn: impl BrokerConnection + Clone,
     correlation_id: i32,
     client_id: &str,
     config: SaslConfig,
-) -> Result<()> {
-    let mechanism = String::from("SCRAM-SHA-256");
-    let handshake_response =
-        sasl_handshake(broker_conn.clone(), correlation_id, client_id, mechanism).await?;
-    if handshake_response.error_code != KafkaCode::None {
-        return Err(Error::KafkaError(handshake_response.error_code));
+    mechanism: String,
+) -> Result<(Option<Session>, Vec<String>)> {
+    let config = SASLConfig::with_credentials(None, config.username, config.password).unwrap();
+    let mut maybe_session = None;
+    let mut maybe_next = vec![];
+
+    let handshake_response = sasl_handshake(
+        broker_conn.clone(),
+        correlation_id,
+        client_id,
+        mechanism.clone(),
+    )
+    .await?;
+
+    let suggested = parse_mechanisms(&handshake_response)?;
+
+    match handshake_response.error_code {
+        KafkaCode::None => {
+            tracing::trace!("good handshake, starting a session");
+
+            let sasl = rsasl::prelude::SASLClient::new(config.clone());
+            maybe_session = Some(sasl.start_suggested(&suggested).map_err(|e| {
+                tracing::error!("{:?}", e);
+                Error::InvalidSaslMechanism
+            })?);
+        }
+        KafkaCode::UnsupportedSaslMechanism => {
+            tracing::debug!(
+                "bad handshake: {:?} not supported; will try others",
+                mechanism
+            );
+            tracing::debug!("supported mechanisms {:?}", suggested);
+            assert!(!suggested.is_empty());
+
+            maybe_next.extend(suggested.iter().map(|x| String::from(x.as_str())));
+        }
+        _ => {
+            tracing::error!("handshake failed with: {:?}", handshake_response);
+            return Err(Error::KafkaError(handshake_response.error_code));
+        }
     }
 
-    let config = SASLConfig::with_credentials(None, config.username, config.password).unwrap();
-    let sasl = rsasl::prelude::SASLClient::new(config);
-    let mechanisms = handshake_response
+    Ok((maybe_session, maybe_next))
+}
+
+fn parse_mechanisms(handshake_response: &SaslHandshakeResponse) -> Result<Vec<&Mechname>> {
+    handshake_response
         .mechanisms
         .iter()
         .map(|mech| {
             tracing::debug!("{:?}", mech);
-            Mechname::parse(mech).map_err(|e| {
+            Mechname::parse(mech.as_ref()).map_err(|e| {
                 tracing::error!("{:?}", e);
                 Error::InvalidSaslMechanism
             })
         })
-        .collect::<Result<Vec<&Mechname>>>()?;
-    tracing::debug!("mechanisms {:?}", mechanisms);
+        .collect::<Result<Vec<&Mechname>>>()
+}
+
+//
+// We need the factory because kafka drops our connection
+// if the given mechanism is unsupported.
+// In this case we retry with another mechanism, suggested by kafka.
+// rsasl supports this kind of workflow.
+//
+pub async fn do_sasl_v2<T, FT, FFT>(
+    conn_factory: FFT,
+    correlation_id: i32,
+    client_id: &str,
+    config: SaslConfig,
+) -> Result<T>
+where
+    T: BrokerConnection + Clone,
+    FT: Future<Output = Result<T>>,
+    FFT: Fn() -> FT,
+{
+    let mut tried = HashSet::new();
+
+    let mut maybe_session;
+    // Start with the given mechanism.
+    // Server will respond with a list of supported ones.
+    let mut next_mech = Some(String::from("SCRAM-SHA-256"));
+
+    let mut broker_conn;
+
+    loop {
+        tried.insert(next_mech.clone().unwrap());
+        broker_conn = conn_factory().await?;
+
+        let suggested;
+        (maybe_session, suggested) = start_sasl_session(
+            broker_conn.clone(),
+            correlation_id,
+            client_id,
+            config.clone(),
+            next_mech.clone().unwrap(),
+        )
+        .await?;
 
-    let mut session = sasl.start_suggested(&mechanisms).unwrap();
-    let selected_mechanism = session.get_mechname();
-    tracing::debug!("Using {:?} for our SASL Mechanism", selected_mechanism);
+        next_mech = HashSet::from_iter(suggested)
+            .difference(&tried)
+            .map(|x| x.to_owned())
+            .next();
+
+        if maybe_session.is_some() || next_mech.is_none() {
+            break;
+        }
+        tracing::debug!("will retry with mechanism {:?}", next_mech);
+    }
 
-    let mut data: Option<Vec<u8>> = None;
+    if maybe_session.is_none() {
+        tracing::error!("failed to start a sasl session");
+        return Err(Error::InvalidSaslMechanism);
+    }
+
+    let session = maybe_session.unwrap();
+    tracing::info!("Using {:?} for our SASL Mechanism", session.get_mechname());
+
+    match do_sasl_chit_chat(session, &broker_conn, correlation_id, client_id).await {
+        Ok(_) => Ok(broker_conn),
+        Err(e) => {
+            tracing::error!("chit_chat failed: {:?}", e);
+            Err(e)
+        }
+    }
+}
+
+async fn do_sasl_chit_chat<T>(
+    mut session: Session,
+    broker_conn: &T,
+    correlation_id: i32,
+    client_id: &str,
+) -> Result<()>
+where
+    T: BrokerConnection + Clone,
+{
+    let mut data_in: Option<Vec<u8>> = None;
+
+    tracing::trace!("start sasl chit-chat");
 
     // stepping the authentication exchange to completion
-    while {
+    loop {
         let mut out = Cursor::new(Vec::new());
+
         // each call to step writes the generated auth data into the provided writer.
-        // Normally this data would then have to be sent to the other party, but this goes
-        // beyond the scope of this example
         let state = session
-            .step(data.as_deref(), &mut out)
+            .step(data_in.as_deref(), &mut out)
             .expect("step errored!");
 
-        data = Some(out.into_inner());
+        let data_out = out.into_inner();
+
+        tracing::trace!("outgoing: {:?}", data_out);
 
-        // returns `true` if step needs to be called again with another batch of data
-        state.is_running()
-    } {
-        let authentication_response = sasl_authentication(
+        let response = sasl_authentication(
             broker_conn.clone(),
             correlation_id,
             client_id,
-            Bytes::from(data.unwrap()),
+            Bytes::from(data_out),
         )
         .await?;
-        data = Some(authentication_response.auth_bytes.to_vec());
+
+        tracing::trace!("incoming: {:?}", response);
+
+        match response.error_code {
+            KafkaCode::None => {
+                let auth_bytes = response.auth_bytes.to_vec();
+
+                data_in = if !auth_bytes.is_empty() {
+                    Some(auth_bytes)
+                } else {
+                    None
+                };
+            }
+            KafkaCode::SaslAuthenticationFailed => {
+                let msg = response
+                    .error_message
+                    .map(|x| String::from_utf8_lossy(&x).into_owned())
+                    .unwrap_or("".to_owned());
+                tracing::info!("auth failed: {:?}: {:?}", response.error_code, msg);
+                return Err(Error::SaslAuthFailed(msg));
+            }
+            _ => return Err(Error::KafkaError(response.error_code)),
+        }
+
+        if data_in.is_none() && state.is_finished() {
+            break;
+        }
     }
 
     Ok(())

src/network/tcp.rs

@@ -12,7 +12,7 @@ use crate::{
     error::{Error, Result},
 };
 
-use super::sasl::{do_sasl, SaslConfig};
+use super::sasl::{do_sasl_v2, SaslConfig};
 use super::{BrokerAddress, BrokerConnection};
 
 /// TCP connection to a Kafka/Redpanda broker.
@@ -84,6 +84,10 @@ impl TcpConnection {
             // Try to read data, this may still fail with `WouldBlock`
             // if the readiness event is a false positive.
             match self.stream.try_read(&mut buf[index..]) {
+                Ok(0) => {
+                    tracing::info!("Empty read: connection was closed by server");
+                    return Err(Error::MissingData("Connection closed".to_owned()));
+                }
                 Ok(n) => {
                     index += n;
                     tracing::trace!("Read {} bytes", n);
@@ -239,27 +243,24 @@ impl BrokerConnection for SaslTcpConnection {
         self.tcp_conn.receive_response_().await
     }
 
+    #[instrument(name = "sasl-new", level = "trace")]
     async fn new(p: Self::ConnConfig) -> Result<Self> {
-        let conn = TcpConnection::new_(p.tcp_config).await?;
-        do_sasl(
-            conn.clone(),
+        let conn = do_sasl_v2(
+            async || TcpConnection::new_(p.tcp_config.clone()).await,
             p.sasl_config.correlation_id,
             &p.sasl_config.client_id,
             p.sasl_config.clone(),
         )
         .await?;
+
         Ok(Self { tcp_conn: conn })
     }
 
     async fn from_addr(p: Self::ConnConfig, addr: BrokerAddress) -> Result<Self> {
-        let conn = TcpConnection::new_(vec![addr]).await?;
-        do_sasl(
-            conn.clone(),
-            p.sasl_config.correlation_id,
-            &p.sasl_config.client_id,
-            p.sasl_config.clone(),
-        )
-        .await?;
-        Ok(Self { tcp_conn: conn })
+        let tc = Self::ConnConfig {
+            tcp_config: vec![addr],
+            ..p
+        };
+        return Self::new(tc).await;
     }
 }

src/network/tls.rs

@@ -19,7 +19,7 @@ use crate::{
     error::{Error, Result},
 };
 
-use super::sasl::do_sasl;
+use super::sasl::do_sasl_v2;
 use super::sasl::SaslConfig;
 use super::{BrokerAddress, BrokerConnection};
 
@@ -296,9 +296,8 @@ impl BrokerConnection for SaslTlsConnection {
 
     /// Connect to a Kafka/Redpanda broker
     async fn new(p: Self::ConnConfig) -> Result<Self> {
-        let conn = TlsConnection::new_(p.tls_config).await?;
-        do_sasl(
-            conn.clone(),
+        let conn = do_sasl_v2(
+            async || TlsConnection::new_(p.tls_config.clone()).await,
             p.sasl_config.correlation_id,
             &p.sasl_config.client_id,
             p.sasl_config.clone(),
@@ -308,22 +307,14 @@ impl BrokerConnection for SaslTlsConnection {
     }
 
     async fn from_addr(p: Self::ConnConfig, addr: BrokerAddress) -> Result<Self> {
-        let cafile = p.tls_config.cafile.clone();
-
         let options = TlsConnectionOptions {
             broker_options: vec![addr],
-            cert: p.tls_config.cert,
-            key: p.tls_config.key,
-            cafile,
+            ..p.tls_config
         };
-        let conn = TlsConnection::new_(options).await?;
-        do_sasl(
-            conn.clone(),
-            p.sasl_config.correlation_id,
-            &p.sasl_config.client_id,
-            p.sasl_config.clone(),
-        )
-        .await?;
-        Ok(Self { tls_conn: conn })
+        let tc = Self::ConnConfig {
+            tls_config: options,
+            ..p
+        };
+        return Self::new(tc).await;
     }
 }