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 Index(SafeUrl),
21 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)] pub struct CredentialsCache {
36 realms: RwLock<FxHashMap<(Realm, Username), Arc<Authentication>>>,
38 pub(crate) fetches: FxOnceMap<(FetchUrl, Username), Option<Arc<Authentication>>>,
40 urls: RwLock<UrlTrie<Arc<Authentication>>>,
42}
43
44impl Default for CredentialsCache {
45 fn default() -> Self {
46 Self::new()
47 }
48}
49
50impl CredentialsCache {
51 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 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 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 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 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 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 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 pub(crate) fn insert(&self, url: &SafeUrl, credentials: Arc<Authentication>) {
146 if credentials.is_empty() {
148 return;
149 }
150
151 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 self.insert_realm((Realm::from(url), Username::none()), &credentials);
160
161 let mut urls = self.urls.write().unwrap();
163 urls.insert(url, credentials);
164 }
165
166 fn insert_realm(
170 &self,
171 key: (Realm, Username),
172 credentials: &Arc<Authentication>,
173 ) -> Option<Arc<Authentication>> {
174 if credentials.is_empty() {
176 return None;
177 }
178
179 let mut realms = self.realms.write().unwrap();
180
181 if credentials.is_authenticated() {
183 return realms.insert(key, credentials.clone());
184 }
185
186 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 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 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 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 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}