Support cargo_bazel_bootstrap for bzlmod

MODULE.bazel

@@ -20,8 +20,15 @@ use_repo(
     "rules_rust_tinyjson",
 )
 
-# This is the default host tools configuration, but if a user defines it in
-# their repo, it's overridden.
 rust = use_extension("//rust:extensions.bzl", "rust")
-rust.host_tools(edition = "2021")
+
+# Allow us to run, for example, "bazel build //tools/runfiles" with bzlmod.
+# Register it as a dev dependency so that we don't force this toolchain on
+# downstream users.
+rust.toolchain(edition = "2021")
+use_repo(rust, "rust_toolchains")
+register_toolchains("@rust_toolchains//:all", dev_dependency = True)
+
 use_repo(rust, "rust_host_tools")
+cargo_bazel_bootstrap = use_extension("//crate_universe/private/module_extensions:cargo_bazel_bootstrap.bzl", "cargo_bazel_bootstrap")
+use_repo(cargo_bazel_bootstrap, "cargo_bazel_bootstrap")

crate_universe/deps_bootstrap.bzl

@@ -6,12 +6,13 @@ load("//crate_universe/private:srcs.bzl", "CARGO_BAZEL_SRCS")
 # buildifier: disable=bzl-visibility
 load("//rust/private:common.bzl", "rust_common")
 
