.bazel-patchelf-rules

Rules for invoking patchelf in Bazel

defs.bzl

"""Rules for using [patchelf] in Bazel.

[patchelf]: https://github.com/NixOS/patchelf
"""

load("@rules_sh//sh:sh.bzl", "ShBinariesInfo")
load("//build/bazel/utils:strings.bzl", "error")

################################################################################

# Toolchain:

TOOLCHAIN_TYPE = Label("//build/bazel/rules/patchelf:toolchain_type")

def _patchelf_toolchain_impl(ctx):
    patchelf = ctx.attr.patchelf
    sh_info = patchelf[ShBinariesInfo]

    # Check that there's a single binary named patchelf:
    list_of_executables = sh_info.executables.keys()
    if list_of_executables != ["patchelf"]: error("""
        Expected the `patchelf` attribute to point to an `sh_binaries`
        containing a single executable named `patchelf`; instead it contains: {}
    """, list_of_executables)

    toolchain_info = platform_common.ToolchainInfo(
        sh_info = sh_info,
        # Exposed for runfiles; see:
        # https://github.com/tweag/rules_sh/blob/e7d0ae305694fd3b80534968d42acc0a6142649c/sh/experimental/posix_hermetic.bzl#L171-L177
        tool = patchelf,
        make_variables = platform_common.TemplateVariableInfo({
            "PATCHELF": sh_info.executables["patchelf"].path,
            "PATCHELF_PATH": sh_info.paths.to_list()[0],
        })
    )

    return [toolchain_info]

patchelf_toolchain = rule(
    implementation = _patchelf_toolchain_impl,
    attrs = {
        "patchelf": attr.label(
            executable = True,
            mandatory = True,
            providers = [ShBinariesInfo],
            # Like `bash`, we want a binary that we can run on the execution
            # platform.
            cfg = "exec",
            doc = "TODO",
        ),
    },
    provides = [platform_common.ToolchainInfo],
    doc = "TODO",
)

def _current_patchelf_toolchain_impl(ctx):
    toolchain = ctx.toolchains[TOOLCHAIN_TYPE]

    return [
        toolchain,
        toolchain.make_variables,
        toolchain.sh_info,
        # We cannot re-export `DefaultInfo` verbatim because it contains an
        # executable that was not created by this rule.
        DefaultInfo(
            runfiles = toolchain.tool[DefaultInfo].default_runfiles,
            files = toolchain.tool[DefaultInfo].files,
        )

    ]

# Workaround for https://github.com/bazelbuild/bazel/issues/14009#issuecomment-921960766
current_patchelf_toolchain = rule(
    implementation = _current_patchelf_toolchain_impl,
    toolchains = [TOOLCHAIN_TYPE],
    provides = [
        platform_common.ToolchainInfo,
        platform_common.TemplateVariableInfo,
        ShBinariesInfo,
        DefaultInfo,
    ],
    doc = "TODO",
)

################################################################################

# Rules:

# NOTE: the major caveat here is this only works if the runfiles directory is
# next to the binary!
#
# not sure that this is the case if, for example, the produced binary is invoked
# from another rule (instead of just `run`ing the target)
#
# initial experiments (simple genrules) indicate that this is an okay assumption

