Buck2 OCI rules sketch

BUCK

# Root BUCK file

http_archive(
  name="dive",
  urls=["https://github.com/wagoodman/dive/releases/download/v0.12.0/dive_0.12.0_linux_arm64.tar.gz"],
  sha256="a2a1470302cdfa367a48f80b67bbf11c0cd8039af9211e39515bd2bbbda58fea",
  strip_prefix='dive',
  sub_targets=["dive"],
)

tar_file(
  name = "dive-tar",
  srcs = [":dive[dive]"],
  out = "dive.tar",
)

oci_pull(
  name = "distroless",
  image = "gcr.io/distroless/base",
  digest = "sha256:ccaef5ee2f1850270d453fdf700a5392534f8d1a8ca2acda391fbb6a06b81c86",
  platforms = ["linux/arm64"],
)

oci_image(
  name = "image",
  base = ":distroless",
  tars = [":dive-tar"],
  entrypoint = ["/dive"],

image.py

# prelude-replay/oci/image.py

import argparse
from io import BufferedRandom
import subprocess
import sys
import tempfile

REGISTRY_PORT = 61978

def start_registry(crane_path: str, log_file: BufferedRandom):
    """Starts a local crane registry and logs its output."""
    log = log_file
    registry_process = subprocess.Popen([crane_path, "registry", "serve", "--address", ":{}".format(REGISTRY_PORT)], stdout=log, stderr=log)
    return registry_process

def stop_registry(registry_process):
    """Stops the local crane registry."""
    registry_process.terminate()
    registry_process.wait()

def build_image(crane_path, base_image_path, tar_files, entrypoint, output, name):
    # get last part of base_image path
    base_image = base_image_path.split("/")[-1]
    fully_qualified_base_image = f"localhost:{REGISTRY_PORT}/{base_image}"
    push_base_image_command = [crane_path, 'push', base_image_path, fully_qualified_base_image]
    print(f"Pushing base image: {push_base_image_command}")
    subprocess.run(push_base_image_command, check=True)

    fully_qualified_new_image = f"localhost:{REGISTRY_PORT}/{name}"
    append_layer_command = [crane_path, 'append', '-t', fully_qualified_new_image, '-f', ",".join(tar_files), '-b', fully_qualified_base_image]
    print(f"Appending layers: {append_layer_command}")
    subprocess.run(append_layer_command, check=True)

    # Set the entrypoint if provided
    if entrypoint:
        config_command = [crane_path, 'mutate', '--entrypoint', entrypoint, fully_qualified_new_image, '-o', output]
        print(f"Setting entrypoint: {config_command}")
        subprocess.run(config_command, check=True)

    # Save the final image to a tar file
    # export_command = [crane_path, 'export', fully_qualified_new_image, output]
    # print(f"Exporting image: {export_command}")
    # subprocess.run(export_command, check=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Build OCI image using Crane")
    parser.add_argument("--crane", required=True, help="Path to the crane binary")
    parser.add_argument("--base", required=True, help="Base OCI image")
    parser.add_argument("--tars", nargs='+', required=True, help="Paths to tar files representing layers")
    parser.add_argument("--entrypoint", nargs='+', help="Entrypoint for the OCI image")
    parser.add_argument("--output", required=True, help="Path to the output tar file")
    parser.add_argument("--name", required=True, help="Name of the OCI image")
    args = parser.parse_args()

    log_file = tempfile.TemporaryFile()
    registry_process = None

    try:
        registry_process = start_registry(args.crane, log_file)
        build_image(args.crane, args.base, args.tars, ' '.join(args.entrypoint) if args.entrypoint else None, args.output, args.name)
    except subprocess.CalledProcessError as e:
        print(f"Error: {e}", file=sys.stderr)
    finally:
        if registry_process:
            stop_registry(registry_process)
            print(log_file.read())
                

oci.BUCK

# prelude-replay/oci/BUCK
export_file(
  name = "pull.py",
  src = "pull.py",
  visibility = ["PUBLIC"],
)

export_file(
  name = "image.py",
  src = "image.py",
  visibility = ["PUBLIC"],
)

oci.bzl

load("@prelude//python:toolchain.bzl", "PythonToolchainInfo")
load("//oci:toolchain.bzl", "OciToolchainInfo")


def oci_pull_impl(ctx: AnalysisContext) -> list[Provider]:
  image = ctx.attrs.image
  output = ctx.actions.declare_output("{}.tar".format(ctx.attrs.name))
  platform = ctx.attrs.platforms[0]

  python = ctx.attrs._python_toolchain[PythonToolchainInfo].interpreter
  pull_py = ctx.attrs._oci_toolchain[OciToolchainInfo].pull_py[DefaultInfo].default_outputs
  crane = ctx.attrs._oci_toolchain[OciToolchainInfo].crane[RunInfo]

  cmd = cmd_args(
    python,
    pull_py,
    "--crane",
    crane,
    "--platform",
    platform,
    "--output", output.as_output(),
    "--image",
    image,
    "--digest",
    ctx.attrs.digest,
  )

  ctx.actions.run(cmd, category = "oci")

  return [DefaultInfo(default_output = output)]


oci_pull = rule(
  impl = oci_pull_impl,
  attrs = {
    "digest": attrs.string(),
    "image": attrs.string(),
    "platforms": attrs.list(attrs.string()),
    "_python_toolchain": attrs.toolchain_dep(
      default = "toolchains//:python",
      providers = [PythonToolchainInfo]
    ),
    "_oci_toolchain": attrs.toolchain_dep(
      default = "toolchains//:oci",
      providers = [OciToolchainInfo]
    ),
  },
)

def oci_image_impl(ctx: AnalysisContext) -> list[Provider]:
  base = ctx.attrs.base[DefaultInfo].default_outputs
  # map all of these tars to their outputs using a list comprehension
  tars = ctx.attrs.tars
  tar_outputs = [tar[DefaultInfo].default_outputs for tar in tars]
  entrypoint = ctx.attrs.entrypoint

  output = ctx.actions.declare_output("{}.tar".format(ctx.attrs.name))
  python = ctx.attrs._python_toolchain[PythonToolchainInfo].interpreter
  image_py = ctx.attrs._oci_toolchain[OciToolchainInfo].image_py[DefaultInfo].default_outputs
  crane = ctx.attrs._oci_toolchain[OciToolchainInfo].crane[RunInfo]

  cmd = cmd_args(
    python,
    image_py,
    "--crane",
    crane,
    "--output", output.as_output(),
    "--base", base,
    "--tars", tar_outputs,
    "--entrypoint", cmd_args(entrypoint, delimiter = ","),
    "--name", ctx.attrs.name,
  )

  ctx.actions.run(cmd, category = "oci")

  return [DefaultInfo(default_output = output)]

oci_image = rule(
  impl = oci_image_impl,
  attrs = {
    "base": attrs.dep(),
    "tars": attrs.list(attrs.dep()),
    # TODO(dmiller): I'm not sure if this data type is correct
    "entrypoint": attrs.option(attrs.list(attrs.string())),
    "_python_toolchain": attrs.toolchain_dep(
      default = "toolchains//:python",
      providers = [PythonToolchainInfo]
    ),
    "_oci_toolchain": attrs.toolchain_dep(
      default = "toolchains//:oci",
      providers = [OciToolchainInfo]
    ),
  }
)

pull.py

# prelude-replay/oci/pull.py

import argparse
import subprocess
import sys

def pull_image(crane_path, image, digest, platform, output):
    full_image = f"{image}@{digest}"

    # Construct and execute the crane pull command
    command = [crane_path, 'pull', '--platform', platform, full_image, output]
    subprocess.run(command, check=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Pull OCI image using Crane")
    parser.add_argument("--crane", required=True, help="Path to the crane binary")
    parser.add_argument("--image", required=True, help="OCI image to pull")
    parser.add_argument("--digest", required=True, help="Digest of the OCI image")
    parser.add_argument("--platform", required=True, help="Platform for which to pull the image")
    parser.add_argument("--output", required=True, help="Path to the output tar file")

    args = parser.parse_args()

    try:
        pull_image(args.crane, args.image, args.digest, args.platform, args.output)
    except subprocess.CalledProcessError as e:
        print(f"Error pulling image: {e}", file=sys.stderr)
        sys.exit(1)

releases.bzl

# prelude-replay/oci/releases.bzl
releases = {
  "0.19.0": {
    "Darwin-arm64": {
      "url": "https://github.com/google/go-containerregistry/releases/download/v0.19.0/go-containerregistry_Darwin_arm64.tar.gz",
      "sha256": "0d34767cebd3c2e2e3e227ba301bd867f951ebd85b6079aca93b326152ff667f",
    }
  }
}

toolchain.bzl

# prelude-replay/oci/toolchain.bzl
load(
  "@prelude//rules.bzl",
  "http_archive",
)

load(
  ":crane_releases.bzl",
  crane_releases = "releases",
)

CraneReleaseInfo = provider(
  fields = {
    "version": provider_field(typing.Any, default = None),
    "url": provider_field(typing.Any, default = None),
    "sha256": provider_field(typing.Any, default = None),
  }
)

def _get_crane_release(version: str, platform: str) -> CraneReleaseInfo:
  if version not in crane_releases:
    fail("Unsupported crane version: {}".format(version))
  if platform not in crane_releases[version]:
    fail("Unsupported crane platform: {}".format(platform))

  return CraneReleaseInfo(
    version = version,
    url = crane_releases[version][platform]["url"],
    sha256 = crane_releases[version][platform]["sha256"],
  )

def _host_arch() -> str:
  arch = host_info().arch
  if arch.is_x86_64:
      return "x86_64"
  elif host_info().arch.is_aarch64:
      return "arm64"
  else:
      fail("Unsupported host architecture.")

def _host_os() -> str:
  os = host_info().os
  if os.is_linux:
      return "Linux"
  elif os.is_macos:
      return "Darwin"
  elif os.is_windows:
      return "Windows"
  else:
      fail("Unsupported host os.")

CraneInfo = provider(
  fields = {
    "version": provider_field(typing.Any, default = None),
    "arch": provider_field(typing.Any, default = None),
    "os": provider_field(typing.Any, default = None),
  }
)


def _crane_binary_impl(ctx: AnalysisContext) -> list[Provider]:
  dst = ctx.actions.declare_output("crane")
  src = ctx.attrs.bin[DefaultInfo].default_outputs[0]
  ctx.actions.run(["cp", cmd_args(src, format="{}/crane"), dst.as_output()], category = "cp_crane")

  crane = cmd_args([dst])
  crane.hidden()
  crane.hidden(ctx.attrs.bin[DefaultInfo].default_outputs)
  crane.hidden(ctx.attrs.bin[DefaultInfo].other_outputs)

  return [
    ctx.attrs.bin[DefaultInfo],
    RunInfo(args = crane),
    CraneInfo(
      version = ctx.attrs.version,
      arch = ctx.attrs.arch,
      os = ctx.attrs.os,
    )
  ]


crane_binary = rule(
  impl = _crane_binary_impl,
  attrs = {
    "version": attrs.string(),
    "arch": attrs.string(),
    "os": attrs.string(),
    "bin": attrs.dep(providers = [DefaultInfo]),
  }
)

def download_crane_binary(
    name: str,
    version: str,
    arch: [None, str] = None,
    os: [None, str] = None):
  if arch == None:
    arch = _host_arch()
  if os == None:
    os = _host_os()

  archive_name = "go_containerregistry_{}_{}.tar.gz".format(os, arch)
  release = _get_crane_release(version, "{}-{}".format(os, arch))
  http_archive(
    name = archive_name,
    urls = [release.url],
    sha256 = release.sha256,
  )
  crane_binary(
    name = name,
    version = version,
    arch = arch,
    os = os,
    bin = ":{}".format(archive_name)
  )

OciToolchainInfo = provider(fields = {
  "crane": typing.Any,
  "pull_py": typing.Any,
  "image_py": typing.Any,
})

def oci_toolchain_impl(ctx) -> list[[DefaultInfo, OciToolchainInfo]]:
  return [
    DefaultInfo(),
    OciToolchainInfo(
      crane=ctx.attrs.crane,
      pull_py=ctx.attrs._pull_py, 
      image_py = ctx.attrs._image_py,
    )
  ]

oci_toolchain = rule(
  impl = oci_toolchain_impl,
  attrs = {
    "crane": attrs.exec_dep(providers = [RunInfo]),
    "_pull_py": attrs.dep(
      default = "prelude-replay//oci:pull.py",
    ),
    "_image_py": attrs.dep(
      default = "prelude-replay//oci:image.py",
    ),
  },
  is_toolchain_rule = True,
)

toolchains.BUCK

# toolchains/BUCK
load("@prelude-replay//oci:toolchain.bzl", "oci_toolchain", "download_crane_binary")
download_crane_binary(
    name = "crane-0.19.0",
    version = "0.19.0",
)
oci_toolchain(
    name = "oci",
    crane = ":crane-0.19.0",
    visibility = ["PUBLIC"],
)