rocketmq_auth/authentication/builder/
default_authentication_context_builder.rs1use 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
34const CREDENTIAL: &str = "Credential";
36const SIGNATURE: &str = "Signature";
37const ACCESS_KEY: &str = "AccessKey";
38const SIGNATURE_FIELD: &str = "Signature";
39
40const SPACE: char = ' ';
42const COMMA: char = ',';
43const EQUAL: char = '=';
44const SLASH: char = '/';
45
46#[derive(Debug, Default, Clone, Copy)]
78pub struct DefaultAuthenticationContextBuilder;
79
80impl DefaultAuthenticationContextBuilder {
81 pub fn new() -> Self {
83 Self
84 }
85
86 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 fn combine_request_content(
119 request: &RemotingCommand,
120 sorted_fields: &BTreeMap<CheetahString, CheetahString>,
121 ) -> Vec<u8> {
122 let mut content = Vec::new();
123
124 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 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 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 let rpc_code = std::any::type_name_of_val(request);
157 context.base.set_rpc_code(Some(CheetahString::from(rpc_code)));
158
159 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 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 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 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 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 let base64_signature = Self::hex_to_base64(value)?;
208 context.set_signature(CheetahString::from(base64_signature));
209 }
210 _ => {
211 }
213 }
214 }
215
216 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 if let Some(id) = channel_id {
231 context.base.set_channel_id(Some(CheetahString::from(id)));
232 }
233
234 context
236 .base
237 .set_rpc_code(Some(CheetahString::from(request.code().to_string())));
238
239 let fields = match request.ext_fields() {
241 Some(fields) => fields,
242 None => return Ok(context), };
244
245 let access_key = match fields.get(&CheetahString::from(ACCESS_KEY)) {
247 Some(key) => key,
248 None => return Ok(context), };
250
251 context.set_username(access_key.clone());
253
254 if let Some(signature) = fields.get(&CheetahString::from(SIGNATURE_FIELD)) {
256 context.set_signature(signature.clone());
257 }
258
259 let mut sorted_map = BTreeMap::new();
261
262 for (key, value) in fields.iter() {
263 if request.version() <= RocketMqVersion::V4_9_3.ordinal() as i32 && key.as_str() == UNIQUE_MSG_QUERY_FLAG {
265 continue;
266 }
267
268 if key.as_str() != SIGNATURE_FIELD {
270 sorted_map.insert(key.clone(), value.clone());
271 }
272 }
273
274 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}