rocketmq_auth/authentication/strategy/
stateful_authentication_strategy.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//! Stateful authentication strategy with caching support.
16//!
17//! This strategy caches authentication results per channel/user combination
18//! to avoid redundant authentication checks for the same connection.
19
20use std::collections::HashSet;
21use std::sync::atomic::AtomicU64;
22use std::sync::atomic::Ordering;
23use std::sync::Arc;
24use std::time::Duration;
25
26use crate::authentication::context::default_authentication_context::DefaultAuthenticationContext;
27use crate::authentication::provider::AuthenticationProvider;
28use crate::authentication::strategy::abstract_authentication_strategy::AbstractAuthenticationStrategy;
29use crate::authentication::strategy::authenticate_with_provider;
30use crate::authentication::strategy::authentication_strategy::AuthenticationStrategy;
31use crate::authentication::strategy::AuthenticationFuture;
32use crate::authorization::context::authentication_context::AuthenticationContext;
33use crate::config::AuthConfig;
34use crate::AuthFailureKind;
35use crate::AuthMetrics;
36use crate::AuthOperation;
37use crate::AuthServiceError;
38use crate::AuthServiceResult;
39use moka::sync::Cache;
40
41const POUND: &str = "#";
42
43/// Authentication cache entry storing the result of an authentication attempt.
44#[derive(Clone, Debug)]
45struct AuthCacheEntry {
46    /// Whether authentication succeeded
47    success: bool,
48    /// Bounded error kind if authentication failed.
49    error_kind: Option<AuthFailureKind>,
50}
51
52impl AuthCacheEntry {
53    fn success() -> Self {
54        Self {
55            success: true,
56            error_kind: None,
57        }
58    }
59
60    fn failure(error: &AuthServiceError) -> Self {
61        Self {
62            success: false,
63            error_kind: Some(error.kind()),
64        }
65    }
66}
67
68/// Stateful authentication strategy with caching.
69///
70/// Caches authentication results per channel ID and username combination
71/// to avoid redundant authentication checks for the same connection/user.
72///
73/// - Configurable cache size and expiration
74/// - Falls back to stateless evaluation if channel ID is not available
75pub struct StatefulAuthenticationStrategy<P>
76where
77    P: AuthenticationProvider<Context = DefaultAuthenticationContext>,
78{
79    /// Authentication configuration
80    auth_config: AuthConfig,
81    /// Whitelist of RPC codes that bypass authentication
82    authentication_white_set: HashSet<String>,
83    /// Authentication provider
84    authentication_provider: Option<Arc<P>>,
85    /// Cache for authentication results
86    auth_cache: Cache<String, AuthCacheEntry>,
87    /// Shared ACL snapshot generation. A successful ACL reload advances this value.
88    acl_generation: Arc<AtomicU64>,
89    /// Last generation observed by this strategy, used to drop old entries promptly.
90    last_seen_acl_generation: AtomicU64,
91    metrics: AuthMetrics,
92}
93
94impl<P> StatefulAuthenticationStrategy<P>
95where
96    P: AuthenticationProvider<Context = DefaultAuthenticationContext> + Send + Sync + 'static,
97{
98    /// Create a new stateful authentication strategy.
99    pub fn new(auth_config: AuthConfig, provider: Option<Arc<P>>) -> Self {
100        Self::new_with_acl_generation(auth_config, provider, Arc::new(AtomicU64::new(0)))
101    }
102
103    /// Create a new stateful authentication strategy bound to a shared ACL generation counter.
104    pub fn new_with_acl_generation(
105        auth_config: AuthConfig,
106        provider: Option<Arc<P>>,
107        acl_generation: Arc<AtomicU64>,
108    ) -> Self {
109        Self::new_with_acl_generation_and_metrics(auth_config, provider, acl_generation, AuthMetrics::default())
110    }
111
112    pub fn new_with_acl_generation_and_metrics(
113        auth_config: AuthConfig,
114        provider: Option<Arc<P>>,
115        acl_generation: Arc<AtomicU64>,
116        metrics: AuthMetrics,
117    ) -> Self {
118        let mut authentication_white_set = HashSet::new();
119
120        let whitelist = auth_config.authentication_whitelist.to_string();
121        if !whitelist.is_empty() {
122            for rpc_code in whitelist.split(',') {
123                let trimmed = rpc_code.trim();
124                if !trimmed.is_empty() {
125                    authentication_white_set.insert(trimmed.to_string());
126                }
127            }
128        }
129
130        let auth_cache = Cache::builder()
131            .time_to_live(Duration::from_secs(
132                auth_config.stateful_authentication_cache_expired_second as u64,
133            ))
134            .max_capacity(auth_config.stateful_authentication_cache_max_num as u64)
135            .build();
136        let initial_generation = acl_generation.load(Ordering::Acquire);
137
138        Self {
139            auth_config,
140            authentication_white_set,
141            authentication_provider: provider,
142            auth_cache,
143            acl_generation,
144            last_seen_acl_generation: AtomicU64::new(initial_generation),
145            metrics,
146        }
147    }
148
149    /// Build cache key from authentication context.
150    ///
151    /// Key format:
152    /// - `{generation}#{channel_id}` if username is not available
153    /// - `{generation}#{channel_id}#{username}` if username is available
154    fn build_cache_key(&self, context: &DefaultAuthenticationContext) -> AuthServiceResult<String> {
155        let channel_id = context
156            .base
157            .channel_id()
158            .ok_or_else(|| AuthServiceError::new(AuthOperation::Authenticate, AuthFailureKind::InvalidInput))?;
159        let generation = self.refresh_acl_generation();
160
161        if let Some(username) = context.username() {
162            Ok(format!("{}{}{}{}{}", generation, POUND, channel_id, POUND, username))
163        } else {
164            Ok(format!("{}{}{}", generation, POUND, channel_id))
165        }
166    }
167
168    fn refresh_acl_generation(&self) -> u64 {
169        let current = self.acl_generation.load(Ordering::Acquire);
170        let previous = self.last_seen_acl_generation.load(Ordering::Acquire);
171        if previous != current
172            && self
173                .last_seen_acl_generation
174                .compare_exchange(previous, current, Ordering::AcqRel, Ordering::Acquire)
175                .is_ok()
176        {
177            self.auth_cache.invalidate_all();
178            self.metrics.record_cache_invalidation();
179        }
180        current
181    }
182
183    /// Perform actual authentication without caching.
184    async fn do_authenticate_internal(&self, context: &dyn AuthenticationContext) -> AuthServiceResult<()> {
185        if !self.auth_config.authentication_enabled {
186            return Ok(());
187        }
188
189        let default_context = context
190            .as_any()
191            .downcast_ref::<DefaultAuthenticationContext>()
192            .ok_or_else(|| AuthServiceError::new(AuthOperation::Authenticate, AuthFailureKind::InvalidInput))?;
193
194        if let Some(rpc_code) = default_context.base.rpc_code() {
195            if self.is_whitelisted(rpc_code.as_str()) {
196                return Ok(());
197            }
198        }
199
200        let provider = match &self.authentication_provider {
201            Some(p) => p,
202            None => return Ok(()),
203        };
204
205        authenticate_with_provider(provider.as_ref(), default_context).await
206    }
207
208    pub fn provider(&self) -> Option<&P> {
209        self.authentication_provider.as_deref()
210    }
211}
212
213impl<P> AbstractAuthenticationStrategy for StatefulAuthenticationStrategy<P>
214where
215    P: AuthenticationProvider<Context = DefaultAuthenticationContext> + Send + Sync + 'static,
216{
217    fn authentication_white_set(&self) -> &HashSet<String> {
218        &self.authentication_white_set
219    }
220}
221
222impl<P> AuthenticationStrategy for StatefulAuthenticationStrategy<P>
223where
224    P: AuthenticationProvider<Context = DefaultAuthenticationContext> + Send + Sync + 'static,
225{
226    fn authenticate<'a>(&'a self, context: &'a dyn AuthenticationContext) -> AuthenticationFuture<'a> {
227        Box::pin(async move {
228            let default_context = context
229                .as_any()
230                .downcast_ref::<DefaultAuthenticationContext>()
231                .ok_or_else(|| AuthServiceError::new(AuthOperation::Authenticate, AuthFailureKind::InvalidInput))?;
232
233            if default_context.base.channel_id().is_none() {
234                return self.do_authenticate_internal(context).await;
235            }
236
237            let cache_key = self.build_cache_key(default_context)?;
238
239            let result = if let Some(result) = self.auth_cache.get(&cache_key) {
240                self.metrics.record_cache_hit();
241                result
242            } else {
243                self.metrics.record_cache_miss();
244                let result = match self.do_authenticate_internal(context).await {
245                    Ok(()) => AuthCacheEntry::success(),
246                    Err(error) => AuthCacheEntry::failure(&error),
247                };
248                self.auth_cache.insert(cache_key, result.clone());
249                result
250            };
251
252            if !result.success {
253                return Err(AuthServiceError::new(
254                    AuthOperation::Authenticate,
255                    result.error_kind.unwrap_or(AuthFailureKind::Unauthenticated),
256                ));
257            }
258
259            Ok(())
260        })
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use std::any::Any;
267
268    use super::*;
269    use crate::authentication::provider::AuthenticationProvider;
270    use crate::authentication::provider::DefaultAuthenticationProvider;
271    use cheetah_string::CheetahString;
272
273    struct CountingAuthenticationProvider {
274        calls: Arc<std::sync::atomic::AtomicUsize>,
275        should_succeed: Arc<std::sync::atomic::AtomicBool>,
276    }
277
278    impl AuthenticationProvider for CountingAuthenticationProvider {
279        type Context = DefaultAuthenticationContext;
280
281        async fn initialize(
282            &mut self,
283            _config: AuthConfig,
284            _metadata_service: Option<Arc<dyn Any + Send + Sync>>,
285        ) -> AuthServiceResult<()> {
286            Ok(())
287        }
288
289        async fn authenticate(&self, _context: &Self::Context) -> AuthServiceResult<()> {
290            self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
291            if self.should_succeed.load(std::sync::atomic::Ordering::SeqCst) {
292                Ok(())
293            } else {
294                Err(AuthServiceError::new(
295                    AuthOperation::Authenticate,
296                    AuthFailureKind::Unauthenticated,
297                ))
298            }
299        }
300
301        fn new_context_from_metadata(
302            &self,
303            _metadata: &std::collections::HashMap<String, String>,
304            _request: Box<dyn Any + Send>,
305        ) -> Self::Context {
306            DefaultAuthenticationContext::new()
307        }
308
309        fn new_context_from_command(
310            &self,
311            _command: &rocketmq_protocol::protocol::remoting_command::RemotingCommand,
312        ) -> Self::Context {
313            DefaultAuthenticationContext::new()
314        }
315    }
316
317    #[test]
318    fn test_stateful_strategy_creation() {
319        let config = AuthConfig::default();
320        let provider = Arc::new(DefaultAuthenticationProvider::new());
321
322        let strategy = StatefulAuthenticationStrategy::new(config, Some(provider));
323
324        assert!(strategy.provider().is_some());
325        assert!(strategy.authentication_white_set().is_empty());
326    }
327
328    #[test]
329    fn test_build_cache_key_with_username() {
330        let config = AuthConfig::default();
331        let strategy: StatefulAuthenticationStrategy<DefaultAuthenticationProvider> =
332            StatefulAuthenticationStrategy::new(config, None);
333
334        let mut context = DefaultAuthenticationContext::new();
335        context.base.set_channel_id(Some(CheetahString::from("channel-123")));
336        context.set_username(CheetahString::from("user-456"));
337
338        let key = strategy.build_cache_key(&context).unwrap();
339        assert_eq!(key, "0#channel-123#user-456");
340    }
341
342    #[test]
343    fn test_build_cache_key_without_username() {
344        let config = AuthConfig::default();
345        let strategy: StatefulAuthenticationStrategy<DefaultAuthenticationProvider> =
346            StatefulAuthenticationStrategy::new(config, None);
347
348        let mut context = DefaultAuthenticationContext::new();
349        context.base.set_channel_id(Some(CheetahString::from("channel-789")));
350
351        let key = strategy.build_cache_key(&context).unwrap();
352        assert_eq!(key, "0#channel-789");
353    }
354
355    #[test]
356    fn test_build_cache_key_includes_acl_generation() {
357        let config = AuthConfig::default();
358        let acl_generation = Arc::new(AtomicU64::new(7));
359        let strategy: StatefulAuthenticationStrategy<DefaultAuthenticationProvider> =
360            StatefulAuthenticationStrategy::new_with_acl_generation(config, None, acl_generation.clone());
361
362        let mut context = DefaultAuthenticationContext::new();
363        context.base.set_channel_id(Some(CheetahString::from("channel-789")));
364
365        let key = strategy.build_cache_key(&context).unwrap();
366        assert_eq!(key, "7#channel-789");
367
368        acl_generation.store(8, Ordering::Release);
369        let key = strategy.build_cache_key(&context).unwrap();
370        assert_eq!(key, "8#channel-789");
371    }
372
373    #[tokio::test]
374    async fn test_authentication_disabled() {
375        let config = AuthConfig {
376            authentication_enabled: false,
377            ..Default::default()
378        };
379
380        let strategy: StatefulAuthenticationStrategy<DefaultAuthenticationProvider> =
381            StatefulAuthenticationStrategy::new(config, None);
382
383        let context = DefaultAuthenticationContext::new();
384        let result = strategy.authenticate(&context).await;
385        assert!(result.is_ok());
386    }
387
388    #[test]
389    fn test_cache_expiration_config() {
390        let config = AuthConfig {
391            stateful_authentication_cache_max_num: 500,
392            stateful_authentication_cache_expired_second: 30,
393            ..Default::default()
394        };
395
396        let strategy: StatefulAuthenticationStrategy<DefaultAuthenticationProvider> =
397            StatefulAuthenticationStrategy::new(config, None);
398
399        assert_eq!(strategy.auth_config.stateful_authentication_cache_max_num, 500);
400        assert_eq!(strategy.auth_config.stateful_authentication_cache_expired_second, 30);
401    }
402
403    #[test]
404    fn test_whitelist_parsing() {
405        let config = AuthConfig {
406            authentication_whitelist: "SEND_MESSAGE,PULL_MESSAGE".into(),
407            ..Default::default()
408        };
409
410        let strategy: StatefulAuthenticationStrategy<DefaultAuthenticationProvider> =
411            StatefulAuthenticationStrategy::new(config, None);
412
413        assert!(strategy.authentication_white_set().contains("SEND_MESSAGE"));
414        assert!(strategy.authentication_white_set().contains("PULL_MESSAGE"));
415    }
416
417    #[tokio::test]
418    async fn test_authentication_without_channel_id() {
419        let config = AuthConfig {
420            authentication_enabled: false,
421            ..Default::default()
422        };
423
424        let strategy: StatefulAuthenticationStrategy<DefaultAuthenticationProvider> =
425            StatefulAuthenticationStrategy::new(config, None);
426
427        let context = DefaultAuthenticationContext::new();
428
429        let result = strategy.authenticate(&context).await;
430        assert!(result.is_ok());
431    }
432
433    #[tokio::test(flavor = "multi_thread")]
434    async fn test_acl_generation_change_invalidates_cached_authentication_result() {
435        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
436        let should_succeed = Arc::new(std::sync::atomic::AtomicBool::new(true));
437        let provider = Arc::new(CountingAuthenticationProvider {
438            calls: calls.clone(),
439            should_succeed: should_succeed.clone(),
440        });
441        let acl_generation = Arc::new(AtomicU64::new(0));
442        let strategy = StatefulAuthenticationStrategy::new_with_acl_generation(
443            AuthConfig {
444                authentication_enabled: true,
445                ..AuthConfig::default()
446            },
447            Some(provider),
448            acl_generation.clone(),
449        );
450        let mut context = DefaultAuthenticationContext::new();
451        context.base.set_channel_id(Some(CheetahString::from("channel-1")));
452        context.set_username(CheetahString::from("alice"));
453
454        assert!(strategy.authenticate(&context).await.is_ok());
455        assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
456
457        should_succeed.store(false, std::sync::atomic::Ordering::SeqCst);
458        assert!(strategy.authenticate(&context).await.is_ok());
459        assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
460
461        acl_generation.fetch_add(1, Ordering::AcqRel);
462        let result = strategy.authenticate(&context).await;
463
464        assert!(result.is_err());
465        assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2);
466    }
467
468    #[tokio::test(flavor = "current_thread")]
469    async fn authenticates_inside_current_thread_runtime_without_block_in_place_panic() {
470        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
471        let should_succeed = Arc::new(std::sync::atomic::AtomicBool::new(true));
472        let provider = Arc::new(CountingAuthenticationProvider {
473            calls: calls.clone(),
474            should_succeed,
475        });
476        let strategy = StatefulAuthenticationStrategy::new(
477            AuthConfig {
478                authentication_enabled: true,
479                ..AuthConfig::default()
480            },
481            Some(provider),
482        );
483        let mut context = DefaultAuthenticationContext::new();
484        context
485            .base
486            .set_channel_id(Some(CheetahString::from("channel-current-thread")));
487        context.set_username(CheetahString::from("alice"));
488        assert!(strategy.authenticate(&context).await.is_ok());
489
490        assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
491    }
492}
rocketmq_auth/authentication/strategy/stateful_authentication_strategy.rs

rocketmq_auth/authentication/strategy/
stateful_authentication_strategy.rs
