Skip to main content

crates_docs/server/auth/
token.rs

1//! Token store and token information
2
3use std::collections::HashMap;
4
5use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
6use tokio::time::{timeout, Duration};
7
8/// Default lock acquisition timeout (500ms)
9const DEFAULT_LOCK_TIMEOUT_MS: u64 = 500;
10
11/// Token store operations error types
12#[derive(Debug, thiserror::Error)]
13pub enum TokenStoreError {
14    /// Lock acquisition timed out
15    #[error("Token store lock timeout after {}ms", ms)]
16    LockTimeout {
17        /// Timeout duration in milliseconds
18        ms: u64,
19    },
20}
21
22/// Module-level result type for token store operations
23pub type TokenStoreResult<T> = std::result::Result<T, TokenStoreError>;
24
25/// Simple async in-memory token store (production should use Redis or database)
26///
27/// # Differences from previous implementation:
28///
29/// 1. Uses `tokio::sync::RwLock` instead of `std::sync::RwLock` for async operations
30/// 2. Lock acquisition has timeout protection (default 500ms)
31/// 3. Returns proper errors instead of panicking on lock failures
32/// 4. All methods are now async
33/// 5. Does not have poison mechanics (tokio `RwLock` is panic-safe)
34#[derive(Default)]
35pub struct TokenStore {
36    tokens: RwLock<HashMap<String, TokenInfo>>,
37}
38
39/// OAuth token information
40#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
41pub struct TokenInfo {
42    /// Access token
43    pub access_token: String,
44    /// Refresh token (optional)
45    pub refresh_token: Option<String>,
46    /// Token expiration time
47    pub expires_at: chrono::DateTime<chrono::Utc>,
48    /// Authorization scopes
49    pub scopes: Vec<String>,
50    /// User ID (optional)
51    pub user_id: Option<String>,
52    /// User email (optional)
53    pub user_email: Option<String>,
54}
55
56impl TokenStore {
57    /// Create a new token store
58    #[must_use]
59    pub fn new() -> Self {
60        Self::default()
61    }
62
63    /// Store token in the store
64    ///
65    /// # Errors
66    ///
67    /// Returns `TokenStoreError::LockTimeout` if write lock cannot be acquired within timeout.
68    pub async fn store_token(&self, key: String, token: TokenInfo) -> TokenStoreResult<()> {
69        let mut guard = self.acquire_write_lock().await?;
70        guard.insert(key, token);
71        Ok(())
72    }
73
74    /// Get a token from the store
75    ///
76    /// # Errors
77    ///
78    /// Returns `TokenStoreError::LockTimeout` if read lock cannot be acquired within timeout.
79    pub async fn get_token(&self, key: &str) -> TokenStoreResult<Option<TokenInfo>> {
80        let guard = self.acquire_read_lock().await?;
81        Ok(guard.get(key).cloned())
82    }
83
84    /// Remove a token from the store
85    ///
86    /// # Errors
87    ///
88    /// Returns `TokenStoreError::LockTimeout` if write lock cannot be acquired within timeout.
89    pub async fn remove_token(&self, key: &str) -> TokenStoreResult<()> {
90        let mut guard = self.acquire_write_lock().await?;
91        guard.remove(key);
92        Ok(())
93    }
94
95    /// Cleanup all expired tokens from the store
96    ///
97    /// # Errors
98    ///
99    /// Returns `TokenStoreError::LockTimeout` if write lock cannot be acquired within timeout.
100    pub async fn cleanup_expired(&self) -> TokenStoreResult<()> {
101        let now = chrono::Utc::now();
102        let mut guard = self.acquire_write_lock().await?;
103        guard.retain(|_, token| token.expires_at > now);
104        Ok(())
105    }
106
107    /// Acquire a read lock with timeout protection
108    ///
109    /// Uses `tokio::time::timeout` to prevent indefinite blocking.
110    /// This is critical for concurrent async workloads where lock contention might occur.
111    ///
112    /// # Errors
113    ///
114    /// - `TokenStoreError::LockTimeout` - Lock not acquired within `DEFAULT_LOCK_TIMEOUT_MS` (500ms)
115    async fn acquire_read_lock(
116        &self,
117    ) -> TokenStoreResult<RwLockReadGuard<'_, HashMap<String, TokenInfo>>> {
118        match timeout(
119            Duration::from_millis(DEFAULT_LOCK_TIMEOUT_MS),
120            self.tokens.read(),
121        )
122        .await
123        {
124            Ok(guard) => Ok(guard),
125            Err(_elapsed) => Err(TokenStoreError::LockTimeout {
126                ms: DEFAULT_LOCK_TIMEOUT_MS,
127            }),
128        }
129    }
130
131    /// Acquire a write lock with timeout protection
132    ///
133    /// Uses `tokio::time::timeout` to prevent indefinite blocking.
134    /// This is critical for concurrent async workloads where lock contention might occur.
135    ///
136    /// # Errors
137    ///
138    /// - `TokenStoreError::LockTimeout` - Lock not acquired within `DEFAULT_LOCK_TIMEOUT_MS` (500ms)
139    async fn acquire_write_lock(
140        &self,
141    ) -> TokenStoreResult<RwLockWriteGuard<'_, HashMap<String, TokenInfo>>> {
142        match timeout(
143            Duration::from_millis(DEFAULT_LOCK_TIMEOUT_MS),
144            self.tokens.write(),
145        )
146        .await
147        {
148            Ok(guard) => Ok(guard),
149            Err(_elapsed) => Err(TokenStoreError::LockTimeout {
150                ms: DEFAULT_LOCK_TIMEOUT_MS,
151            }),
152        }
153    }
154}