Repository: robustmq/robustmq (language: rust)

Issue:
RBIP-03: MQTT Broker supports the Auth command word in MQTT5
An AUTH packet is sent from Client to Server or Server to Client as part of an extended authentication exchange, such as challenge / response authentication. It is a Protocol Error for the Client or Server to send an AUTH packet if the CONNECT packet did not contain the same Authentication Method.


Discussion / hints:

Fix patch:
diff --git a/src/protocol/src/mqtt/codec.rs b/src/protocol/src/mqtt/codec.rs
index 0db04d1a3..c429e28e8 100644
--- a/src/protocol/src/mqtt/codec.rs
+++ b/src/protocol/src/mqtt/codec.rs
@@ -314,6 +314,9 @@ impl MqttCodec {
                         buffer,
                     )?
                 }
+                MqttPacket::Auth(auth, auth_properties) => {
+                    crate::mqtt::mqttv5::auth::write(&auth, &auth_properties, buffer)?
+                }
             };
         }
         Ok(())
@@ -433,6 +436,9 @@ fn calc_mqtt_packet_len(
                     &mut buffer,
                 )?
             }
+            MqttPacket::Auth(auth, auth_properties) => {
+                crate::mqtt::mqttv5::auth::write(&auth, &auth_properties, &mut buffer)?
+            }
         };
     }
     Ok(size)
@@ -454,6 +460,7 @@ pub fn parse_mqtt_packet_to_name(packet: MqttPacket) -> String {
         MqttPacket::PingReq(_) => "ping",
         MqttPacket::PingResp(_) => "pong",
         MqttPacket::Disconnect(_, _) => "disconnect",
+        MqttPacket::Auth(_, _) => "auth",
     };
     name.to_string()
 }
diff --git a/src/protocol/src/mqtt/common.rs b/src/protocol/src/mqtt/common.rs
index 36135f098..2b560919a 100644
--- a/src/protocol/src/mqtt/common.rs
+++ b/src/protocol/src/mqtt/common.rs
@@ -89,6 +89,7 @@ pub enum PacketType {
     PingReq,
     PingResp,
     Disconnect,
+    Auth,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -114,6 +115,7 @@ pub enum MqttPacket {
     Unsubscribe(Unsubscribe, Option<UnsubscribeProperties>),
     UnsubAck(UnsubAck, Option<UnsubAckProperties>),
     Disconnect(Disconnect, Option<DisconnectProperties>),
+    Auth(Auth, Option<AuthProperties>),
 }
 
 /// Packet type from a byte
@@ -1047,6 +1049,53 @@ pub struct DisconnectProperties {
     pub server_reference: Option<String>,
 }
 
+//--------------------------- AUTH packet -------------------------------
+
+/// AUTH packet is for authentication, but only used in MQTT 5.0
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Auth {
+    pub reason: Option<AuthReason>,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[repr(u8)]
+pub enum AuthReason {
+    Success,
+    ContinueAuthentication,
+    ReAuthenticate,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Default)]
+pub struct AuthProperties {
+    /// authentication method
+    pub authentication_method: Option<String>,
+
+    /// authentication data
+    pub authentication_data: Option<Bytes>,
+
+    /// Human readable reason for the disconnect
+    pub reason_string: Option<String>,
+
+    /// List of user properties
+    pub user_properties: Vec<(String, String)>,
+}
+
+impl fmt::Display for Auth {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "reason_code:{:?}", self.reason)
+    }
+}
+
+impl fmt::Display for AuthProperties {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "authentication_method:{:?}, authentication_data:{:?},reason_string:{:?}, user_properties:{:?}",
+            self.authentication_method, self.authentication_data, self.reason_string, self.user_properties
+        )
+    }
+}
+
 /// Error during serialization and deserialization
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
diff --git a/src/protocol/src/mqtt/mqttv4/codec.rs b/src/protocol/src/mqtt/mqttv4/codec.rs
index 9e90ea5be..f705aa924 100644
--- a/src/protocol/src/mqtt/mqttv4/codec.rs
+++ b/src/protocol/src/mqtt/mqttv4/codec.rs
@@ -101,6 +101,7 @@ impl codec::Decoder for Mqtt4Codec {
             // MQTT V4 Disconnect packet gets handled in the previous check, this branch gets
             // hit when Disconnect packet has properties which are only valid for MQTT V5
             PacketType::Disconnect => return Err(Error::InvalidProtocol),
+            PacketType::Auth => return Err(Error::InvalidProtocol),
         };
         Ok(Some(packet))
     }