# Returns: (generated file, runfiles)
def _apply_patches(ctx, input_file, input_target):
    # NOTE: https://groups.google.com/g/bazel-discuss/c/lom4hWm5wdQ
    # see: https://github.com/tweag/rules_sh/blob/ab645c259bd6d30b3595ace5d632468fcbc79a22/sh/sh.bzl#L246-L255
    patchelf = ctx.toolchains[TOOLCHAIN_TYPE].tool
    tool_inputs, tool_manifests = ctx.resolve_tools(tools = [patchelf])

    out = ctx.actions.declare_file(ctx.attr.name)

    removes = ctx.attr.remove_needed
    adds = ctx.attr.add_needed
    replaces = ctx.attr.replace_needed

    ## Gather the runfiles:
    runfiles = input_target[DefaultInfo].default_runfiles

    # Returns the runfiles and the file:
    def _check_for_single_file_and_get_runfiles(target, attr_name):
        files = target[DefaultInfo].files.to_list()

        okay = True
        if len(files) == 0: okay = False
        elif len(files) == 1: pass
        else:
            # We'll also allow multiple files if the names of the extra files
            # are all strict superset or subset of the first file; i.e.
            # `libfoo.so` and `libfoo.so.1`, `libfoo.so.2`:
            first = files[0]
            okay = all([
                f.path.startswith(first.path) or first.path.startswith(f.path)
                for f in files[1:]
            ])

        if not okay: error("""
            Attrs in `{}` should contain a single library file but target `{}`
            contained {} files: {}.
        """, attr_name, target.label, len(files), files)

        file = files[0]
        runfiles = ctx.runfiles(transitive_files = target[DefaultInfo].files)
        runfiles = runfiles.merge(target[DefaultInfo].default_runfiles)

        return file, runfiles

    def _get_runfiles_relative_path(file):
        # https://bazel.build/extending/rules#runfiles_location
        runfiles_dir = out.basename + ".runfiles"
        workspace_name = ctx.workspace_name
        runfile_path = file.short_path
        return "{}/{}/{}".format(
            runfiles_dir, workspace_name, runfile_path
        )

    def _get_runfiles_relative_rpath_entry(file):
        relative_path = _get_runfiles_relative_path(file)
        # do "dirname":
        dir_path = relative_path.removesuffix("/" + file.basename)

        return "$ORIGIN" + "/" + dir_path

    # Process removes, adds, and replaces:
    args = ctx.actions.args()
    args.add(input_file)
    args.add("--output", out)

    new_runfiles = []

    args.add_all(removes, before_each = "--remove-needed")

    for add in adds:
        file, extra_runfiles = _check_for_single_file_and_get_runfiles(add, "add_needed")
        new_runfiles.append(extra_runfiles)

        rpath = _get_runfiles_relative_rpath_entry(file)
        lib_name = file.basename

        args.add("--add-needed", lib_name)
        args.add("--add-rpath", rpath)

    for new_library_target, orig_library_name in replaces.items():
        file, extra_runfiles = _check_for_single_file_and_get_runfiles(new_library_target, "replace_needed")
        new_runfiles.append(extra_runfiles)

        rpath = _get_runfiles_relative_rpath_entry(file)
        lib_name = file.basename

        args.add_all(["--replace-needed", orig_library_name, lib_name])
        args.add("--add-rpath", rpath)

    runfiles = runfiles.merge_all(new_runfiles)

    ## Invoke!
    if ctx.attr.shrink_rpath:
        args.add("--shrink-rpath")

    args.add_all(ctx.attr.extra_flags)

    ctx.actions.run(
        outputs = [out],
        # TODO: input_file[DefaultInfo].files? runfiles? runfiles of the extra inputs?
        inputs = depset(direct=[input_file], transitive = [tool_inputs] + [runfiles.files]),
        input_manifests = tool_manifests,
        executable = patchelf[ShBinariesInfo].executables["patchelf"],
        arguments = [args],
        mnemonic = "Patchelf",
        progress_message = "Patching {input} → %{output} (%{label})".replace("{input}", input_file.basename),
        use_default_shell_env = False,
        toolchain = TOOLCHAIN_TYPE,
    )

    # TODO: map in all the runfiles so that we can use shrink-rpath effectively!

    return out, runfiles

# TODO: the tricky bit is the file paths for the shared objects
#  - ideally we'd be able to rewrite to use runfiles relative paths...

# all libraries added are added via rpath!
# runfiles relative!
_common_attrs = dict(
    remove_needed = attr.string_list(
        doc = "names of shared objects to remove dependencies on; corresponds to `--remove-needed`",
    ),
    add_needed = attr.label_list(
        allow_files = True,
        cfg = "target",
        doc = "TODO",
    ),
    replace_needed = attr.label_keyed_string_dict(
        allow_files = True,
        cfg = "target",
        doc = "TODO",
    ),
    extra_flags = attr.string_list(
        doc = "TODO",
    ),
    shrink_rpath = attr.bool(default = False, doc = "TODO"),
)

