Skip to main content

uv_auth/
cache.rs

1use std::fmt::{self, Display, Formatter};
2use std::hash::BuildHasherDefault;
3use std::sync::Arc;
4use std::sync::RwLock;
5
6use rustc_hash::{FxHashMap, FxHasher};
7use tracing::trace;
8
9use uv_once_map::OnceMap;
10use uv_redacted::SafeUrl;
11
12use crate::credentials::{Authentication, CredentialsFromUrlError, Username};
13use crate::{Credentials, Realm};
14
15type FxOnceMap<K, V> = OnceMap<K, V, BuildHasherDefault<FxHasher>>;
16
17#[derive(Debug, Clone, PartialEq, Eq, Hash)]
18pub(crate) enum FetchUrl {
19    /// A full index URL
20    Index(SafeUrl),
21    /// A realm URL
22    Realm(Realm),
23}
24
25impl Display for FetchUrl {
26    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
27        match self {
28            Self::Index(index) => Display::fmt(index, f),
29            Self::Realm(realm) => Display::fmt(realm, f),
30        }
31    }
32}
33
34#[derive(Debug)] // All internal types are redacted.
35pub struct CredentialsCache {
36    /// A cache per realm and username
37    realms: RwLock<FxHashMap<(Realm, Username), Arc<Authentication>>>,
38    /// A cache tracking the result of realm or index URL fetches from external services
39    pub(crate) fetches: FxOnceMap<(FetchUrl, Username), Option<Arc<Authentication>>>,
40    /// A cache per URL, uses a trie for efficient prefix queries.
41    urls: RwLock<UrlTrie<Arc<Authentication>>>,
42}
43
44impl Default for CredentialsCache {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50impl CredentialsCache {
51    /// Create a new cache.
52    pub fn new() -> Self {
53        Self {
54            fetches: FxOnceMap::default(),
55            realms: RwLock::new(FxHashMap::default()),
56            urls: RwLock::new(UrlTrie::new()),
57        }
58    }
59
60    /// Populate the global authentication store with credentials on a URL, if there are any.
61    ///
62    /// Returns `true` if the store was updated.
63    pub fn store_credentials_from_url(
64        &self,
65        url: &SafeUrl,
66    ) -> Result<bool, CredentialsFromUrlError> {
67        if let Some(credentials) = Credentials::from_url(url)? {
68            trace!("Caching credentials for `{url}`");
69            self.insert(url, Arc::new(Authentication::from(credentials)));
70            Ok(true)
71        } else {
72            Ok(false)
73        }
74    }
75
76    /// Populate the global authentication store with credentials on a URL, if there are any.
77    ///
78    /// Returns `true` if the store was updated.
79    pub fn store_credentials(&self, url: &SafeUrl, credentials: Credentials) {
80        trace!("Caching credentials for `{url}`");
81        self.insert(url, Arc::new(Authentication::from(credentials)));
82    }
83
84    /// Return the credentials that should be used for a realm and username, if any.
85    pub(crate) fn get_realm(
86        &self,
87        realm: Realm,
88        username: Username,
89    ) -> Option<Arc<Authentication>> {
90        let realms = self.realms.read().unwrap();
91        let given_username = username.is_some();
92        let key = (realm, username);
93        let realm_username = fmt::from_fn(|f| {
94            let (realm, username) = &key;
95            if let Some(username) = username.as_deref() {
96                write!(f, "{username}@{realm}")
97            } else {
98                write!(f, "{realm}")
99            }
100        });
101
102        let Some(credentials) = realms.get(&key).cloned() else {
103            trace!("No credentials in cache for realm {realm_username}");
104            return None;
105        };
106
107        if given_username && credentials.password().is_none() {
108            // If given a username, don't return password-less credentials
109            trace!("No password in cache for realm {realm_username}");
110            return None;
111        }
112
113        trace!("Found cached credentials for realm {realm_username}");
114        Some(credentials)
115    }
116
117    /// Return the cached credentials for a URL and username, if any.
118    ///
119    /// Note we do not cache per username, but if a username is passed we will confirm that the
120    /// cached credentials have a username equal to the provided one — otherwise `None` is returned.
121    /// If multiple usernames are used per URL, the realm cache should be queried instead.
122    pub(crate) fn get_url(
123        &self,
124        url: &SafeUrl,
125        username: &Username,
126    ) -> Option<Arc<Authentication>> {
127        let urls = self.urls.read().unwrap();
128        let credentials = urls.get(url);
129        if let Some(credentials) = credentials {
130            if username.is_none() || username.as_deref() == credentials.username() {
131                if username.is_some() && credentials.password().is_none() {
132                    // If given a username, don't return password-less credentials
133                    trace!("No password in cache for URL `{url}`");
134                    return None;
135                }
136                trace!("Found cached credentials for URL `{url}`");
137                return Some(credentials.clone());
138            }
139        }
140        trace!("No credentials in cache for URL `{url}`");
141        None
142    }
143
144    /// Update the cache with the given credentials.
145    pub(crate) fn insert(&self, url: &SafeUrl, credentials: Arc<Authentication>) {
146        // Do not cache empty credentials
147        if credentials.is_empty() {
148            return;
149        }
150
151        // Insert an entry for requests including the username
152        let username = credentials.to_username();
153        if username.is_some() {
154            let realm = (Realm::from(url), username);
155            self.insert_realm(realm, &credentials);
156        }
157
158        // Insert an entry for requests with no username
159        self.insert_realm((Realm::from(url), Username::none()), &credentials);
160
161        // Insert an entry for the URL
162        let mut urls = self.urls.write().unwrap();
163        urls.insert(url, credentials);
164    }
165
166    /// Private interface to update a realm cache entry.
167    ///
168    /// Returns replaced credentials, if any.
169    fn insert_realm(
170        &self,
171        key: (Realm, Username),
172        credentials: &Arc<Authentication>,
173    ) -> Option<Arc<Authentication>> {
174        // Do not cache empty credentials
175        if credentials.is_empty() {
176            return None;
177        }
178
179        let mut realms = self.realms.write().unwrap();
180
181        // Always replace existing entries if we have a password or token
182        if credentials.is_authenticated() {
183            return realms.insert(key, credentials.clone());
184        }
185
186        // If we only have a username, add a new entry or replace an existing entry if it doesn't have a password
187        let existing = realms.get(&key);
188        if existing.is_none_or(|credentials| credentials.password().is_none()) {
189            return realms.insert(key, credentials.clone());
190        }
191
192        None
193    }
194}
195
196#[derive(Debug)]
197struct UrlTrie<T> {
198    states: Vec<TrieState<T>>,
199}
200
201#[derive(Debug)]
202struct TrieState<T> {
203    children: Vec<(String, usize)>,
204    value: Option<T>,
205}
206
207impl<T> Default for TrieState<T> {
208    fn default() -> Self {
209        Self {
210            children: vec![],
211            value: None,
212        }
213    }
214}
215
216impl<T> UrlTrie<T> {
217    fn new() -> Self {
218        let mut trie = Self { states: vec![] };
219        trie.alloc();
220        trie
221    }
222
223    fn get(&self, url: &SafeUrl) -> Option<&T> {
224        let segments = url.path_segments()?;
225        let mut state = 0;
226        let realm = Realm::from(url).to_string();
227        for component in [realm.as_str()]
228            .into_iter()
229            .chain(segments.filter(|item| !item.is_empty()))
230        {
231            state = self.states[state].get(component)?;
232            if let Some(ref value) = self.states[state].value {
233                return Some(value);
234            }
235        }
236        self.states[state].value.as_ref()
237    }
238
239    fn insert(&mut self, url: &SafeUrl, value: T) {
240        // Opaque URLs have no path hierarchy for prefix matching.
241        let Some(segments) = url.path_segments() else {
242            return;
243        };
244        let mut state = 0;
245        let realm = Realm::from(url).to_string();
246        for component in [realm.as_str()]
247            .into_iter()
248            .chain(segments.filter(|item| !item.is_empty()))
249        {
250            match self.states[state].index(component) {
251                Ok(i) => state = self.states[state].children[i].1,
252                Err(i) => {
253                    let new_state = self.alloc();
254                    self.states[state]
255                        .children
256                        .insert(i, (component.to_string(), new_state));
257                    state = new_state;
258                }
259            }
260        }
261        self.states[state].value = Some(value);
262    }
263
264    fn alloc(&mut self) -> usize {
265        let id = self.states.len();
266        self.states.push(TrieState::default());
267        id
268    }
269}
270
271impl<T> TrieState<T> {
272    fn get(&self, component: &str) -> Option<usize> {
273        let i = self.index(component).ok()?;
274        Some(self.children[i].1)
275    }
276
277    fn index(&self, component: &str) -> Result<usize, usize> {
278        self.children
279            .binary_search_by(|(label, _)| label.as_str().cmp(component))
280    }
281}
282
283#[cfg(test)]
284mod tests {
285    use uv_redacted::SafeUrlError;
286
287    use crate::Credentials;
288    use crate::credentials::Password;
289
290    use super::*;
291
292    #[test]
293    fn test_trie() {
294        let credentials1 =
295            Credentials::basic(Some("username1".to_string()), Some("password1".to_string()));
296        let credentials2 =
297            Credentials::basic(Some("username2".to_string()), Some("password2".to_string()));
298        let credentials3 =
299            Credentials::basic(Some("username3".to_string()), Some("password3".to_string()));
300        let credentials4 =
301            Credentials::basic(Some("username4".to_string()), Some("password4".to_string()));
302
303        let mut trie = UrlTrie::new();
304        trie.insert(
305            &SafeUrl::parse("https://burntsushi.net").unwrap(),
306            credentials1.clone(),
307        );
308        trie.insert(
309            &SafeUrl::parse("https://astral.sh").unwrap(),
310            credentials2.clone(),
311        );
312        trie.insert(
313            &SafeUrl::parse("https://example.com/foo").unwrap(),
314            credentials3.clone(),
315        );
316        trie.insert(
317            &SafeUrl::parse("https://example.com/bar").unwrap(),
318            credentials4.clone(),
319        );
320
321        let url = SafeUrl::parse("https://burntsushi.net/regex-internals").unwrap();
322        assert_eq!(trie.get(&url), Some(&credentials1));
323
324        let url = SafeUrl::parse("https://burntsushi.net/").unwrap();
325        assert_eq!(trie.get(&url), Some(&credentials1));
326
327        let url = SafeUrl::parse("https://astral.sh/about").unwrap();
328        assert_eq!(trie.get(&url), Some(&credentials2));
329
330        let url = SafeUrl::parse("https://example.com/foo").unwrap();
331        assert_eq!(trie.get(&url), Some(&credentials3));
332
333        let url = SafeUrl::parse("https://example.com/foo/").unwrap();
334        assert_eq!(trie.get(&url), Some(&credentials3));
335
336        let url = SafeUrl::parse("https://example.com/foo/bar").unwrap();
337        assert_eq!(trie.get(&url), Some(&credentials3));
338
339        let url = SafeUrl::parse("https://example.com/bar").unwrap();
340        assert_eq!(trie.get(&url), Some(&credentials4));
341
342        let url = SafeUrl::parse("https://example.com/bar/").unwrap();
343        assert_eq!(trie.get(&url), Some(&credentials4));
344
345        let url = SafeUrl::parse("https://example.com/bar/foo").unwrap();
346        assert_eq!(trie.get(&url), Some(&credentials4));
347
348        let url = SafeUrl::parse("https://example.com/about").unwrap();
349        assert_eq!(trie.get(&url), None);
350
351        let url = SafeUrl::parse("https://example.com/foobar").unwrap();
352        assert_eq!(trie.get(&url), None);
353    }
354
355    #[test]
356    fn test_trie_opaque_url() -> Result<(), SafeUrlError> {
357        let mut trie = UrlTrie::new();
358        let url = SafeUrl::parse("git+https:foo")?;
359        let credentials =
360            Credentials::basic(Some("username".to_string()), Some("password".to_string()));
361
362        assert_eq!(trie.get(&url), None);
363        trie.insert(&url, credentials.clone());
364        assert_eq!(trie.get(&url), None);
365
366        // Opaque URLs must not share credentials with hierarchical URLs in the same realm.
367        let base_url = SafeUrl::parse("git+https:/")?;
368        assert_eq!(trie.get(&base_url), None);
369        trie.insert(&base_url, credentials.clone());
370        assert_eq!(trie.get(&url), None);
371        assert_eq!(trie.get(&base_url), Some(&credentials));
372
373        Ok(())
374    }
375
376    #[test]
377    fn test_url_with_credentials() {
378        let username = Username::new(Some(String::from("username")));
379        let password = Password::new(String::from("password"));
380        let credentials = Arc::new(Authentication::from(Credentials::Basic {
381            username: username.clone(),
382            password: Some(password),
383        }));
384        let cache = CredentialsCache::default();
385        // Insert with URL with credentials and get with redacted URL.
386        let url = SafeUrl::parse("https://username:password@example.com/foobar").unwrap();
387        cache.insert(&url, credentials.clone());
388        assert_eq!(cache.get_url(&url, &username), Some(credentials.clone()));
389        // Insert with redacted URL and get with URL with credentials.
390        let url = SafeUrl::parse("https://username:password@second-example.com/foobar").unwrap();
391        cache.insert(&url, credentials.clone());
392        assert_eq!(cache.get_url(&url, &username), Some(credentials.clone()));
393    }
394}