fix: resolve registry auth tokens for cargo upgrade

Cargo.toml

@@ -89,6 +89,7 @@ version = "1.0.28"
 trycmd = "1.2.0"
 snapbox = { version = "0.6.24", features = ["cmd", "path"] }
 cargo-test-support = "0.10.0"
+tempfile = "3"
 
 [profile.release]
 panic = "abort"
@@ -128,7 +129,7 @@ create_dir = "warn"
 dbg_macro = "warn"
 debug_assert_with_mut_call = "warn"
 doc_markdown = "warn"
-empty_enum = "warn"
+empty_enums = "warn"
 enum_glob_use = "warn"
 expl_impl_clone_on_copy = "warn"
 explicit_deref_methods = "warn"

src/bin/upgrade/upgrade.rs

@@ -6,8 +6,8 @@ use std::path::PathBuf;
 use anyhow::Context as _;
 use cargo_edit::{
     CargoResult, CertsSource, CrateSpec, Dependency, IndexCache, LocalManifest, RustVersion,
-    Source, find_compatible_version, find_latest_version, registry_url, set_dep_version,
-    shell_note, shell_status, shell_warn, shell_write_stdout,
+    Source, find_compatible_version, find_latest_version, registry_token, registry_url,
+    set_dep_version, shell_note, shell_status, shell_warn, shell_write_stdout,
 };
 use clap::Args;
 use indexmap::IndexMap;
@@ -225,6 +225,8 @@ fn exec(args: UpgradeArgs) -> CargoResult<()> {
     let mut pinned_present = false;
     let mut incompatible_present = false;
     let mut uninteresting_crates = BTreeSet::new();
+    let mut token_cache: std::collections::HashMap<String, Option<String>> =
+        std::collections::HashMap::new();
     for (pkg_name, manifest_path, rust_version) in manifests {
         let mut manifest = LocalManifest::try_new(&manifest_path)?;
         let mut crate_modified = false;
@@ -296,15 +298,23 @@ fn exec(args: UpgradeArgs) -> CargoResult<()> {
                     }
                 };
 
-                let (latest_compatible, latest_incompatible) = if dependency
+                let is_registry_dep = dependency
                     .source
                     .as_ref()
-                    .and_then(|s| s.as_registry())
-                    .is_some()
-                {
+                    .map(|s| s.as_registry().is_some())
+                    .unwrap_or(true); // None source means default registry (crates-io)
+                let (latest_compatible, latest_incompatible) = if is_registry_dep {
                     // Update indices for any alternative registries, unless
                     // we're offline.
                     let registry_url = registry_url(&manifest_path, dependency.registry())?;
+                    let cache_key =
+                        format!("{}:{:?}", manifest_path.display(), dependency.registry());
+                    let token = token_cache
+                        .entry(cache_key)
+                        .or_insert_with(|| registry_token(&manifest_path, dependency.registry()));
+                    if let Some(token) = token {
+                        index.set_token(&registry_url, token.clone());
+                    }
                     let krate = index.krate(&registry_url, &dependency.name)?;
                     let versions = krate
                         .as_ref()

src/index.rs

@@ -17,6 +17,7 @@ pub enum CertsSource {
 pub struct IndexCache {
     certs_source: CertsSource,
     index: std::collections::HashMap<Url, AnyIndexCache>,
+    tokens: std::collections::HashMap<String, String>,
 }
 
 impl IndexCache {
@@ -25,9 +26,15 @@ impl IndexCache {
         Self {
             certs_source,
             index: Default::default(),
+            tokens: Default::default(),
         }
     }
 
+    /// Set an auth token for a registry URL
+    pub fn set_token(&mut self, registry_url: &Url, token: String) {
+        self.tokens.insert(registry_url.to_string(), token);
+    }
+
     /// Determines if the specified crate exists in the crates.io index
     #[inline]
     pub fn has_krate(&mut self, registry: &Url, name: &str) -> CargoResult<bool> {
@@ -65,7 +72,8 @@ impl IndexCache {
 
     fn index<'s>(&'s mut self, registry: &Url) -> CargoResult<&'s mut AnyIndexCache> {
         if !self.index.contains_key(registry) {
-            let index = AnyIndex::open(registry, self.certs_source)?;
+            let token = self.tokens.get(&registry.to_string()).cloned();
+            let index = AnyIndex::open(registry, self.certs_source, token)?;
             let index = AnyIndexCache::new(index);
             self.index.insert(registry.clone(), index);
         }
@@ -122,13 +130,13 @@ enum AnyIndex {
 }
 
 impl AnyIndex {
-    fn open(url: &Url, certs_source: CertsSource) -> CargoResult<Self> {
+    fn open(url: &Url, certs_source: CertsSource, token: Option<String>) -> CargoResult<Self> {
         if url.scheme() == "file" {
             LocalIndex::open(url)
                 .map(Self::Local)
                 .with_context(|| format!("invalid local registry {url:?}"))
         } else {
-            RemoteIndex::open(url, certs_source)
+            RemoteIndex::open(url, certs_source, token)
                 .map(Self::Remote)
                 .with_context(|| format!("invalid registry {url:?}"))
         }
@@ -182,10 +190,11 @@ struct RemoteIndex {
     client: tame_index::external::reqwest::blocking::Client,
     lock: FileLock,
     etags: Vec<(String, String)>,
+    token: Option<String>,
 }
 
 impl RemoteIndex {
-    fn open(url: &Url, certs_source: CertsSource) -> CargoResult<Self> {
+    fn open(url: &Url, certs_source: CertsSource, token: Option<String>) -> CargoResult<Self> {
         log::trace!("opening index entry for {url}");
         let url = url.to_string();
         let url = tame_index::IndexUrl::NonCratesIo(std::borrow::Cow::Owned(url));
@@ -209,6 +218,7 @@ impl RemoteIndex {
             client,
             lock,
             etags: Vec::new(),
+            token,
         })
     }
 
@@ -237,6 +247,12 @@ impl RemoteIndex {
         let mut req = self.client.request(method, uri.to_string());
         req = req.version(version);
         req = req.headers(headers);
+        if let Some(token) = &self.token {
+            req = req.header(
+                tame_index::external::reqwest::header::AUTHORIZATION,
+                token.as_str(),
+            );
+        }
         let res = self.client.execute(req.build()?)?;
 
         // Grab the etag if it exists for future requests

src/lib.rs

@@ -44,6 +44,7 @@ pub use fetch::{RustVersion, find_compatible_version, find_latest_version};
 pub use index::*;
 pub use manifest::{LocalManifest, Manifest, find, get_dep_version, set_dep_version};
 pub use metadata::manifest_from_pkgid;
+pub use registry::registry_token;
 pub use registry::registry_url;
 pub use util::{
     Color, ColorChoice, colorize_stderr, shell_note, shell_print, shell_status, shell_warn,