diff --git a/src/protocol/src/mqtt/mqttv5/auth.rs b/src/protocol/src/mqtt/mqttv5/auth.rs
index 2f907a43e..36957bb76 100644
--- a/src/protocol/src/mqtt/mqttv5/auth.rs
+++ b/src/protocol/src/mqtt/mqttv5/auth.rs
@@ -12,3 +12,279 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use super::*;
+
+pub fn len(auth: &Auth, properties: &Option<AuthProperties>) -> usize {
+    // The Reason Code and Property Length can be ommited if the Reason Code is 0x00(Success)
+    // and there are no properties. In this case the AUTH packet has a remaining length of 2.
+    // <https://docs.oasis-open.org/mqtt/mqtt/v5.0/os/mqtt-v5.0-os.html#_Toc3901217>
+    if auth.reason.unwrap() == AuthReason::Success && properties.is_none() {
+        return 2; // Packet type + 0x00
+    }
+
+    let mut len = 0;
+    if let Some(p) = properties {
+        let properties_len = properties::len(p);
+        let properties_len_len = len_len(properties_len);
+        len += properties_len_len + properties_len;
+    } else {
+        // just 1 byte representing 0 len
+        len += 1;
+    }
+
+    len
+}
+
+pub fn write(
+    auth: &Auth,
+    properties: &Option<AuthProperties>,
+    buffer: &mut BytesMut,
+) -> Result<usize, Error> {
+    let len = len(auth, properties);
+    buffer.put_u8(0b1111_0000);
+
+    if len == 2 {
+        buffer.put_u8(0x00); // reason code Success (0x00)
+        return Ok(len);
+    }
+    let count = write_remaining_length(buffer, len)?;
+
+    buffer.put_u8(code(auth.reason.unwrap()));
+
+    if let Some(p) = &properties {
+        properties::write(p, buffer)?;
+    } else {
+        write_remaining_length(buffer, 0)?;
+    }
+
+    Ok(1 + count + len)
+}
+
+pub fn read(
+    fixed_header: FixedHeader,
+    mut bytes: Bytes,
+) -> Result<(Auth, Option<AuthProperties>), Error> {
+    let packet_type = fixed_header.byte1 >> 4;
+    let flags = fixed_header.byte1 & 0b0000_1111;
+
+    bytes.advance(fixed_header.fixed_header_len);
+
+    if packet_type != PacketType::Auth as u8 {
+        return Err(Error::InvalidPacketType(packet_type));
+    };
+
+    if flags != 0x00 {
+        return Err(Error::MalformedPacket);
+    };
+
+    if fixed_header.remaining_len == 0 {
+        return Ok((
+            Auth {
+                reason: Some(AuthReason::Success),
+            },
+            None,
+        ));
+    }
+
+    let reason_code = read_u8(&mut bytes)?;
+
+    let auth = Auth {
+        reason: Some(reason(reason_code)?),
+    };
+    let properties = properties::read(&mut bytes)?;
+
+    Ok((auth, properties))
+}
+
+mod properties {
+
+    use super::*;
+
+    pub fn len(properties: &AuthProperties) -> usize {
+        let mut len = 0;
+
+        if let Some(authentication_method) = &properties.authentication_method {
+            len += 1 + 2 + authentication_method.len();
+        }
+
+        if let Some(authentication_data) = &properties.authentication_data {
+            len += 1 + 2 + authentication_data.len();
+        }
+
+        if let Some(reason) = &properties.reason_string {
+            len += 1 + 2 + reason.len();
+        }
+
+        for (key, value) in properties.user_properties.iter() {
+            len += 1 + 2 + key.len() + 2 + value.len();
+        }
+
+        len
+    }
+
+    pub fn write(properties: &AuthProperties, buffer: &mut BytesMut) -> Result<(), Error> {
+        let len = len(properties);
+        write_remaining_length(buffer, len)?;
+
+        // write the authentication method
+        if let Some(authentication_method) = &properties.authentication_method {
+            // write the identifier of authentication method 21(0x15) - one byte
+            buffer.put_u8(PropertyType::AuthenticationMethod as u8);
+            // write the content of authentication method
+            write_mqtt_string(buffer, authentication_method);
+        }
+
+        if let Some(authentication_data) = &properties.authentication_data {
+            // write the identifier of authentication data 22(0x16) - one byte
+            buffer.put_u8(PropertyType::AuthenticationData as u8);
+            // write the content of authentication data
+            write_mqtt_bytes(buffer, authentication_data);
+        }
+
+        if let Some(reason) = &properties.reason_string {
+            // write the identifier of reason string 31(0x16) - one byte
+            buffer.put_u8(PropertyType::ReasonString as u8);
+            // write the content of reason string
+            write_mqtt_string(buffer, reason);
+        }
+
+        for (key, value) in properties.user_properties.iter() {
+            // write the identifier of user properties 38(0x26) - one byte
+            buffer.put_u8(PropertyType::UserProperty as u8);
+            // write the key of user properties
+            write_mqtt_string(buffer, key);
+            // write the value of user properties
+            write_mqtt_string(buffer, value);
+        }
+
+        Ok(())
+    }
+
+    pub fn read(bytes: &mut Bytes) -> Result<Option<AuthProperties>, Error> {
+        let mut authentication_method = None;
+        let mut authentication_data = None;
+        let mut reason_string = None;
+        let mut user_properties = Vec::new();
+
+        let (properties_len_len, properties_len) = length(bytes.iter())?;
+        bytes.advance(properties_len_len);
+        if properties_len == 0 {
+            return Ok(None);
+        }
+
+        let mut cursor = 0;
+        // read until cursor reaches property length. properties_len = 0 will skip this loop
+        while cursor < properties_len {
+            let prop = read_u8(bytes)?;
+            cursor += 1;
+
+            match property(prop)? {
+                PropertyType::AuthenticationMethod => {
+                    let method = read_mqtt_string(bytes)?;
+                    cursor += 2 + method.len();
+                    authentication_method = Some(method);
+                }
+
+                PropertyType::AuthenticationData => {
+                    let data = read_mqtt_bytes(bytes)?;
+                    cursor += 2 + data.len();
+                    authentication_data = Some(data);
+                }
+
+                PropertyType::ReasonString => {
+                    let reason = read_mqtt_string(bytes)?;
+                    cursor += 2 + reason.len();
+                    reason_string = Some(reason);
+                }
+
+                PropertyType::UserProperty => {
+                    let key = read_mqtt_string(bytes)?;
+                    let value = read_mqtt_string(bytes)?;
+                    cursor += 2 + key.len() + 2 + value.len();
+                    user_properties.push((key, value));
+                }
+
+                _ => return Err(Error::InvalidPropertyType(prop)),
+            }
+        }
+
+        Ok(Some(AuthProperties {
+            authentication_method,
+            authentication_data,
+            reason_string,
+            user_properties,
+        }))
+    }
+}
+
+fn code(reason: AuthReason) -> u8 {
+    match reason {
+        AuthReason::Success => 0x00,
+        AuthReason::ContinueAuthentication => 0x18,
+        AuthReason::ReAuthenticate => 0x19,
+    }
+}
+
+fn reason(code: u8) -> Result<AuthReason, Error> {
+    let v = match code {
+        0x00 => AuthReason::Success,
+        0x18 => AuthReason::ContinueAuthentication,
+        0x19 => AuthReason::ReAuthenticate,
+        other => return Err(Error::InvalidConnectReturnCode(other)),
+    };
+    Ok(v)
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_auth_v5() {
+        use super::*;
+
+        let mut buffer = BytesMut::new();
+        let auth = Auth {
+            reason: Some(AuthReason::ContinueAuthentication),
+        };
+
+        let authentication_method: String = "SCRAM-SHA-256".to_string();
+        let authentication_data: Bytes = Bytes::from("client-first-data");
+        let user_perperties: Vec<(String, String)> = vec![
+            ("username".into(), "Justin".into()),
+            ("tag".to_string(), "middleware".to_string()),
+        ];
+
+        let properties = AuthProperties {
+            authentication_method: Some(authentication_method),
+            authentication_data: Some(authentication_data),
+            reason_string: Some(String::from("Not Success")),
+            user_properties: user_perperties,
+        };
+
+        // test the write function of Auth in v5
+        write(&auth, &Some(properties), &mut buffer).unwrap();
+
+        // test the fixed_header part
+        let fixed_header: FixedHeader = parse_fixed_header(buffer.iter()).unwrap();
+        assert_eq!(fixed_header.byte1, 0b1111_0000);
+        assert_eq!(fixed_header.fixed_header_len, 2);
+        assert_eq!(fixed_header.remaining_len, 88);
+
+        // test the read function of pubrec packet and check the result of write function in MQTT v5
+        let (x, y) = read(fixed_header, buffer.copy_to_bytes(buffer.len())).unwrap();
+        assert_eq!(x.reason.unwrap(), AuthReason::ContinueAuthentication);
+
+        let auth_properties = y.unwrap();
+        assert_eq!(
+            auth_properties.reason_string,
+            Some("Not Success".to_string())
+        );
+        assert_eq!(
+            auth_properties.user_properties.first(),
+            Some(&("username".to_string(), "Justin".to_string()))
+        );
+
+        // test display of puback and puback_properties in v5
+        println!("auth is {}", auth);
+        println!("auth_properties is {}", auth_properties);
+    }
+}
diff --git a/src/protocol/src/mqtt/mqttv5/codec.rs b/src/protocol/src/mqtt/mqttv5/codec.rs
index c85f05ccc..0c19e82ec 100644
--- a/src/protocol/src/mqtt/mqttv5/codec.rs
+++ b/src/protocol/src/mqtt/mqttv5/codec.rs
@@ -16,8 +16,8 @@ use bytes::BytesMut;
 use tokio_util::codec;
 
 use super::{
-    check, connack, connect, disconnect, ping, puback, pubcomp, publish, pubrec, pubrel, suback,
-    subscribe, unsuback, unsubscribe, ConnectReadOutcome, MqttPacket, PacketType,
+    auth, check, connack, connect, disconnect, ping, puback, pubcomp, publish, pubrec, pubrel,
+    suback, subscribe, unsuback, unsubscribe, ConnectReadOutcome, MqttPacket, PacketType,
 };
 
 #[derive(Clone, Debug)]