def _patch_binary_impl(ctx):
    out, runfiles = _apply_patches(ctx, ctx.file.binary, ctx.attr.binary)
    return [
        DefaultInfo(
            runfiles = runfiles,
            executable = out,
        )
    ]

patch_binary = rule(
    implementation = _patch_binary_impl,
    attrs = {
        "binary": attr.label(
            executable = True,
            allow_single_file = True,
            mandatory = True,
            cfg = "target",
            doc = "binary to patch",
        ),
    } | _common_attrs,
    executable = True,
    toolchains = [TOOLCHAIN_TYPE],
    doc = "TODO",
)

def _patch_library_impl(ctx):
    _, runfiles = _apply_patches(ctx, ctx.file.library, ctx.attr.library)
    return [
        DefaultInfo(runfiles = runfiles)
    ]

# Same as above but without `executable = True`
patch_library = rule(
    implementation = _patch_library_impl,
    attrs = {
        "library": attr.label(
            allow_single_file = True,
            mandatory = True,
            cfg = "target",
            doc = "shared object to patch",
        ),
    } | _common_attrs,
    toolchains = [TOOLCHAIN_TYPE],
    doc = "TODO",
)

repos.bzl

"""TODO"""

# Patchelf provides binaries; conveniently these are statically linked.
#
# i.e. https://github.com/NixOS/patchelf/releases/tag/0.18.0

load("//build/bazel/utils:strings.bzl", "format")
load("//build/bazel/utils:mod_ext_sibling_repos.bzl", "get_sibling_repo_in_module_extension")
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")

PATCHELF_VER = "0.18.0"

def _bin(name, sha256): return struct(
    name = name,
    url = "https://github.com/NixOS/patchelf/releases/download/0.18.0/" + name.format(ver = PATCHELF_VER),
    sha256 = sha256,
)

# Matches constraint names in `@platforms`; see:
# https://github.com/bazelbuild/platforms
#
# (os, cpu)
BINARIES = {
    ("linux", "aarch64"):  _bin("patchelf-{ver}-aarch64.tar.gz",  "ae13e2effe077e829be759182396b931d8f85cfb9cfe9d49385516ea367ef7b2"),
    ("linux", "armv7"):    _bin("patchelf-{ver}-armv7l.tar.gz",   "6817cc3f55811eb6631cfb0c23b667d0b9693ec81b02f0e7685411c4aa70e555"),
    ("linux", "x86_32"):   _bin("patchelf-{ver}-i686.tar.gz",     "809371067871a4237ae25e5b4efcf93803e53ed1b5fb64d5a064ebdcc46c6bb7"),
    ("linux", "ppc"):      _bin("patchelf-{ver}-ppc64le.tar.gz",  "b01501473ed7652d9334a21af0416aaa661db0b1aa42bb2900409daa7c0acc0b"),
    ("linux", "riscv64"):  _bin("patchelf-{ver}-riscv64.tar.gz",  "1ea3e669c0fa26eacde205a24d49f9b7b921ce3445c33421c6eb66d24580ed15"),
    ("linux", "s390x"):    _bin("patchelf-{ver}-s390x.tar.gz",    "5eb9bb58f0eabace75af463803ab32d43c32480d601d4b015411282cd9aaa6f3"),
    ("linux", "x86_64"):   _bin("patchelf-{ver}-x86_64.tar.gz",   "ce84f2447fb7a8679e58bc54a20dc2b01b37b5802e12c57eece772a6f14bf3f0"),
    ("windows", "x86_32"): _bin("patchelf-win32-{ver}.exe",       "2db93a71658089841e48d19ccde3bcfed3b2945753ce40f10cf9886b2c00773e"),
    ("windows", "x86_64"): _bin("patchelf-win64-{ver}.exe",       "ea5293833b6a547612ce4b073ac84fd603011ce3455f488a1017fabc8bd170ff"),
}

