Implement image publishing in resolver (#434)

* Implement image publishing in resolver

Also implement ability to switch between Go & python resolver

* buildifier

k8s/go/cmd/resolver/BUILD

@@ -24,6 +24,9 @@ go_library(
     importpath = "github.com/bazelbuild/rules_k8s/k8s/go/cmd/resolver",
     visibility = ["//visibility:private"],
     deps = [
+        "@com_github_google_go_containerregistry//pkg/authn:go_default_library",
+        "@com_github_google_go_containerregistry//pkg/name:go_default_library",
+        "@com_github_google_go_containerregistry//pkg/v1/remote:go_default_library",
         "@io_bazel_rules_docker//container/go/pkg/compat:go_default_library",
         "@io_bazel_rules_docker//container/go/pkg/utils:go_default_library",
     ],

k8s/go/cmd/resolver/resolver.go

@@ -15,38 +15,196 @@ package main
 
 import (
 	"flag"
+	"fmt"
 	"log"
+	"path"
+	"strings"
 
 	"github.com/bazelbuild/rules_docker/container/go/pkg/compat"
 	"github.com/bazelbuild/rules_docker/container/go/pkg/utils"
+	"github.com/google/go-containerregistry/pkg/authn"
+	"github.com/google/go-containerregistry/pkg/name"
+	"github.com/google/go-containerregistry/pkg/v1/remote"
 )
 
 var (
-	imgTarball   = flag.String("tarball", "", "Path to the image tarball as generated by docker save. Required if --config was not specified.")
-	imgConfig    = flag.String("config", "", "Path to the image config JSON file. Required if --tarball was not specified.")
-	baseManifest = flag.String("manifest", "", "Path to the manifest of the base image. This should be the very first image in the chain of images and is only really required for Windows images with a base image that has foreign layers.")
-	format       = flag.String("format", "", "The format of the uploaded image (Docker or OCI).")
-	layers       utils.ArrayStringFlags
+	imgChroot         = flag.String("image_chroot", "", "The repository under which to chroot image references when publishing them.")
+	k8sTemplate       = flag.String("template", "", "The k8s YAML template file to resolve.")
+	allowUnusedImages = flag.Bool("allow_unused_images", false, "Allow images that don't appear in the JSON. This is useful when generating multiple SKUs of a k8s_object, only some of which use a particular image.")
+	stampInfoFile     utils.ArrayStringFlags
+	imgSpecs          utils.ArrayStringFlags
 )
 
-func main() {
-	flag.Var(&layers, "layer", "One or more layers with the following comma separated values (Compressed layer tarball, Uncompressed layer tarball, digest file, diff ID file). e.g., --layer layer.tar.gz,layer.tar,<file with digest>,<file with diffID>.")
-	flag.Parse()
+// imageSpec describes the differents parts of an image generated by
+// rules_docker.
+type imageSpec struct {
+	// name is the name of the image.
+	name string
+	// imgTarball is the image in the `docker save` tarball format.
+	imgTarball string
+	// imgConfig if the config JSON file of the image.
+	imgConfig string
+	// digests is a list of files with the sha256 digests of the compressed
+	// layers.
+	digests []string
+	// diffIDs is a list of files with the sha256 digests of the uncompressed
+	// layers.
+	diffIDs []string
+	// compressedLayers are the paths to the compressed layer tarballs.
+	compressedLayers []string
+	// uncompressedLayers are the paths to the uncompressed layer tarballs.
+	uncomressedLayers []string
+}
 
-	if *imgConfig == "" {
-		log.Fatalln("Option --config is required.")
+// layers returns a list of strings that can be passed to the image reader in
+// the compatiblity package of rules_docker to read the layers of an image in
+// the format "va11,val2,val3,val4" where:
+// val1 is the compressed layer tarball.
+// val2 is the uncompressed layer tarball.
+// val3 is the digest file.
+// val4 is the diffID file.
+func (s *imageSpec) layers() ([]string, error) {
+	result := []string{}
+	if len(s.digests) != len(s.diffIDs) || len(s.diffIDs) != len(s.compressedLayers) || len(s.compressedLayers) != len(s.uncomressedLayers) {
+		return nil, fmt.Errorf("digest, diffID, compressed blobs & uncompressed blobs had unequal lengths for image %s, got %d, %d, %d, %d, want all of the lengths to be equal", s.name, len(s.digests), len(s.diffIDs), len(s.compressedLayers), len(s.uncomressedLayers))
 	}
-	imgParts, err := compat.ImagePartsFromArgs(*imgConfig, *baseManifest, *imgTarball, layers)
+	for i, digest := range s.digests {
+		diffID := s.diffIDs[i]
+		compressedLayer := s.compressedLayers[i]
+		uncompressedLayer := s.uncomressedLayers[i]
+		result = append(result, fmt.Sprintf("%s,%s,%s,%s", compressedLayer, uncompressedLayer, digest, diffID))
+	}
+	return result, nil
+}
+
+// parseImageSpec parses the differents parts of a single docker image specified
+// as string in the format "key1=val1;key2=val2" where the expected keys are:
+// 1. "name": Name of the image.
+// 2. "tarball": docker save tarball of the image.
+// 3. "config": JSON config file of the image.
+// 4. "diff_id": Files with sha256 digest of uncompressed layers.
+// 5. "digest": Files with sha256 digest of compressed layers.
+// 6. "compressed_layer": Path to compressed layer tarballs.
+// 7. "uncompressed_layer": Path to uncompressed layer tarballs.
+func parseImageSpec(spec string) (imageSpec, error) {
+	result := imageSpec{}
+	splitSpec := strings.Split(spec, ";")
+	for _, s := range splitSpec {
+		splitFields := strings.SplitN(s, "=", 2)
+		if len(splitFields) != 2 {
+			return imageSpec{}, fmt.Errorf("image spec item %q split by '=' into unexpected fields, got %d, want 2", s, len(splitFields))
+		}
+		switch splitFields[0] {
+		case "name":
+			result.name = splitFields[1]
+		case "tarball":
+			result.imgTarball = splitFields[1]
+		case "config":
+			result.imgConfig = splitFields[1]
+		case "diff_id":
+			result.diffIDs = strings.Split(splitFields[1], ",")
+		case "digest":
+			result.digests = strings.Split(splitFields[1], ",")
+		case "compressed_layer":
+			result.compressedLayers = strings.Split(splitFields[1], ",")
+		case "uncompressed_layer":
+			result.uncomressedLayers = strings.Split(splitFields[1], ",")
+		default:
+			return imageSpec{}, fmt.Errorf("unknown image spec field %q", splitFields[0])
+		}
+	}
+	return result, nil
+}
+
+// publishSingle publishes a docker image with the given spec to the remote
+// registry indicated in the image name. The image name is stamped with the
+// given stamper.
+// The stamped image name is returned referenced by its sha256 digest.
+func publishSingle(spec imageSpec, stamper *compat.Stamper) (string, error) {
+	layers, err := spec.layers()
 	if err != nil {
-		log.Fatalf("Unable to determine parts of the image from the specified arguments: %v", err)
+		return "", fmt.Errorf("unable to convert the layer parts in image spec for %s into a single comma separated argument: %v", spec.name, err)
+	}
+
+	imgParts, err := compat.ImagePartsFromArgs(spec.imgConfig, "", spec.imgTarball, layers)
+	if err != nil {
+		return "", fmt.Errorf("unable to determine parts of the image from the specified arguments: %v", err)
 	}
 	img, err := compat.ReadImage(imgParts)
 	if err != nil {
-		log.Fatalf("Error reading image: %v", err)
+		return "", fmt.Errorf("error reading image: %v", err)
+	}
+	stampedName := stamper.Stamp(spec.name)
+
+	var ref name.Reference
+	if *imgChroot != "" {
+		n := path.Join(*imgChroot, stampedName)
+		t, err := name.NewTag(n, name.WeakValidation)
+		if err != nil {
+			return "", fmt.Errorf("unable to create a docker tag from stamped name %q: %v", n, err)
+		}
+		ref = t
+	} else {
+		t, err := name.NewTag(stampedName, name.WeakValidation)
+		if err != nil {
+			return "", fmt.Errorf("unable to create a docker tag from stamped name %q: %v", stampedName, err)
+		}
+		ref = t
+	}
+	auth, err := authn.DefaultKeychain.Resolve(ref.Context())
+	if err != nil {
+		return "", fmt.Errorf("unable to get authenticator for image %v", ref.Name())
 	}
+
+	if err := remote.Write(ref, img, remote.WithAuth(auth)); err != nil {
+		return "", fmt.Errorf("unable to push image %v: %v", ref.Name(), err)
+	}
+
 	d, err := img.Digest()
 	if err != nil {
-		log.Fatalf("Unable to get digest of image: %v", err)
+		return "", fmt.Errorf("unable to get digest of image %v", ref.Name())
+	}
+
+	return fmt.Sprintf("%s/%s@%v", ref.Context().RegistryStr(), ref.Context().RepositoryStr(), d), nil
+}
+
+// publish publishes the image with the given spec. It returns:
+// 1. A map from the unstamped & tagged image name to the stamped image name
+//    referenced by its sha256 digest.
+// 2. A set of unstamped & tagged image names that were pushed to the registry.
+func publish(spec []imageSpec, stamper *compat.Stamper) (map[string]string, map[string]bool, error) {
+	overrides := make(map[string]string)
+	unseen := make(map[string]bool)
+	for _, s := range spec {
+		digestRef, err := publishSingle(s, stamper)
+		if err != nil {
+			return nil, nil, fmt.Errorf("unable to publish image %s", s.name)
+		}
+		overrides[s.name] = digestRef
+		unseen[s.name] = true
+	}
+	return overrides, unseen, nil
+}
+
+func main() {
+	flag.Var(&imgSpecs, "image_spec", "Associative lists of the constitutent elements of a docker image.")
+	flag.Var(&stampInfoFile, "stamp-info-file", "One or more Bazel stamp info files.")
+	flag.Parse()
+
+	stamper, err := compat.NewStamper(stampInfoFile)
+	if err != nil {
+		log.Fatalf("Failed to initialize the stamper: %v", err)
+	}
+
+	specs := []imageSpec{}
+	for _, s := range imgSpecs {
+		spec, err := parseImageSpec(s)
+		if err != nil {
+			log.Fatalf("Unable to parse image spec %q: %v", s, err)
+		}
+		specs = append(specs, spec)
+	}
+	if _, _, err := publish(specs, stamper); err != nil {
+		log.Fatalf("Unable to publish images: %v", err)
 	}
-	log.Printf("Successfully loaded image with digest %v.", d)
 }

k8s/object.bzl

@@ -76,9 +76,24 @@ def _impl(ctx):
             image_spec["digest"] = ",".join([_runfiles(ctx, f) for f in blobsums])
             all_inputs += blobsums
 
-            blobs = image.get("zipped_layer", [])
-            image_spec["layer"] = ",".join([_runfiles(ctx, f) for f in blobs])
-            all_inputs += blobs
+            if not ctx.attr.use_legacy_resolver:
+                # Add additional files about the image used by the Go resolver
+                # to load the image more efficiently.
+                diff_ids = image.get("diff_id", [])
+                image_spec["diff_id"] = ",".join([_runfiles(ctx, f) for f in diff_ids])
+                all_inputs += diff_ids
+
+                blobs = image.get("zipped_layer", [])
+                image_spec["compressed_layer"] = ",".join([_runfiles(ctx, f) for f in blobs])
+                all_inputs += blobs
+
+                uncompressed_blobs = image.get("unzipped_layer", [])
+                image_spec["uncompressed_layer"] = ",".join([_runfiles(ctx, f) for f in uncompressed_blobs])
+                all_inputs += uncompressed_blobs
+            else:
+                blobs = image.get("zipped_layer", [])
+                image_spec["layer"] = ",".join([_runfiles(ctx, f) for f in blobs])
+                all_inputs += blobs
 
             image_spec["config"] = _runfiles(ctx, image["config"])
             all_inputs += [image["config"]]
@@ -121,7 +136,7 @@ def _impl(ctx):
                 for spec in image_specs
             ]),
             "%{resolver_args}": " ".join(ctx.attr.resolver_args or []),
-            "%{resolver}": _runfiles(ctx, ctx.executable.resolver),
+            "%{resolver}": _runfiles(ctx, ctx.executable.resolver if ctx.attr.use_legacy_resolver else ctx.executable.go_resolver),
             "%{stamp_args}": stamp_args,
             "%{yaml}": _runfiles(ctx, ctx.outputs.substituted),
         },
@@ -133,6 +148,7 @@ def _impl(ctx):
             runfiles = ctx.runfiles(
                 files = [
                     ctx.executable.resolver,
+                    ctx.executable.go_resolver,
                     ctx.outputs.substituted,
                 ] + all_inputs,
                 transitive_files = ctx.attr.resolver[DefaultInfo].default_runfiles.files,
@@ -156,7 +172,7 @@ def _resolve(ctx, string, output):
     )
 
 def _common_impl(ctx):
-    files = [ctx.executable.resolver]
+    files = [ctx.executable.resolver, ctx.executable.go_resolver]
 
     cluster_arg = ctx.attr.cluster
     cluster_arg = ctx.expand_make_variables("cluster", cluster_arg, {})
@@ -256,6 +272,12 @@ _common_attrs = {
     # don't expose the extra actions.
     "cluster": attr.string(),
     "context": attr.string(),
+    "go_resolver": attr.label(
+        default = Label("//k8s/go/cmd/resolver"),
+        cfg = "host",
+        executable = True,
+        allow_files = True,
+    ),
     "image_chroot": attr.string(),
     # This is only needed for describe.
     "kind": attr.string(),
@@ -271,6 +293,11 @@ _common_attrs = {
     ),
     # Extra arguments to pass to the resolver.
     "resolver_args": attr.string_list(),
+    "use_legacy_resolver": attr.bool(
+        default = True,
+        doc = "Use the legacy python resolver if True. Use the experimental" +
+              " Go resolver if false.",
+    ),
     "user": attr.string(),
     "_stamper": attr.label(
         default = Label("//k8s:stamper"),