@@ -84,6 +84,9 @@ impl codec::Encoder<MqttPacket> for Mqtt5Codec {
             MqttPacket::Disconnect(disconnect, disconnect_properties) => {
                 disconnect::write(&disconnect, &disconnect_properties, buffer)?
             }
+            MqttPacket::Auth(auth, auth_properties) => {
+                auth::write(&auth, &auth_properties, buffer)?
+            }
         };
         Ok(())
     }
@@ -166,6 +169,10 @@ impl codec::Decoder for Mqtt5Codec {
                 let (disconnect, disconnect_properties) = disconnect::read(fixed_header, packet)?;
                 MqttPacket::Disconnect(disconnect, disconnect_properties)
             }
+            PacketType::Auth => {
+                let (auth, auth_properties) = auth::read(fixed_header, packet)?;
+                MqttPacket::Auth(auth, auth_properties)
+            }
         };
         Ok(Some(packet))
     }
diff --git a/src/protocol/src/mqtt/mqttv5/mod.rs b/src/protocol/src/mqtt/mqttv5/mod.rs
index 7b941055e..5fdf89174 100644
--- a/src/protocol/src/mqtt/mqttv5/mod.rs
+++ b/src/protocol/src/mqtt/mqttv5/mod.rs
@@ -16,6 +16,7 @@ use bytes::{Buf, BufMut, Bytes, BytesMut};
 
 use crate::mqtt::common::*;
 
+pub mod auth;
 pub mod codec;
 pub mod connack;
 pub mod connect;


Test patch: