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