-def cargo_bazel_bootstrap(name = "cargo_bazel_bootstrap", rust_version = rust_common.default_version):
+def cargo_bazel_bootstrap(name = "cargo_bazel_bootstrap", rust_version = rust_common.default_version, **kwargs):
     """An optional repository which bootstraps `cargo-bazel` for use with `crates_repository`
 
     Args:
         name (str, optional): The name of the `cargo_bootstrap_repository`.
         rust_version (str, optional): The rust version to use. Defaults to the default of `cargo_bootstrap_repository`.
+        **kwargs: kwargs to pass through to cargo_bootstrap_repository.
     """
     cargo_bootstrap_repository(
         name = name,
@@ -22,4 +23,5 @@ def cargo_bazel_bootstrap(name = "cargo_bazel_bootstrap", rust_version = rust_co
         version = rust_version,
         # The increased timeout helps avoid flakes in CI
         timeout = 900,
+        **kwargs
     )

crate_universe/private/module_extensions/cargo_bazel_bootstrap.bzl

@@ -0,0 +1,14 @@
+"""Module extension for bootstrapping cargo-bazel."""
+
+load("//crate_universe:deps_bootstrap.bzl", _cargo_bazel_bootstrap_repo_rule = "cargo_bazel_bootstrap")
+
+def _cargo_bazel_bootstrap_impl(_):
+    _cargo_bazel_bootstrap_repo_rule(
+        rust_toolchain_cargo_template = "@rust_host_tools//:bin/{tool}",
+        rust_toolchain_rustc_template = "@rust_host_tools//:bin/{tool}",
+    )
+
+cargo_bazel_bootstrap = module_extension(
+    implementation = _cargo_bazel_bootstrap_impl,
+    doc = """Module extension to generate the cargo_bazel binary.""",
+)

rust/extensions.bzl

@@ -10,30 +10,90 @@ load(
     "DEFAULT_STATIC_RUST_URL_TEMPLATES",
 )
 
+_HOST_TOOL_ERR = """When %s, host tools must be explicitly defined. For example:
+
+rust = use_extension("@rules_rust//rust:extensions.bzl", "rust")
+rust.host_tools(
+    edition = "2021",
+    version = "1.70.2",
+)
+"""
+
+_EXAMPLE_TOOLCHAIN = """
+rust = use_extension("@rules_rust//rust:extensions.bzl", "rust")
+rust.toolchain(
+    edition = "2021",
+    versions = ["1.70.2"],
+)
+use_repo(rust, "rust_toolchains")
+register_toolchains("@rust_toolchains//:all")"""
+
+_TRANSITIVE_DEP_ERR = """
+Your transitive dependency %s is using rules_rust, so you need to define a rust toolchain.
+To do so, you will need to add the following to your root MODULE.bazel. For example:
+
+bazel_dep(name = "rules_rust", version = "<rules_rust version>")
+""" + _EXAMPLE_TOOLCHAIN
+
+_TOOLCHAIN_ERR = """
+Please add at least one toolchain to your root MODULE.bazel. For example:
+""" + _EXAMPLE_TOOLCHAIN
+
 def _rust_impl(module_ctx):
-    # Allow the root module to define host tools. Otherwise, we'll fall back to
-    # the one defined in rules_rust.
+    # Toolchain configuration is only allowed in the root module.
+    # It would be very confusing (and a security concern) if I was using the
+    # default rust toolchains, then when I added a module built on rust, I was
+    # suddenly using a custom rustc.
+    root = None
     for mod in module_ctx.modules:
-        if mod.tags.host_tools:
-            host_tools = mod.tags.host_tools[0]
-            host_triple = get_host_triple(module_ctx)
-
-            rust_toolchain_tools_repository(
-                name = "rust_host_tools",
-                exec_triple = host_triple.str,
-                target_triple = host_triple.str,
-                allocator_library = host_tools.allocator_library,
-                dev_components = host_tools.dev_components,
-                edition = host_tools.edition,
-                rustfmt_version = host_tools.rustfmt_version,
-                sha256s = host_tools.sha256s,
-                urls = host_tools.urls,
-                version = host_tools.version,
-            )
-            break
-
-    mod = module_ctx.modules[0]
-    for toolchain in mod.tags.toolchain:
+        if mod.is_root:
+            root = mod
+    if not root:
+        fail(_TRANSITIVE_DEP_ERR % module_ctx.modules[0].name)
+
+    toolchains = root.tags.toolchain
+    if not toolchains:
+        fail(_TOOLCHAIN_ERR)
+
+    if len(root.tags.host_tools) == 1:
+        host_tools = root.tags.host_tools[0]
+    elif not root.tags.host_tools:
+        if len(toolchains) != 1:
+            fail(_HOST_TOOL_ERR % "multiple toolchains are provided")
+        toolchain = toolchains[0]
+        if len(toolchain.versions) == 1:
+            version = toolchain.versions[0]
+        elif not toolchain.versions:
+            version = None
+        else:
+            fail(_HOST_TOOL_ERR % "multiple toolchain versions are provided")
+        host_tools = struct(
+            allocator_library = toolchain.allocator_library,
+            dev_components = toolchain.dev_components,
+            edition = toolchain.edition,
+            rustfmt_version = toolchain.rustfmt_version,
+            sha256s = toolchain.sha256s,
+            urls = toolchain.urls,
+            version = version,
+        )
+    else:
+        fail("Multiple host_tools were defined in your root MODULE.bazel")
+
+    host_triple = get_host_triple(module_ctx)
+    rust_toolchain_tools_repository(
+        name = "rust_host_tools",
+        exec_triple = host_triple.str,
+        target_triple = host_triple.str,
+        allocator_library = host_tools.allocator_library,
+        dev_components = host_tools.dev_components,
+        edition = host_tools.edition,
+        rustfmt_version = host_tools.rustfmt_version,
+        sha256s = host_tools.sha256s,
+        urls = host_tools.urls,
+        version = host_tools.version or rust_common.default_version,
+    )
+
+    for toolchain in toolchains:
         rust_register_toolchains(
             dev_components = toolchain.dev_components,
             edition = toolchain.edition,
@@ -64,7 +124,7 @@ _RUST_TOOLCHAIN_TAG = tag_class(attrs = dict(
 ))
 
 _RUST_HOST_TOOLS_TAG = tag_class(attrs = dict(
-    version = attr.string(default = rust_common.default_version),
+    version = attr.string(),
     **_COMMON_TAG_KWARGS
 ))