pipedash_core/infrastructure/
token_store.rs

1use std::collections::HashMap;
2use std::sync::RwLock;
3
4use async_trait::async_trait;
5
6use crate::domain::{
7    DomainError,
8    DomainResult,
9};
10
11#[async_trait]
12pub trait TokenStore: Send + Sync {
13    async fn store_token(&self, provider_id: i64, token: &str) -> DomainResult<()>;
14
15    async fn get_token(&self, provider_id: i64) -> DomainResult<String>;
16
17    async fn delete_token(&self, provider_id: i64) -> DomainResult<()>;
18
19    async fn get_all_tokens(&self) -> DomainResult<HashMap<i64, String>>;
20
21    async fn get_token_by_name(&self, _name: &str) -> DomainResult<String> {
22        Err(DomainError::InternalError(
23            "This token store does not support name-based lookups".into(),
24        ))
25    }
26
27    async fn export_encrypted(&self, _password: &str) -> DomainResult<Vec<u8>> {
28        Err(DomainError::InternalError(
29            "This token store does not support export".into(),
30        ))
31    }
32
33    async fn import_encrypted(&self, _data: &[u8], _password: &str) -> DomainResult<()> {
34        Err(DomainError::InternalError(
35            "This token store does not support import".into(),
36        ))
37    }
38
39    async fn warmup(&self) -> DomainResult<()> {
40        Ok(())
41    }
42}
43
44pub struct MemoryTokenStore {
45    tokens: RwLock<HashMap<i64, String>>,
46}
47
48impl MemoryTokenStore {
49    pub fn new() -> Self {
50        Self {
51            tokens: RwLock::new(HashMap::new()),
52        }
53    }
54
55    pub fn with_tokens(tokens: HashMap<i64, String>) -> Self {
56        Self {
57            tokens: RwLock::new(tokens),
58        }
59    }
60}
61
62impl Default for MemoryTokenStore {
63    fn default() -> Self {
64        Self::new()
65    }
66}
67
68#[async_trait]
69impl TokenStore for MemoryTokenStore {
70    async fn store_token(&self, provider_id: i64, token: &str) -> DomainResult<()> {
71        let mut tokens = self
72            .tokens
73            .write()
74            .map_err(|e| DomainError::InternalError(format!("Lock poisoned: {}", e)))?;
75        tokens.insert(provider_id, token.to_string());
76        Ok(())
77    }
78
79    async fn get_token(&self, provider_id: i64) -> DomainResult<String> {
80        let tokens = self
81            .tokens
82            .read()
83            .map_err(|e| DomainError::InternalError(format!("Lock poisoned: {}", e)))?;
84        tokens.get(&provider_id).cloned().ok_or_else(|| {
85            DomainError::NotFound(format!("Token not found for provider {}", provider_id))
86        })
87    }
88
89    async fn delete_token(&self, provider_id: i64) -> DomainResult<()> {
90        let mut tokens = self
91            .tokens
92            .write()
93            .map_err(|e| DomainError::InternalError(format!("Lock poisoned: {}", e)))?;
94        tokens.remove(&provider_id);
95        Ok(())
96    }
97
98    async fn get_all_tokens(&self) -> DomainResult<HashMap<i64, String>> {
99        let tokens = self
100            .tokens
101            .read()
102            .map_err(|e| DomainError::InternalError(format!("Lock poisoned: {}", e)))?;
103        Ok(tokens.clone())
104    }
105}
106
107pub struct EnvTokenStore {
108    prefix: String,
109    cache: RwLock<HashMap<i64, String>>,
110}
111
112impl EnvTokenStore {
113    pub fn new() -> Self {
114        Self::with_prefix("PIPEDASH_TOKEN_")
115    }
116
117    pub fn with_prefix(prefix: &str) -> Self {
118        Self {
119            prefix: prefix.to_string(),
120            cache: RwLock::new(HashMap::new()),
121        }
122    }
123
124    fn env_var_name(&self, provider_id: i64) -> String {
125        format!("{}{}", self.prefix, provider_id)
126    }
127}
128
129impl Default for EnvTokenStore {
130    fn default() -> Self {
131        Self::new()
132    }
133}
134
135#[async_trait]
136impl TokenStore for EnvTokenStore {
137    async fn store_token(&self, provider_id: i64, token: &str) -> DomainResult<()> {
138        let mut cache = self
139            .cache
140            .write()
141            .map_err(|e| DomainError::InternalError(format!("Lock poisoned: {}", e)))?;
142        cache.insert(provider_id, token.to_string());
143        Ok(())
144    }
145
146    async fn get_token(&self, provider_id: i64) -> DomainResult<String> {
147        {
148            let cache = self
149                .cache
150                .read()
151                .map_err(|e| DomainError::InternalError(format!("Lock poisoned: {}", e)))?;
152            if let Some(token) = cache.get(&provider_id) {
153                return Ok(token.clone());
154            }
155        }
156
157        let var_name = self.env_var_name(provider_id);
158        std::env::var(&var_name).map_err(|e| match e {
159            std::env::VarError::NotPresent => DomainError::NotFound(format!(
160                "Token not found for provider {} (env var: {})",
161                provider_id, var_name
162            )),
163            std::env::VarError::NotUnicode(_) => DomainError::DatabaseError(format!(
164                "Token for provider {} is not valid UTF-8 (env var: {})",
165                provider_id, var_name
166            )),
167        })
168    }
169
170    async fn delete_token(&self, provider_id: i64) -> DomainResult<()> {
171        let mut cache = self
172            .cache
173            .write()
174            .map_err(|e| DomainError::InternalError(format!("Lock poisoned: {}", e)))?;
175        cache.remove(&provider_id);
176        Ok(())
177    }
178
179    async fn get_all_tokens(&self) -> DomainResult<HashMap<i64, String>> {
180        let cache = self
181            .cache
182            .read()
183            .map_err(|e| DomainError::InternalError(format!("Lock poisoned: {}", e)))?;
184        Ok(cache.clone())
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    #[tokio::test]
193    async fn test_memory_store() {
194        let store = MemoryTokenStore::new();
195
196        store.store_token(1, "secret_token").await.unwrap();
197
198        let token = store.get_token(1).await.unwrap();
199        assert_eq!(token, "secret_token");
200
201        let all = store.get_all_tokens().await.unwrap();
202        assert_eq!(all.len(), 1);
203
204        store.delete_token(1).await.unwrap();
205
206        assert!(store.get_token(1).await.is_err());
207    }
208
209    #[tokio::test]
210    async fn test_memory_store_with_initial_tokens() {
211        let mut initial = HashMap::new();
212        initial.insert(1, "token1".to_string());
213        initial.insert(2, "token2".to_string());
214
215        let store = MemoryTokenStore::with_tokens(initial);
216
217        assert_eq!(store.get_token(1).await.unwrap(), "token1");
218        assert_eq!(store.get_token(2).await.unwrap(), "token2");
219    }
220}
pipedash_core/infrastructure/token_store.rs

pipedash_core/infrastructure/
token_store.rs
