gcr_repo_digest.bzl

load("@io_bazel_rules_docker//container:providers.bzl", "BundleInfo", "ImageInfo")
load("//deployment/docker:image_name.bzl", "image_name")
load(":providers.bzl", "DigestInfo", "RepoInfo")

def _bundle_impl(target, ctx):
    bundle_info = target[BundleInfo]
    return RepoInfo(image_digests = {
        image_name(bundle_info, x[ImageInfo]): x[DigestInfo]
        for x in ctx.rule.attr.image_targets
    })

def _image_impl(target, ctx):
    image_info = target[ImageInfo]
    container_parts = image_info.container_parts

    args = ctx.actions.args()
    args.add("--config", container_parts["config"])
    args.add("--config-digest", container_parts["config_digest"])
    args.add_all(container_parts["blobsum"], before_each = "--blobsum")
    args.add_all(container_parts["zipped_layer"], before_each = "--zipped-layer")

    manifest = ctx.actions.declare_file("{}.manifest.json".format(ctx.rule.attr.name))
    args.add("--manifest", manifest)

    digest = ctx.actions.declare_file("{}.sha256".format(ctx.rule.attr.name))
    args.add("--digest", digest)

    inputs = depset(
        direct = [container_parts["config"], container_parts["config_digest"]],
        transitive = [depset(container_parts["blobsum"]), depset(container_parts["zipped_layer"])],
    )

    ctx.actions.run(
        executable = ctx.executable._gcr_repo_digest,
        arguments = [args],
        inputs = inputs,
        outputs = [manifest, digest],
    )

    return [
        DigestInfo(
            digest_path = digest,
        ),
    ]

def _impl(target, ctx):
    if ImageInfo in target:
        return _image_impl(target, ctx)
    elif BundleInfo in target:
        return _bundle_impl(target, ctx)
    else:
        fail("Unexpected type {}".format(target))

gcr_repo_digest = aspect(
    implementation = _impl,
    attr_aspects = ["image_targets"],  # traverse from container_bundle to container_image
    attrs = {
        "_gcr_repo_digest": attr.label(
            executable = True,
            cfg = "exec",
            default = "gcr_repo_digest",
        ),
    },
    doc = """Attaches a GCR repository digest to an image, or to all images in a bundle.

    container_bundle targets gain a `RepoInfo` provider containing a `image name -> DigestInfo` mapping.
    container_image targets gain a `DigestInfo` provider containing a path to a file containing the repo digest.
    """,
)

gcr_repo_digest.py

#!/usr/bin/env python3
"""
Calculate an image RepoDigest by generating a repository manifest and hashing it.

NOTE: this only works for gcr.io! since we're only pushing to gcr that's fine.

The difference between digests created on gcr.io vs docker.io is the json formatting...
If you need a docker.io repo hash set `indent=3` and `separators=None`.

The best tool I've found to inspect these is `gcrane`:

```console
$ go get github.com/google/go-containerregistry/cmd/gcrane
```

If you fetch a manifest by repo digest the hash matches:

``` console
$ gcrane manifest us.gcr.io/build-artifacts-xxx/environments/production/foo.image@sha256:89df6d6d6f07de185d89a3ae1fd37cb99b593f33cb7cce8c0f047930a3e47ba7 | sha256sum
89df6d6d6f07de185d89a3ae1fd37cb99b593f33cb7cce8c0f047930a3e47ba7
```
"""

import argparse
import hashlib
import json
from pathlib import Path


def _main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, required=True)
    parser.add_argument('--config-digest', type=str, required=True)
    parser.add_argument('--blobsum', type=str, action='append', required=True)
    parser.add_argument('--zipped-layer', type=str, action='append', required=True)
    parser.add_argument('--manifest', type=str, required=True)  # destination
    parser.add_argument('--digest', type=str, required=True)  # destination
    args = parser.parse_args()

    doc = {
        'schemaVersion': 2,
        'mediaType': 'application/vnd.docker.distribution.manifest.v2+json',
        'config': {
            'mediaType': 'application/vnd.docker.container.image.v1+json',
            'size': Path(args.config).stat().st_size,
            'digest': f'sha256:{Path(args.config_digest).read_text()}',
        },
        'layers': [
            {
                'mediaType': 'application/vnd.docker.image.rootfs.diff.tar.gzip',
                'size': Path(zipped_layer).stat().st_size,
                'digest': f'sha256:{Path(blobsum).read_text()}',
            } for zipped_layer, blobsum in zip(args.zipped_layer, args.blobsum)
        ],
    }

    # write out the manifest for debugging
    manifest = json.dumps(doc, separators=(',', ':')).encode('utf-8')
    Path(args.manifest).write_bytes(manifest)

    # write out the distribution digest for image stamping
    h = hashlib.sha256()
    h.update(manifest)
    Path(args.digest).write_text(h.hexdigest(), encoding='utf-8')


if __name__ == '__main__':
    _main()

providers.bzl

DigestInfo = provider(fields = {
    "digest_path": "for ImageInfo targets, the digest path",
})

RepoInfo = provider(fields = {
    "image_digests": "for BundleInfo targets, a dict from image name to digest path",
})