Skip to main content

rocketmq_auth/authentication/builder/
default_authentication_context_builder.rs

1// Copyright 2023 The RocketMQ Rust Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Default implementation of authentication context builder.
16//!
17//! This module provides the standard implementation for building authentication contexts
18//! from both gRPC and Remoting protocol requests.
19
20use std::collections::BTreeMap;
21
22use base64::Engine;
23use cheetah_string::CheetahString;
24use rocketmq_common::common::mix_all::UNIQUE_MSG_QUERY_FLAG;
25use rocketmq_common::common::mq_version::RocketMqVersion;
26use rocketmq_error::AuthError;
27use rocketmq_remoting::protocol::remoting_command::RemotingCommand;
28#[cfg(feature = "grpc")]
29use tracing::warn;
30
31use crate::authentication::builder::AuthenticationContextBuilder;
32use crate::authentication::context::default_authentication_context::DefaultAuthenticationContext;
33
34/// Constants for parsing authentication headers
35const CREDENTIAL: &str = "Credential";
36const SIGNATURE: &str = "Signature";
37const ACCESS_KEY: &str = "AccessKey";
38const SIGNATURE_FIELD: &str = "Signature";
39
40/// Separator constants
41const SPACE: char = ' ';
42const COMMA: char = ',';
43const EQUAL: char = '=';
44const SLASH: char = '/';
45
46/// Default authentication context builder.
47///
48/// This builder implements the standard RocketMQ authentication context construction
49/// logic, supporting both gRPC (with Authorization header parsing) and Remoting protocol
50/// (with extFields parsing) request types.
51///
52/// # Authentication Format
53///
54/// ## gRPC Authorization Header Format:
55/// ```text
56/// MQv2-HMAC-SHA1 Credential=<username>[/<additional>], SignedHeaders=x-mq-date-time, Signature=<hex-signature>
57/// ```
58///
59/// ## Remoting Protocol Format:
60/// ExtFields containing:
61/// - `AccessKey`: The username/access key
62/// - `Signature`: Base64-encoded signature
63/// - Other fields: Combined and sorted for signature verification
64///
65/// # Examples
66///
67/// ```rust,ignore
68/// use rocketmq_auth::authentication::builder::DefaultAuthenticationContextBuilder;
69/// use rocketmq_auth::authentication::builder::AuthenticationContextBuilder;
70///
71/// let builder = DefaultAuthenticationContextBuilder::new();
72///
73/// // From remoting command
74/// let context = builder.build_from_remoting(&command, Some("channel-123"))?;
75/// assert_eq!(context.username(), Some(&CheetahString::from("testuser")));
76/// ```
77#[derive(Debug, Default, Clone, Copy)]
78pub struct DefaultAuthenticationContextBuilder;
79
80impl DefaultAuthenticationContextBuilder {
81    /// Create a new default authentication context builder.
82    pub fn new() -> Self {
83        Self
84    }
85
86    /// Convert hex-encoded string to base64-encoded string.
87    ///
88    /// This is used to convert the hex signature from gRPC Authorization header
89    /// to the base64 format expected by the authentication system.
90    ///
91    /// # Arguments
92    ///
93    /// * `hex_input` - Hex-encoded string
94    ///
95    /// # Returns
96    ///
97    /// * `Ok(String)` - Base64-encoded string
98    /// * `Err(AuthError)` - Invalid hex input
99    fn hex_to_base64(hex_input: &str) -> Result<String, AuthError> {
100        let bytes = hex::decode(hex_input)
101            .map_err(|e| AuthError::InvalidHexSignature(format!("Failed to decode hex: {}", e)))?;
102        Ok(base64::engine::general_purpose::STANDARD.encode(&bytes))
103    }
104
105    /// Combine request content for signature verification (Remoting protocol).
106    ///
107    /// This method combines the RemotingCommand fields into a single byte array
108    /// for signature verification.
109    ///
110    /// # Arguments
111    ///
112    /// * `request` - The remoting command
113    /// * `sorted_fields` - Sorted field map (excluding Signature)
114    ///
115    /// # Returns
116    ///
117    /// Byte vector containing the combined content
118    fn combine_request_content(
119        request: &RemotingCommand,
120        sorted_fields: &BTreeMap<CheetahString, CheetahString>,
121    ) -> Vec<u8> {
122        let mut content = Vec::new();
123
124        // Add all sorted field values
125        for (key, value) in sorted_fields {
126            content.extend_from_slice(key.as_bytes());
127            content.extend_from_slice(value.as_bytes());
128        }
129
130        // Add request body if present
131        if let Some(body) = request.body() {
132            content.extend_from_slice(body);
133        }
134
135        content
136    }
137}
138
139impl AuthenticationContextBuilder<DefaultAuthenticationContext> for DefaultAuthenticationContextBuilder {
140    #[cfg(feature = "grpc")]
141    fn build_from_grpc(
142        &self,
143        metadata: &tonic::metadata::MetadataMap,
144        request: &dyn std::any::Any,
145    ) -> Result<DefaultAuthenticationContext, AuthError> {
146        let mut context = DefaultAuthenticationContext::new();
147
148        // Extract channel ID from metadata
149        if let Some(channel_id) = metadata.get("channel-id") {
150            if let Ok(channel_id_str) = channel_id.to_str() {
151                context.base.set_channel_id(Some(CheetahString::from(channel_id_str)));
152            }
153        }
154
155        // Set RPC code from request type name
156        let rpc_code = std::any::type_name_of_val(request);
157        context.base.set_rpc_code(Some(CheetahString::from(rpc_code)));
158
159        // Extract Authorization header
160        let authorization = metadata
161            .get("authorization")
162            .and_then(|v| v.to_str().ok())
163            .unwrap_or("");
164
165        if authorization.is_empty() {
166            return Ok(context);
167        }
168
169        // Extract DateTime header (required if authorization present)
170        let datetime = metadata
171            .get("x-mq-date-time")
172            .and_then(|v| v.to_str().ok())
173            .ok_or_else(|| AuthError::MissingDateTime("datetime header is required".into()))?;
174
175        // Parse authorization header: "MQv2-HMAC-SHA1 Credential=xxx, Signature=yyy"
176        let parts: Vec<&str> = authorization.splitn(2, SPACE).collect();
177        if parts.len() != 2 {
178            return Err(AuthError::InvalidAuthorizationHeader(
179                "authentication header format is incorrect".into(),
180            ));
181        }
182
183        // Parse key-value pairs: "Credential=xxx, Signature=yyy"
184        let key_values: Vec<&str> = parts[1].split(COMMA).collect();
185        for key_value in key_values {
186            let kv: Vec<&str> = key_value.trim().splitn(2, EQUAL).collect();
187            if kv.len() != 2 {
188                warn!("Skipping invalid authentication key-value pair: {}", key_value);
189                continue;
190            }
191
192            let auth_item = kv[0];
193            let value = kv[1];
194
195            match auth_item {
196                CREDENTIAL => {
197                    // Extract username from credential: "username/region/service" or just
198                    // "username"
199                    let credential_parts: Vec<&str> = value.split(SLASH).collect();
200                    if credential_parts.is_empty() {
201                        return Err(AuthError::InvalidCredential("credential is empty".into()));
202                    }
203                    context.set_username(CheetahString::from(credential_parts[0]));
204                }
205                SIGNATURE => {
206                    // Convert hex signature to base64
207                    let base64_signature = Self::hex_to_base64(value)?;
208                    context.set_signature(CheetahString::from(base64_signature));
209                }
210                _ => {
211                    // Ignore other fields like "SignedHeaders"
212                }
213            }
214        }
215
216        // Set content to datetime bytes for signature verification
217        context.set_content(datetime.as_bytes().to_vec());
218
219        Ok(context)
220    }
221
222    fn build_from_remoting(
223        &self,
224        request: &RemotingCommand,
225        channel_id: Option<&str>,
226    ) -> Result<DefaultAuthenticationContext, AuthError> {
227        let mut context = DefaultAuthenticationContext::new();
228
229        // Set channel ID
230        if let Some(id) = channel_id {
231            context.base.set_channel_id(Some(CheetahString::from(id)));
232        }
233
234        // Set RPC code from request code
235        context
236            .base
237            .set_rpc_code(Some(CheetahString::from(request.code().to_string())));
238
239        // Get extFields
240        let fields = match request.ext_fields() {
241            Some(fields) => fields,
242            None => return Ok(context), // No auth fields, return empty context
243        };
244
245        // Check for AccessKey
246        let access_key = match fields.get(&CheetahString::from(ACCESS_KEY)) {
247            Some(key) => key,
248            None => return Ok(context), // No AccessKey, return empty context
249        };
250
251        // Set username
252        context.set_username(access_key.clone());
253
254        // Set signature if present
255        if let Some(signature) = fields.get(&CheetahString::from(SIGNATURE_FIELD)) {
256            context.set_signature(signature.clone());
257        }
258
259        // Build sorted content map for signature verification
260        let mut sorted_map = BTreeMap::new();
261
262        for (key, value) in fields.iter() {
263            // Filter out UNIQUE_MSG_QUERY_FLAG for old versions
264            if request.version() <= RocketMqVersion::V4_9_3.ordinal() as i32 && key.as_str() == UNIQUE_MSG_QUERY_FLAG {
265                continue;
266            }
267
268            // Exclude Signature field from content
269            if key.as_str() != SIGNATURE_FIELD {
270                sorted_map.insert(key.clone(), value.clone());
271            }
272        }
273
274        // Combine request content for signature verification
275        let content = Self::combine_request_content(request, &sorted_map);
276        context.set_content(content);
277
278        Ok(context)
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285
286    #[test]
287    fn test_hex_to_base64() {
288        let hex = "48656C6C6F";
289        let result = DefaultAuthenticationContextBuilder::hex_to_base64(hex).unwrap();
290        assert_eq!(result, "SGVsbG8=");
291    }
292
293    #[test]
294    fn test_hex_to_base64_invalid() {
295        let hex = "ZZZZ";
296        let result = DefaultAuthenticationContextBuilder::hex_to_base64(hex);
297        assert!(result.is_err());
298    }
299
300    #[test]
301    fn test_build_from_remoting_no_ext_fields() {
302        use rocketmq_remoting::protocol::command_custom_header::CommandCustomHeader;
303
304        #[derive(Debug, Clone)]
305        struct EmptyHeader;
306        impl CommandCustomHeader for EmptyHeader {
307            fn to_map(
308                &self,
309            ) -> Option<std::collections::HashMap<cheetah_string::CheetahString, cheetah_string::CheetahString>>
310            {
311                None
312            }
313        }
314
315        let builder = DefaultAuthenticationContextBuilder::new();
316        let request = RemotingCommand::create_request_command(1, EmptyHeader);
317        let result = builder.build_from_remoting(&request, Some("test-channel")).unwrap();
318
319        assert_eq!(result.base.channel_id(), Some(&CheetahString::from("test-channel")));
320        assert!(result.username().is_none());
321    }
322}