def _windows_toolchain_repo_impl(rctx):
    rctx.download(
        url = rctx.attr.url,
        output = "patchelf.exe",
        sha256 = rctx.attr.sha256,
        executable = True,
        canonical_id = rctx.attr.url,
    )
    rctx.file("BUILD.bazel", executable = False, contents = format("""
        load("@rules_sh//sh:sh.bzl", "sh_binaries")
        load("@//build/bazel/rules/patchelf:defs.bzl", "patchelf_toolchain")

        sh_binaries(
            name = "patchelf",
            srcs = [":patchelf.exe"],
            target_compatible_with = [
                "@platforms//os:{os}",
                "@platforms//cpu:{cpu}",
            ],
            visibility = ["//visibility:public"],
        )

        patchelf_toolchain(
            name = "patchelf_toolchain",
            patchelf = ":patchelf",
        )
    """, os = rctx.attr.os, cpu = rctx.attr.cpu))

_windows_toolchain_repo = repository_rule(
    implementation = _windows_toolchain_repo_impl,
    attrs = {
        "os": attr.string(),
        "arch": attr.string(),
        "canonical_id": attr.string(),
        "url": attr.string(),
        "sha256": attr.string(),
    },
)

def _toolchain_repo_impl(rctx):
    info = json.decode(rctx.attr.json_encoded_info)
    build_file = format("""
        load("@//build/bazel/rules/patchelf:defs.bzl", "TOOLCHAIN_TYPE")

        package(default_visibility = ["//visibility:public"])
    """)

    for repo_name, attrs in info.items():
        build_file += format("""
            toolchain(
                name = "{name}",
                toolchain = "{label}",
                toolchain_type = TOOLCHAIN_TYPE,
                exec_compatible_with = {constraints},
            )
        """,
            name = repo_name,
            constraints = attrs.constraints,
            label = "@@{}//:patchelf_toolchain".format(
                get_sibling_repo_in_module_extension(rctx, repo_name)
            ),
        )

    rctx.file("BUILD.bazel", executable = False, contents = build_file)


_toolchain_repo = repository_rule(
    implementation = _toolchain_repo_impl,
    attrs = {
        "json_encoded_info": attr.string(),
    },
)

# Creates toolchain repos for patchelf and a top-level repo with the actual
# `toolchain` definitions (you can use `@<name>//:all` when registering these
# toolchains).
def create_toolchain_repos(name = "patchelf_toolchains"):
    info = {}

    # We create separate repos for each platform to take advantage of lazy
    # repo fetching:
    for (os, arch), attrs in BINARIES.items():
        common = dict(
            name = "patchelf_{}_{}_toolchain".format(os, arch),
            canonical_id = "patchelf-{}-{}-{}".format(PATCHELF_VER, os, arch),
            sha256 = attrs.sha256,
            url = attrs.url,
        )

        if os == "linux":
            http_archive(
                build_file_content = format("""
                    load("@rules_sh//sh:sh.bzl", "sh_binaries")
                    load(
                        "@//build/bazel/rules/patchelf:defs.bzl",
                        "patchelf_toolchain"
                    )

                    sh_binaries(
                        name = "patchelf",
                        srcs = [":bin/patchelf"],
                        target_compatible_with = [
                            "@platforms//os:{os}",
                            "@platforms//cpu:{arch}",
                        ],
                        visibility = ["//visibility:public"],
                    )

                    patchelf_toolchain(
                        name = "patchelf_toolchain",
                        patchelf = ":patchelf",
                    )
                """, os = os, arch = arch),
                **common,
            )
        elif os == "windows":
            _windows_toolchain_repo(os = os, arch = arch, **common)
        else: fail("unknown os: ", os)

        info[common["name"]] = dict(
            constraints = [
                "@platforms//os:{}".format(os),
                "@platforms//cpu:{}".format(arch),
            ],
            os = os,
            arch = arch,
        )

    _toolchain_repo(name = name, json_encoded_info = json.encode(info))

    return info


# TODO: module extension?
# with the same optional toolchain registration thing