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 Index(DisplaySafeUrl),
22 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)] pub struct CredentialsCache {
37 realms: RwLock<FxHashMap<(Realm, Username), Arc<Authentication>>>,
39 pub(crate) fetches: FxOnceMap<(FetchUrl, Username), Option<Arc<Authentication>>>,
41 urls: RwLock<UrlTrie<Arc<Authentication>>>,
43}
44
45impl Default for CredentialsCache {
46 fn default() -> Self {
47 Self::new()
48 }
49}
50
51impl CredentialsCache {
52 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 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 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 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 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 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 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 pub(crate) fn insert(&self, url: &DisplaySafeUrl, credentials: Arc<Authentication>) {
147 if credentials.is_empty() {
149 return;
150 }
151
152 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 self.insert_realm((Realm::from(url), Username::none()), &credentials);
161
162 let mut urls = self.urls.write().unwrap();
164 urls.insert(url, credentials);
165 }
166
167 fn insert_realm(
171 &self,
172 key: (Realm, Username),
173 credentials: &Arc<Authentication>,
174 ) -> Option<Arc<Authentication>> {
175 if credentials.is_empty() {
177 return None;
178 }
179
180 let mut realms = self.realms.write().unwrap();
181
182 if credentials.is_authenticated() {
184 return realms.insert(key, credentials.clone());
185 }
186
187 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 mut state = 0;
226 let realm = Realm::from(url).to_string();
227 for component in [realm.as_str()]
228 .into_iter()
229 .chain(url.path_segments().unwrap().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: &Url, value: T) {
240 let mut state = 0;
241 let realm = Realm::from(url).to_string();
242 for component in [realm.as_str()]
243 .into_iter()
244 .chain(url.path_segments().unwrap().filter(|item| !item.is_empty()))
245 {
246 match self.states[state].index(component) {
247 Ok(i) => state = self.states[state].children[i].1,
248 Err(i) => {
249 let new_state = self.alloc();
250 self.states[state]
251 .children
252 .insert(i, (component.to_string(), new_state));
253 state = new_state;
254 }
255 }
256 }
257 self.states[state].value = Some(value);
258 }
259
260 fn alloc(&mut self) -> usize {
261 let id = self.states.len();
262 self.states.push(TrieState::default());
263 id
264 }
265}
266
267impl<T> TrieState<T> {
268 fn get(&self, component: &str) -> Option<usize> {
269 let i = self.index(component).ok()?;
270 Some(self.children[i].1)
271 }
272
273 fn index(&self, component: &str) -> Result<usize, usize> {
274 self.children
275 .binary_search_by(|(label, _)| label.as_str().cmp(component))
276 }
277}
278
279#[cfg(test)]
280mod tests {
281 use crate::Credentials;
282 use crate::credentials::Password;
283
284 use super::*;
285
286 #[test]
287 fn test_trie() {
288 let credentials1 =
289 Credentials::basic(Some("username1".to_string()), Some("password1".to_string()));
290 let credentials2 =
291 Credentials::basic(Some("username2".to_string()), Some("password2".to_string()));
292 let credentials3 =
293 Credentials::basic(Some("username3".to_string()), Some("password3".to_string()));
294 let credentials4 =
295 Credentials::basic(Some("username4".to_string()), Some("password4".to_string()));
296
297 let mut trie = UrlTrie::new();
298 trie.insert(
299 &Url::parse("https://burntsushi.net").unwrap(),
300 credentials1.clone(),
301 );
302 trie.insert(
303 &Url::parse("https://astral.sh").unwrap(),
304 credentials2.clone(),
305 );
306 trie.insert(
307 &Url::parse("https://example.com/foo").unwrap(),
308 credentials3.clone(),
309 );
310 trie.insert(
311 &Url::parse("https://example.com/bar").unwrap(),
312 credentials4.clone(),
313 );
314
315 let url = Url::parse("https://burntsushi.net/regex-internals").unwrap();
316 assert_eq!(trie.get(&url), Some(&credentials1));
317
318 let url = Url::parse("https://burntsushi.net/").unwrap();
319 assert_eq!(trie.get(&url), Some(&credentials1));
320
321 let url = Url::parse("https://astral.sh/about").unwrap();
322 assert_eq!(trie.get(&url), Some(&credentials2));
323
324 let url = Url::parse("https://example.com/foo").unwrap();
325 assert_eq!(trie.get(&url), Some(&credentials3));
326
327 let url = Url::parse("https://example.com/foo/").unwrap();
328 assert_eq!(trie.get(&url), Some(&credentials3));
329
330 let url = Url::parse("https://example.com/foo/bar").unwrap();
331 assert_eq!(trie.get(&url), Some(&credentials3));
332
333 let url = Url::parse("https://example.com/bar").unwrap();
334 assert_eq!(trie.get(&url), Some(&credentials4));
335
336 let url = Url::parse("https://example.com/bar/").unwrap();
337 assert_eq!(trie.get(&url), Some(&credentials4));
338
339 let url = Url::parse("https://example.com/bar/foo").unwrap();
340 assert_eq!(trie.get(&url), Some(&credentials4));
341
342 let url = Url::parse("https://example.com/about").unwrap();
343 assert_eq!(trie.get(&url), None);
344
345 let url = Url::parse("https://example.com/foobar").unwrap();
346 assert_eq!(trie.get(&url), None);
347 }
348
349 #[test]
350 fn test_url_with_credentials() {
351 let username = Username::new(Some(String::from("username")));
352 let password = Password::new(String::from("password"));
353 let credentials = Arc::new(Authentication::from(Credentials::Basic {
354 username: username.clone(),
355 password: Some(password),
356 }));
357 let cache = CredentialsCache::default();
358 let url = DisplaySafeUrl::parse("https://username:password@example.com/foobar").unwrap();
360 cache.insert(&url, credentials.clone());
361 assert_eq!(cache.get_url(&url, &username), Some(credentials.clone()));
362 let url =
364 DisplaySafeUrl::parse("https://username:password@second-example.com/foobar").unwrap();
365 cache.insert(&url, credentials.clone());
366 assert_eq!(cache.get_url(&url, &username), Some(credentials.clone()));
367 }
368}