Skip to main content

redis/
auth_management.rs

1//! This module contains utilities for managing token-based authentication
2use std::time::{Duration, SystemTime};
3
4/// Configuration for token refresh behavior
5#[derive(Debug, Clone)]
6#[non_exhaustive]
7pub struct TokenRefreshConfig {
8    /// Fraction of token lifetime after which refresh should be triggered (0.0 to 1.0).
9    /// For example, 0.8 means refresh when 80% of the token's lifetime has elapsed.
10    pub expiration_refresh_ratio: f64,
11    /// Retry configuration for failed refresh attempts
12    pub retry_config: RetryConfig,
13}
14
15impl TokenRefreshConfig {
16    /// Set the expiration refresh ratio
17    pub fn set_expiration_refresh_ratio(mut self, ratio: f64) -> Self {
18        self.expiration_refresh_ratio = ratio;
19        self
20    }
21
22    /// Set the retry configuration
23    pub fn set_retry_config(mut self, retry_config: RetryConfig) -> Self {
24        self.retry_config = retry_config;
25        self
26    }
27}
28
29impl Default for TokenRefreshConfig {
30    fn default() -> Self {
31        Self {
32            expiration_refresh_ratio: 0.8,
33            retry_config: RetryConfig::default(),
34        }
35    }
36}
37
38/// Configuration for handling failed token refresh attempts
39#[derive(Debug, Clone)]
40#[non_exhaustive]
41pub struct RetryConfig {
42    /// The resulting duration is calculated by taking the base to the `n`-th power,
43    /// where `n` denotes the number of past attempts.
44    pub(crate) exponent_base: f32,
45    /// The minimal delay for reconnection attempts
46    pub(crate) min_delay: Duration,
47    /// Apply a maximum delay between connection attempts. The delay between attempts won't be longer than max_delay milliseconds.
48    pub(crate) max_delay: Option<Duration>,
49    /// number_of_retries times, with an exponentially increasing delay
50    pub(crate) number_of_retries: usize,
51}
52
53impl Default for RetryConfig {
54    fn default() -> Self {
55        Self {
56            number_of_retries: 3,
57            min_delay: Duration::from_millis(100),
58            max_delay: Some(Duration::from_secs(30)),
59            exponent_base: 2.0,
60        }
61    }
62}
63
64impl RetryConfig {
65    /// Sets the maximum number of retry attempts for token refresh
66    pub fn set_number_of_retries(mut self, number_of_retries: usize) -> Self {
67        self.number_of_retries = number_of_retries;
68        self
69    }
70
71    /// Sets the minimal delay before attempting to refresh the token after a failure
72    pub fn set_min_delay(mut self, min_delay: Duration) -> Self {
73        self.min_delay = min_delay;
74        self
75    }
76
77    /// Sets the upper bound for retry delays
78    pub fn set_max_delay(mut self, max_delay: Duration) -> Self {
79        self.max_delay = Some(max_delay);
80        self
81    }
82
83    /// Sets the base for exponential backoff
84    pub fn set_exponent_base(mut self, exponent_base: f32) -> Self {
85        self.exponent_base = exponent_base;
86        self
87    }
88}
89
90/// Common logic for credentials management
91pub(crate) mod credentials_management_utils {
92    use super::*;
93
94    /// Calculate the refresh threshold based on the token's lifetime and the refresh ratio
95    #[allow(dead_code)] // Reserved for future use with TokenRefreshConfig
96    pub(crate) fn calculate_refresh_threshold(
97        received_at: SystemTime,
98        expires_at: SystemTime,
99        refresh_ratio: f64,
100    ) -> Option<Duration> {
101        if let Ok(total_lifetime) = expires_at.duration_since(received_at) {
102            Some(Duration::from_secs_f64(
103                total_lifetime.as_secs_f64() * refresh_ratio,
104            ))
105        } else {
106            None
107        }
108    }
109
110    /// Extract the OID claim from a JWT
111    #[cfg(all(feature = "token-based-authentication", feature = "entra-id"))]
112    pub(crate) fn extract_oid_from_jwt(jwt: &str) -> Result<String, String> {
113        use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
114
115        let parts: Vec<&str> = jwt.split('.').collect();
116        if parts.len() != 3 {
117            return Err("Invalid JWT: must have 3 parts".to_string());
118        }
119
120        let payload_bytes = URL_SAFE_NO_PAD
121            .decode(parts[1])
122            .map_err(|e| format!("Failed to decode payload: {e}"))?;
123        let payload_str = String::from_utf8(payload_bytes)
124            .map_err(|e| format!("Payload is not valid UTF-8: {e}"))?;
125
126        if let Some(oid_claim_start_idx) = payload_str.find("\"oid\"")
127            && let Some(colon_idx) = payload_str[oid_claim_start_idx..].find(':')
128        {
129            let oid_value_str = payload_str[oid_claim_start_idx + colon_idx + 1..].trim_start();
130
131            if let Some(stripped_oid_value) = oid_value_str.strip_prefix('"')
132                && let Some(end_quote) = stripped_oid_value.find('"')
133            {
134                return Ok(stripped_oid_value[..end_quote].to_string());
135            }
136        }
137
138        Err("OID claim not found".to_string())
139    }
140}
141
142#[cfg(all(feature = "token-based-authentication", test))]
143mod auth_management_tests {
144    use super::{TokenRefreshConfig, credentials_management_utils};
145    use std::sync::LazyLock;
146
147    const TOKEN_HEADER: &str = "header";
148    const TOKEN_PAYLOAD: &str = "eyJvaWQiOiIxMjM0NTY3OC05YWJjLWRlZi0xMjM0LTU2Nzg5YWJjZGVmMCJ9"; // Payload with "oid" claim
149    const TOKEN_PAYLOAD_NO_OID: &str =
150        "eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNzM1Njg5NjAwfQ"; // Payload with no "oid" claim
151    const TOKEN_SIGNATURE: &str = "signature";
152
153    const OID_CLAIM_VALUE: &str = "12345678-9abc-def-1234-56789abcdef0";
154
155    static TOKEN: LazyLock<String> =
156        LazyLock::new(|| format!("{TOKEN_HEADER}.{TOKEN_PAYLOAD}.{TOKEN_SIGNATURE}"));
157    static TOKEN_WITH_NO_OID: LazyLock<String> =
158        LazyLock::new(|| format!("{TOKEN_HEADER}.{TOKEN_PAYLOAD_NO_OID}.{TOKEN_SIGNATURE}"));
159    static INVALID_TOKEN: LazyLock<String> =
160        LazyLock::new(|| format!("{TOKEN_HEADER}.{TOKEN_PAYLOAD}"));
161
162    #[test]
163    fn test_token_refresh_config() {
164        let config = TokenRefreshConfig::default();
165        assert_eq!(config.expiration_refresh_ratio, 0.8);
166
167        let custom_config = TokenRefreshConfig::default().set_expiration_refresh_ratio(0.9);
168        assert_eq!(custom_config.expiration_refresh_ratio, 0.9);
169    }
170
171    #[test]
172    fn test_refresh_threshold_calculation() {
173        use std::time::{Duration, SystemTime};
174        let config = TokenRefreshConfig::default(); // Default refresh ratio is 0.8
175
176        let received_at = SystemTime::now();
177        let expires_at = received_at + Duration::from_secs(3600); // 1 hour
178
179        let threshold = credentials_management_utils::calculate_refresh_threshold(
180            received_at,
181            expires_at,
182            config.expiration_refresh_ratio,
183        );
184        assert!(threshold.is_some());
185        assert_eq!(threshold.unwrap(), Duration::from_secs(2880)); // 80% of 3600
186    }
187
188    #[cfg(all(feature = "token-based-authentication", feature = "entra-id"))]
189    #[test]
190    fn test_extract_oid_from_jwt() {
191        let result = credentials_management_utils::extract_oid_from_jwt(TOKEN.as_str());
192        assert!(result.is_ok());
193        assert_eq!(result.unwrap(), OID_CLAIM_VALUE);
194    }
195
196    #[cfg(all(feature = "token-based-authentication", feature = "entra-id"))]
197    #[test]
198    fn test_extract_oid_from_jwt_with_invalid_token() {
199        let result = credentials_management_utils::extract_oid_from_jwt(INVALID_TOKEN.as_str());
200        assert!(result.is_err());
201        assert_eq!(result.err().unwrap(), "Invalid JWT: must have 3 parts");
202    }
203
204    #[cfg(all(feature = "token-based-authentication", feature = "entra-id"))]
205    #[test]
206    fn test_extract_oid_from_jwt_with_no_oid_claim() {
207        let result = credentials_management_utils::extract_oid_from_jwt(TOKEN_WITH_NO_OID.as_str());
208        assert!(result.is_err());
209        assert_eq!(result.err().unwrap(), "OID claim not found");
210    }
211}