From 999d6534046ea23c7b2a7effe412d9b604f58257 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Mon, 1 Apr 2024 17:57:13 -0400
Subject: [PATCH] Use canonical URL to key redirect map (#2764)

## Summary

This fixes a potential bug that revealed itself in
https://github.com/astral-sh/uv/pull/2761. We don't run into this now
because we always use "allowed URLs", stores the "last" compatible URL
in the map. But the use of the "raw" URL (rather than the "canonical"
URL) means that other writers have to follow that same assumption and
iterate over dependencies in-order.
---
 crates/uv-resolver/src/resolution.rs     |  7 ++++---
 crates/uv-resolver/src/resolver/index.rs |  3 ++-
 crates/uv-resolver/src/resolver/mod.rs   | 13 ++++++++++---
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/crates/uv-resolver/src/resolution.rs b/crates/uv-resolver/src/resolution.rs
index 7edf9c4ef..95e44d818 100644
--- a/crates/uv-resolver/src/resolution.rs
+++ b/crates/uv-resolver/src/resolution.rs
@@ -2,6 +2,7 @@ use std::borrow::Cow;
 use std::hash::BuildHasherDefault;
 
 use anyhow::Result;
+use cache_key::CanonicalUrl;
 use dashmap::DashMap;
 use itertools::Itertools;
 use owo_colors::OwoColorize;
@@ -68,7 +69,7 @@ impl ResolutionGraph {
         pins: &FilePins,
         packages: &OnceMap<PackageName, VersionsResponse>,
         distributions: &OnceMap<PackageId, Metadata23>,
-        redirects: &DashMap<Url, Url>,
+        redirects: &DashMap<CanonicalUrl, Url>,
         state: &State<UvDependencyProvider>,
         preferences: &Preferences,
         editables: Editables,
@@ -119,7 +120,7 @@ impl ResolutionGraph {
                     let pinned_package = if let Some((editable, _)) = editables.get(package_name) {
                         Dist::from_editable(package_name.clone(), editable.clone())?
                     } else {
-                        let url = redirects.get(url).map_or_else(
+                        let url = redirects.get(&CanonicalUrl::new(url)).map_or_else(
                             || url.clone(),
                             |precise| apply_redirect(url, precise.value()),
                         );
@@ -224,7 +225,7 @@ impl ResolutionGraph {
                                 .or_insert_with(Vec::new)
                                 .push(extra.clone());
                         } else {
-                            let url = redirects.get(url).map_or_else(
+                            let url = redirects.get(&CanonicalUrl::new(url)).map_or_else(
                                 || url.clone(),
                                 |precise| apply_redirect(url, precise.value()),
                             );
diff --git a/crates/uv-resolver/src/resolver/index.rs b/crates/uv-resolver/src/resolver/index.rs
index 9a2f45914..64692032a 100644
--- a/crates/uv-resolver/src/resolver/index.rs
+++ b/crates/uv-resolver/src/resolver/index.rs
@@ -1,3 +1,4 @@
+use cache_key::CanonicalUrl;
 use dashmap::DashMap;
 use url::Url;
 
@@ -21,5 +22,5 @@ pub struct InMemoryIndex {
     /// A map from source URL to precise URL. For example, the source URL
     /// `git+https://github.com/pallets/flask.git` could be redirected to
     /// `git+https://github.com/pallets/flask.git@c2f65dd1cfff0672b902fd5b30815f0b4137214c`.
-    pub(crate) redirects: DashMap<Url, Url>,
+    pub(crate) redirects: DashMap<CanonicalUrl, Url>,
 }
diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs
index 8f0973dd1..0fc4ce11c 100644
--- a/crates/uv-resolver/src/resolver/mod.rs
+++ b/crates/uv-resolver/src/resolver/mod.rs
@@ -6,6 +6,7 @@ use std::ops::Deref;
 use std::sync::Arc;
 
 use anyhow::Result;
+use cache_key::CanonicalUrl;
 use dashmap::{DashMap, DashSet};
 use futures::{FutureExt, StreamExt};
 use itertools::Itertools;
@@ -949,13 +950,19 @@ impl<
                     if let Some(precise) = precise {
                         match distribution {
                             SourceDist::DirectUrl(sdist) => {
-                                self.index.redirects.insert(sdist.url.to_url(), precise);
+                                self.index
+                                    .redirects
+                                    .insert(CanonicalUrl::new(&sdist.url), precise);
                             }
                             SourceDist::Git(sdist) => {
-                                self.index.redirects.insert(sdist.url.to_url(), precise);
+                                self.index
+                                    .redirects
+                                    .insert(CanonicalUrl::new(&sdist.url), precise);
                             }
                             SourceDist::Path(sdist) => {
-                                self.index.redirects.insert(sdist.url.to_url(), precise);
+                                self.index
+                                    .redirects
+                                    .insert(CanonicalUrl::new(&sdist.url), precise);
                             }
                             SourceDist::Registry(_) => {}
                         }