Use index URL instead of package URL for keyring credential lookups (#12651)

Some registries (like Azure Artifact) can require you to authenticate
separately for every package URL if you do not authenticate for the
/simple endpoint. These changes make the auth middleware aware of index
URL endpoints and attempts to fetch keyring credentials for such an
index URL when making a request to any URL it's a prefix of.

The current uv behavior is to cache credentials either at the request
URL or realm level. But with these changes, we also need to cache
credentials at the index level. Note that when uv does not detect an
index URL for a request URL, it will continue to apply the old behavior.

Addresses part of #4056
Closes #4583
Closes #11236
Closes #11391
Closes #11507

Author: jtfmumm

crates/uv-auth/src/cache.rs

@@ -17,8 +17,8 @@ type FxOnceMap<K, V> = OnceMap<K, V, BuildHasherDefault<FxHasher>>;
 pub struct CredentialsCache {
     /// A cache per realm and username
     realms: RwLock<FxHashMap<(Realm, Username), Arc<Credentials>>>,
-    /// A cache tracking the result of fetches from external services
-    pub(crate) fetches: FxOnceMap<(Realm, Username), Option<Arc<Credentials>>>,
+    /// A cache tracking the result of realm or index URL fetches from external services
+    pub(crate) fetches: FxOnceMap<(String, Username), Option<Arc<Credentials>>>,
     /// A cache per URL, uses a trie for efficient prefix queries.
     urls: RwLock<UrlTrie>,
 }

crates/uv-auth/src/policy.rs renamed to crates/uv-auth/src/index.rs

@@ -1,6 +1,6 @@
 use std::fmt::{self, Display, Formatter};
 
-use rustc_hash::FxHashMap;
+use rustc_hash::FxHashSet;
 use url::Url;
 
 /// When to use authentication.
@@ -47,39 +47,68 @@ impl Display for AuthPolicy {
         }
     }
 }
+
+// TODO(john): We are not using `uv_distribution_types::Index` directly
+// here because it would cause circular crate dependencies. However, this
+// could potentially make sense for a future refactor.
+#[derive(Debug, Clone, Hash, Eq, PartialEq)]
+pub struct Index {
+    pub url: Url,
+    /// The root endpoint where authentication is applied.
+    /// For PEP 503 endpoints, this excludes `/simple`.
+    pub root_url: Url,
+    pub auth_policy: AuthPolicy,
+}
+
 #[derive(Debug, Default, Clone, Eq, PartialEq)]
-pub struct UrlAuthPolicies(FxHashMap<Url, AuthPolicy>);
+pub struct Indexes(FxHashSet<Index>);
 
-impl UrlAuthPolicies {
+impl Indexes {
     pub fn new() -> Self {
-        Self(FxHashMap::default())
+        Self(FxHashSet::default())
     }
 
-    /// Create a new [`UrlAuthPolicies`] from a list of URL and [`AuthPolicy`]
-    /// tuples.
-    pub fn from_tuples(tuples: impl IntoIterator<Item = (Url, AuthPolicy)>) -> Self {
-        let mut auth_policies = Self::new();
-        for (url, auth_policy) in tuples {
-            auth_policies.add_policy(url, auth_policy);
+    /// Create a new [`AuthIndexUrls`] from an iterator of [`AuthIndexUrl`]s.
+    pub fn from_indexes(urls: impl IntoIterator<Item = Index>) -> Self {
+        let mut index_urls = Self::new();
+        for url in urls {
+            index_urls.0.insert(url);
         }
-        auth_policies
+        index_urls
     }
 
-    /// An [`AuthPolicy`] for a URL.
-    pub fn add_policy(&mut self, url: Url, auth_policy: AuthPolicy) {
-        self.0.insert(url, auth_policy);
+    /// Get the index URL prefix for a URL if one exists.
+    pub fn index_url_for(&self, url: &Url) -> Option<&Url> {
+        // TODO(john): There are probably not many URLs to iterate through,
+        // but we could use a trie instead of a HashSet here for more
+        // efficient search.
+        self.0
+            .iter()
+            .find(|index| is_url_prefix(&index.root_url, url))
+            .map(|index| &index.url)
     }
 
     /// Get the [`AuthPolicy`] for a URL.
     pub fn policy_for(&self, url: &Url) -> AuthPolicy {
         // TODO(john): There are probably not many URLs to iterate through,
         // but we could use a trie instead of a HashMap here for more
         // efficient search.
-        for (auth_url, auth_policy) in &self.0 {
-            if url.as_str().starts_with(auth_url.as_str()) {
-                return *auth_policy;
+        for index in &self.0 {
+            if is_url_prefix(&index.root_url, url) {
+                return index.auth_policy;
             }
         }
         AuthPolicy::Auto
     }
 }
+
+fn is_url_prefix(base: &Url, url: &Url) -> bool {
+    if base.scheme() != url.scheme()
+        || base.host_str() != url.host_str()
+        || base.port_or_known_default() != url.port_or_known_default()
+    {
+        return false;
+    }
+
+    url.path().starts_with(base.path())
+}

crates/uv-auth/src/lib.rs

@@ -5,16 +5,16 @@ use url::Url;
 
 use cache::CredentialsCache;
 pub use credentials::Credentials;
+pub use index::{AuthPolicy, Index, Indexes};
 pub use keyring::KeyringProvider;
 pub use middleware::AuthMiddleware;
-pub use policy::{AuthPolicy, UrlAuthPolicies};
 use realm::Realm;
 
 mod cache;
 mod credentials;
+mod index;
 mod keyring;
 mod middleware;
-mod policy;
 mod realm;
 
 // TODO(zanieb): Consider passing a cache explicitly throughout