Support additional registries (#942)

Author: illicitonion

crate_universe/defs.bzl

@@ -6,9 +6,10 @@ load("//rust:defs.bzl", "rust_common")
 load("//rust:repositories.bzl", "DEFAULT_TOOLCHAIN_TRIPLES")
 load("//rust/platform:triple_mappings.bzl", "system_to_binary_ext", "triple_to_system")
 
-DEFAULT_CRATE_REGISTRY_TEMPLATE = "https://crates.io/api/v1/crates/{crate}/{version}/download"
+# The default template for download URLs from the default registry (i.e. crates.io).
+DEFAULT_DEFAULT_CRATE_DOWNLOAD_URL_TEMPLATE = "https://crates.io/api/v1/crates/{crate}/{version}/download"
 
-def _input_content_template(ctx, name, packages, cargo_toml_files, overrides, registry_template, targets, cargo_bin_path):
+def _input_content_template(ctx, name, packages, cargo_toml_files, overrides, default_registry_download_url_template, additional_registries, targets, cargo_bin_path):
     """Generate json encoded dependency info for the crate resolver.
 
     Args:
@@ -17,7 +18,8 @@ def _input_content_template(ctx, name, packages, cargo_toml_files, overrides, re
         packages (list): A list of json encoded `crate.spec` entries.
         cargo_toml_files (list): A list of `Label`s to Cargo manifests.
         overrides (dict): A dict of crate name (`str`) to json encoded `crate.override` data.
-        registry_template (str): A crate registry url template
+        default_registry_download_url_template (str): The URL template for downloading crates from the default registry.
+        additional_registries (dict): Dict of registry name (`str`) to the URL of the index (`str`).
         targets (list): A list of target platform triples
         cargo_bin_path (path): The label of a Cargo binary.
 
@@ -43,9 +45,10 @@ def _input_content_template(ctx, name, packages, cargo_toml_files, overrides, re
     return "{}\n".format(
         json.encode_indent(
             struct(
+                additional_registries = additional_registries,
                 cargo = str(cargo_bin_path),
                 cargo_toml_files = encodable_cargo_toml_files,
-                crate_registry_template = registry_template,
+                default_registry_download_url_template = default_registry_download_url_template,
                 overrides = encodable_overrides,
                 packages = dcoded_pkgs,
                 repository_name = name,
@@ -104,7 +107,8 @@ def _crate_universe_resolve_impl(repository_ctx):
         packages = repository_ctx.attr.packages,
         cargo_toml_files = repository_ctx.attr.cargo_toml_files,
         overrides = repository_ctx.attr.overrides,
-        registry_template = repository_ctx.attr.crate_registry_template,
+        default_registry_download_url_template = repository_ctx.attr.default_registry_download_url_template,
+        additional_registries = repository_ctx.attr.additional_registries,
         targets = repository_ctx.attr.supported_targets,
         cargo_bin_path = tools.cargo,
     )
@@ -159,13 +163,20 @@ Environment Variables:
 """,
     implementation = _crate_universe_resolve_impl,
     attrs = {
+        "additional_registries": attr.string_dict(
+            doc = """\
+Additional registries used by Cargo (see https://doc.rust-lang.org/cargo/reference/registries.html).
+
+Dict of registry_name: index_url.
+""",
+        ),
         "cargo_toml_files": attr.label_list(
             doc = "A list of Cargo manifests (`Cargo.toml` files).",
             allow_files = True,
         ),
-        "crate_registry_template": attr.string(
+        "default_registry_download_url_template": attr.string(
             doc = "A template for where to download crates from for the default crate registry. This must contain `{version}` and `{crate}` templates.",
-            default = DEFAULT_CRATE_REGISTRY_TEMPLATE,
+            default = DEFAULT_DEFAULT_CRATE_DOWNLOAD_URL_TEMPLATE,
         ),
         "iso_date": attr.string(
             doc = "The iso_date of cargo binary the resolver should use. Note: This can only be set if `version` is `beta` or `nightly`",

crate_universe/src/config.rs

@@ -52,26 +52,32 @@ pub struct Config {
     pub cargo_toml_files: BTreeMap<String, PathBuf>,
     pub overrides: HashMap<String, Override>,
 
-    /// Template of the URL from which to download crates, which are assumed to be gzip'd tar files.
-    /// This string may contain arbitrarily many instances of {crate} and {version} which will be
-    /// replaced by crate names and versions.
-    pub crate_registry_template: String,
-
     pub target_triples: BTreeSet<String>,
+
     pub cargo: PathBuf,
 
     #[serde(default = "default_rules_rust_workspace_name")]
     pub rust_rules_workspace_name: String,
-    #[serde(default = "default_index_url")]
-    pub index_url: Url,
+
+    #[serde(default = "default_default_registry_index_url")]
+    pub default_registry_index_url: Url,
+    /// Template of the URL from which to download crates, which are assumed to be gzip'd tar files.
+    /// This string may contain arbitrarily many instances of {crate} and {version} which will be
+    /// replaced by crate names and versions.
+    #[serde(default = "default_default_registry_download_url_template")]
+    pub default_registry_download_url_template: String,
+
+    pub additional_registries: BTreeMap<String, Url>,
 }
 
 impl Config {
     pub fn preprocess(mut self) -> anyhow::Result<Resolver> {
         self.packages.sort();
 
+        let known_registries: BTreeSet<_> = self.additional_registries.keys().cloned().collect();
+
         let (toml_contents, label_to_crates) =
-            merge_cargo_tomls(self.cargo_toml_files, self.packages)?;
+            merge_cargo_tomls(&known_registries, self.cargo_toml_files, self.packages)?;
 
         let overrides = self
             .overrides
@@ -97,12 +103,13 @@ impl Config {
             toml_contents.into(),
             ResolverConfig {
                 cargo: self.cargo,
-                index_url: self.index_url,
+                default_registry_index_url: self.default_registry_index_url,
+                default_registry_download_url_template: self.default_registry_download_url_template,
+                additional_registries: self.additional_registries,
             },
             ConsolidatorConfig { overrides },
             RenderConfig {
                 repo_rule_name: self.repository_name.clone(),
-                crate_registry_template: self.crate_registry_template.clone(),
                 rules_rust_workspace_name: self.rust_rules_workspace_name.clone(),
             },
             self.target_triples,
@@ -138,6 +145,10 @@ pub fn default_rules_rust_workspace_name() -> String {
     String::from("rules_rust")
 }
 
-pub fn default_index_url() -> Url {
+fn default_default_registry_index_url() -> Url {
     Url::parse("https://github.com/rust-lang/crates.io-index").expect("Invalid default index URL")
 }
+
+fn default_default_registry_download_url_template() -> String {
+    String::from("https://crates.io/api/v1/crates/{crate}/{version}/download")
+}

crate_universe/src/parser.rs

@@ -50,7 +50,10 @@ pub struct DepSpec {
 
 #[derive(Debug, PartialEq, Eq)]
 pub enum VersionSpec {
-    Semver(VersionReq),
+    Semver {
+        version_req: VersionReq,
+        registry: Option<String>,
+    },
     Git {
         url: String,
         rev: Option<String>,
@@ -83,6 +86,7 @@ struct Ignored {
 }
 
 pub fn merge_cargo_tomls(
+    known_registry_names: &BTreeSet<String>,
     label_to_path: BTreeMap<String, PathBuf>,
     packages: Vec<AdditionalPackage>,
 ) -> anyhow::Result<(CargoToml, BTreeMap<String, BTreeSet<String>>)> {
@@ -116,6 +120,26 @@ pub fn merge_cargo_tomls(
             package: _,
             _ignored,
         } = cargo_toml;
+        for dep_spec in dependencies
+            .values()
+            .chain(build_dependencies.values())
+            .chain(dev_dependencies.values())
+        {
+            if let VersionSpec::Semver {
+                registry: Some(registry),
+                ..
+            } = &dep_spec.version
+            {
+                if !known_registry_names.contains(registry) {
+                    anyhow::bail!(
+                        "Saw dep for unknown registry {} - known registry names: {:?}",
+                        registry,
+                        known_registry_names
+                    );
+                }
+            }
+        }
+
         for (dep, dep_spec) in dependencies.into_iter() {
             if let VersionSpec::Local(_) = dep_spec.version {
                 // We ignore local deps.
@@ -211,12 +235,16 @@ pub fn merge_cargo_tomls(
             DepSpec {
                 default_features: true,
                 features: features.into_iter().collect(),
-                version: VersionSpec::Semver(VersionReq::parse(&semver).with_context(|| {
-                    format!(
-                        "Failed to parse semver requirement for package {}, semver: {}",
-                        name, semver
-                    )
-                })?),
+                version: VersionSpec::Semver {
+                    version_req: VersionReq::parse(&semver).with_context(|| {
+                        format!(
+                            "Failed to parse semver requirement for package {}, semver: {}",
+                            name, semver
+                        )
+                    })?,
+                    // TODO: Support custom registries for additional packages
+                    registry: None,
+                },
             },
         );
     }
@@ -339,8 +367,27 @@ impl DepSpec {
 
         match (&mut self.version, &other.version) {
             (v1, v2) if v1 == v2 => {}
-            (VersionSpec::Semver(v1), VersionSpec::Semver(v2)) => {
-                self.version = VersionSpec::Semver(VersionReq::parse(&format!("{}, {}", v1, v2))?)
+            (
+                VersionSpec::Semver {
+                    version_req: v1,
+                    registry: registry1,
+                },
+                VersionSpec::Semver {
+                    version_req: v2,
+                    registry: registry2,
+                },
+            ) => {
+                if registry1 != registry2 {
+                    return Err(anyhow!(
+                        "Can't merge the same package from different registries (saw registries {:?} and {:?})",
+                        registry1,
+                        registry2,
+                    ));
+                }
+                self.version = VersionSpec::Semver {
+                    version_req: VersionReq::parse(&format!("{}, {}", v1, v2))?,
+                    registry: registry1.clone(),
+                };
             }
             (v1 @ VersionSpec::Git { .. }, v2 @ VersionSpec::Git { .. }) => {
                 return Err(anyhow!(
@@ -377,11 +424,17 @@ impl DepSpec {
             toml::Value::Array(features.into_iter().map(toml::Value::String).collect()),
         );
         match version {
-            VersionSpec::Semver(version) => {
+            VersionSpec::Semver {
+                version_req,
+                registry,
+            } => {
                 v.insert(
                     String::from("version"),
-                    toml::Value::String(format!("{}", version)),
+                    toml::Value::String(format!("{}", version_req)),
                 );
+                if let Some(registry) = registry {
+                    v.insert(String::from("registry"), toml::Value::String(registry));
+                }
             }
             VersionSpec::Git { url, rev, tag } => {
                 v.insert(String::from("git"), toml::Value::String(url));

crate_universe/src/renderer.rs

@@ -17,7 +17,6 @@ use crate::{config, resolver::Dependencies};
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize)]
 pub struct RenderConfig {
     pub repo_rule_name: String,
-    pub crate_registry_template: String,
     pub rules_rust_workspace_name: String,
 }
 
@@ -406,10 +405,6 @@ impl Renderer {
         context.insert("lockfile_hash", &self.context.hash);
         context.insert("crates", &self.context.transitive_renderable_packages);
         context.insert("repo_rule_name", &self.context.config.repo_rule_name);
-        context.insert(
-            "repository_http_template",
-            &self.context.config.crate_registry_template,
-        );
         let rendered_repository_rules = self
             .internal_renderer
             .render("templates/defs.bzl.template", &context)?;
@@ -716,9 +711,6 @@ mod tests {
         let renderer = Renderer::new(
             RenderConfig {
                 repo_rule_name: String::from("rule_prefix"),
-                crate_registry_template: String::from(
-                    "https://crates.io/api/v1/crates/{crate}/{version}/download",
-                ),
                 rules_rust_workspace_name: String::from("rules_rust"),
             },
             String::from("598"),
@@ -892,9 +884,6 @@ mod tests {
     fn default_render_config() -> RenderConfig {
         RenderConfig {
             repo_rule_name: String::from("rule_prefix"),
-            crate_registry_template: String::from(
-                "https://crates.io/api/v1/crates/{crate}/{version}/download",
-            ),
             rules_rust_workspace_name: String::from("rules_rust"),
         }
     }