feat: bzlmod support (#50)

* feat: bzlmod now working

* refactor: don't use index.bzl as we'd need to document a bzlmod argument. Buildifier checks now passing

* chore: remove rules_distroless version from bzlmod example

* fix: lock updating under bzlmod

* fix: resolve transitive dependencies with package lock by default under bzlmod

* chore: buildifier

* fix: get package_name to build package index function name. Now works under bzlmod and workspace

* refactor: cleanup setting lockfile repo under bzlmod

* chore: remove bzlmod example and add it into e2e/smoke

* chore: buildifier

MODULE.bazel

@@ -24,3 +24,4 @@ use_repo(bazel_lib_toolchains, "yq_windows_amd64")
 bazel_dep(name = "gazelle", version = "0.34.0", dev_dependency = True, repo_name = "bazel_gazelle")
 bazel_dep(name = "bazel_skylib_gazelle_plugin", version = "1.5.0", dev_dependency = True)
 bazel_dep(name = "buildifier_prebuilt", version = "6.1.2", dev_dependency = True)
+bazel_dep(name = "platforms", version = "0.0.10", dev_dependency = True)

apt/extensions.bzl

@@ -0,0 +1,46 @@
+"extensions for bzlmod"
+
+load("//apt/private:index.bzl", _deb_package_index = "deb_package_index", _deb_package_index_bzlmod = "deb_package_index_bzlmod")
+load("//apt/private:resolve.bzl", _deb_resolve = "deb_resolve")
+
+deb_index = tag_class(attrs = {
+    "name": attr.string(doc = "Name of the generated repository"),
+    "lock": attr.label(doc = """The lock file to use for the index."""),
+    "manifest": attr.label(doc = """The file used to generate the lock file"""),
+    "resolve_transitive": attr.bool(
+        doc = """Whether dependencies of dependencies should be resolved and added to the lockfile.""",
+        default = True,
+    ),
+})
+
+def _distroless_extension(module_ctx):
+    for mod in module_ctx.modules:
+        for deb_index in mod.tags.deb_index:
+            _deb_package_index_bzlmod(
+                module_ctx = module_ctx,
+                name = deb_index.name,
+                lock = deb_index.lock,
+            )
+
+            _deb_resolve(
+                name = deb_index.name + "_resolution",
+                manifest = deb_index.manifest,
+                resolve_transitive = deb_index.resolve_transitive,
+            )
+
+            if not deb_index.lock:
+                # buildifier: disable=print
+                print("\nNo lockfile was given, please run `bazel run @%s//:lock` to create the lockfile." % deb_index.name)
+
+            _deb_package_index(
+                name = deb_index.name,
+                lock = deb_index.lock if deb_index.lock else "@" + deb_index.name + "_resolution//:lock.json",
+                bzlmod = True,
+            )
+
+apt = module_extension(
+    implementation = _distroless_extension,
+    tag_classes = {
+        "deb_index": deb_index,
+    },
+)

apt/index.bzl

@@ -34,7 +34,13 @@ def deb_index(
     ### BZLMOD Example
 
     ```starlark
-    # TODO: support BZLMOD
+    apt = use_extension("@rules_distroless//apt:extensions.bzl", "apt")
+    apt.deb_index(
+        name = "bullseye",
+        lock = "//examples/apt:bullseye.lock.json",
+        manifest = "//examples/apt:bullseye.yaml",
+    )
+    use_repo(apt, "bullseye")
     ```
 
     ### Macro Expansion
@@ -85,4 +91,5 @@ def deb_index(
         name = name,
         lock = lock if lock else "@" + name + "_resolution//:lock.json",
         package_template = package_template,
+        bzlmod = False,
     )

apt/private/BUILD.bazel

@@ -23,7 +23,11 @@ bzl_library(
     name = "index",
     srcs = ["index.bzl"],
     visibility = ["//apt:__subpackages__"],
-    deps = [":lockfile"],
+    deps = [
+        ":lockfile",
+        "@bazel_tools//tools/build_defs/repo:http.bzl",
+        "@bazel_tools//tools/build_defs/repo:utils.bzl",
+    ],
 )
 
 bzl_library(

apt/private/index.bzl

@@ -1,6 +1,7 @@
 "apt-get"
 
 load(":lockfile.bzl", "lockfile")
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
 _DEB_IMPORT_TMPL = '''\
     http_archive(
@@ -51,13 +52,16 @@ filegroup(
 '''
 
 def _deb_package_index_impl(rctx):
+    # Get the package name with either bzlmod enabled or disabled
+    package_name = rctx.attr.name.split("~")[-1]
+
     package_defs = [
         '"""Generated by rules_distroless. DO NOT EDIT."""',
         """load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")""",
         "",
         "",
         "# buildifier: disable=function-docstring",
-        "def {}_packages():".format(rctx.attr.name),
+        "def {}_packages():".format(package_name),
     ]
 
     lockf = lockfile.from_json(rctx, rctx.attr.lock)
@@ -66,25 +70,28 @@ def _deb_package_index_impl(rctx):
         package_defs.append("   pass")
 
     for (package) in lockf.packages():
-        name = lockfile.make_package_key(
+        package_key = lockfile.make_package_key(
             package["name"],
             package["version"],
             package["arch"],
         )
+
         package_defs.append(
             _DEB_IMPORT_TMPL.format(
-                name = "%s_%s" % (rctx.attr.name, package["key"]),
+                name = "%s_%s" % (rctx.attr.name, package_key),
                 package_name = package["name"],
                 urls = [package["url"]],
                 sha256 = package["sha256"],
             ),
         )
 
+        repo_name = "%s%s_%s" % ("@" if rctx.attr.bzlmod else "", rctx.attr.name, package_key)
+
         rctx.file(
             "%s/%s/BUILD.bazel" % (package["name"], package["arch"]),
             rctx.attr.package_template.format(
                 target_name = package["arch"],
-                src = '"@%s_%s//:data"' % (rctx.attr.name, name),
+                src = '"@%s//:data"' % repo_name,
                 deps = ",\n        ".join([
                     '"//%s/%s"' % (dep["name"], package["arch"])
                     for dep in package["dependencies"]
@@ -93,24 +100,92 @@ def _deb_package_index_impl(rctx):
                 name = package["name"],
                 arch = package["arch"],
                 sha256 = package["sha256"],
-                repo_name = "%s_%s" % (rctx.attr.name, name),
+                repo_name = "%s" % repo_name,
             ),
         )
 
     rctx.file("packages.bzl", "\n".join(package_defs))
     rctx.file("BUILD.bazel", """
 alias(
     name = "lock",
-    actual = "@{}_resolution//:lock",
+    actual = "@{}{}_resolution//:lock",
     tags = ["manual"],
 )
 exports_files(glob(['packages.bzl']))
-""".format(rctx.attr.name))
+""".format("@" if rctx.attr.bzlmod else "", rctx.attr.name))
 
 deb_package_index = repository_rule(
     implementation = _deb_package_index_impl,
     attrs = {
         "lock": attr.label(),
         "package_template": attr.string(default = _PACKAGE_TMPL),
+        "bzlmod": attr.bool(default = False),
     },
 )
+
+def deb_package_index_bzlmod(
+        module_ctx,
+        name,
+        lock,
+        package_template = _PACKAGE_TMPL):
+    """A macro to do everything that `deb_package_index` does, but for bzlmod.
+
+    Args:
+        module_ctx: the module context
+        name: the name of the repository
+        lock: label to a `lock.json`
+        package_template: (EXPERIMENTAL!) a template string for generated BUILD files.
+          Available template replacement keys are: `{target_name}`, `{deps}`, `{urls}`, `{name}`, `{arch}`, `{sha256}`, `{repo_name}`
+    """
+
+    lockf = lockfile.from_json(module_ctx, lock)
+
+    if len(lockf.packages()) < 1:
+        return
+
+    for (package) in lockf.packages():
+        package_key = lockfile.make_package_key(
+            package["name"],
+            package["version"],
+            package["arch"],
+        )
+
+        repo_name = "%s_%s" % (name, package_key)
+
+        http_archive(
+            name = repo_name,
+            urls = [
+                package["url"],
+            ],
+            sha256 = package["sha256"],
+            build_file_content = """\
+filegroup(
+    name = "data",
+    visibility = ["//visibility:public"],
+    srcs = glob(["data.tar.*"]),
+)
+
+filegroup(
+    name = "control",
+    visibility = ["//visibility:public"],
+    srcs = glob(["control.tar.*"]),
+)
+""",
+        )
+
+        module_ctx.file(
+            "%s/%s/BUILD.bazel" % (package["name"], package["arch"]),
+            package_template.format(
+                target_name = package["arch"],
+                src = '"@@%s//:data"' % repo_name,
+                deps = ",\n        ".join([
+                    '"//%s/%s"' % (dep["name"], package["arch"])
+                    for dep in package["dependencies"]
+                ]),
+                urls = [package["url"]],
+                name = package["name"],
+                arch = package["arch"],
+                sha256 = package["sha256"],
+                repo_name = "%s" % repo_name,
+            ),
+        )

apt/private/util.bzl

@@ -11,7 +11,7 @@ def _set_dict(struct, value = None, keys = []):
     struct[keys[-1]] = value
 
 def _sanitize(str):
-    return str.replace("+", "-p-").replace(":", "-")
+    return str.replace("+", "-p-").replace(":", "-").replace("~", "_")
 
 util = struct(
     sanitize = _sanitize,

e2e/smoke/.bazelrc

@@ -0,0 +1 @@
+common --enable_bzlmod

e2e/smoke/.bazelversion

@@ -0,0 +1 @@
+6.4.0