From f7df8106196518df5c6c35432843421e33990329 Mon Sep 17 00:00:00 2001
From: Jason Hall <jasonhall@redhat.com>
Date: Wed, 3 Mar 2021 13:03:31 -0500
Subject: [PATCH] Add --image-label to add labels to built images (#324)

---
 pkg/build/gobuild.go          | 10 ++++++++++
 pkg/build/gobuild_test.go     | 18 ++++++++++++++++++
 pkg/build/options.go          | 11 +++++++++++
 pkg/commands/options/build.go |  3 +++
 pkg/commands/resolver.go      |  7 +++++++
 5 files changed, 49 insertions(+)

diff --git a/pkg/build/gobuild.go b/pkg/build/gobuild.go
index fea6ff0788..4f0565e896 100644
--- a/pkg/build/gobuild.go
+++ b/pkg/build/gobuild.go
@@ -83,6 +83,7 @@ type gobuild struct {
 	mod                  *modules
 	buildContext         buildContext
 	platformMatcher      *platformMatcher
+	labels               map[string]string
 }
 
 // Option is a functional option for NewGo.
@@ -96,6 +97,7 @@ type gobuildOpener struct {
 	mod                  *modules
 	buildContext         buildContext
 	platform             string
+	labels               map[string]string
 }
 
 func (gbo *gobuildOpener) Open() (Interface, error) {
@@ -113,6 +115,7 @@ func (gbo *gobuildOpener) Open() (Interface, error) {
 		disableOptimizations: gbo.disableOptimizations,
 		mod:                  gbo.mod,
 		buildContext:         gbo.buildContext,
+		labels:               gbo.labels,
 		platformMatcher:      matcher,
 	}, nil
 }
@@ -606,6 +609,13 @@ func (g *gobuild) buildOne(ctx context.Context, s string, base v1.Image, platfor
 	cfg.Config.Env = append(cfg.Config.Env, "KO_DATA_PATH="+kodataRoot)
 	cfg.Author = "github.com/google/ko"
 
+	if cfg.Config.Labels == nil {
+		cfg.Config.Labels = map[string]string{}
+	}
+	for k, v := range g.labels {
+		cfg.Config.Labels[k] = v
+	}
+
 	image, err := mutate.ConfigFile(withApp, cfg)
 	if err != nil {
 		return nil, err
diff --git a/pkg/build/gobuild_test.go b/pkg/build/gobuild_test.go
index 60b6a8e74b..85dd4ba76c 100644
--- a/pkg/build/gobuild_test.go
+++ b/pkg/build/gobuild_test.go
@@ -27,6 +27,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/empty"
 	"github.com/google/go-containerregistry/pkg/v1/mutate"
@@ -404,6 +405,8 @@ func TestGoBuild(t *testing.T) {
 		WithCreationTime(creationTime),
 		WithBaseImages(func(context.Context, string) (Result, error) { return base, nil }),
 		withBuilder(writeTempFile),
+		WithLabel("foo", "bar"),
+		WithLabel("hello", "world"),
 	)
 	if err != nil {
 		t.Fatalf("NewGo() = %v", err)
@@ -442,6 +445,21 @@ func TestGoBuild(t *testing.T) {
 		}
 	})
 
+	t.Run("check labels", func(t *testing.T) {
+		cfg, err := img.ConfigFile()
+		if err != nil {
+			t.Fatalf("ConfigFile() = %v", err)
+		}
+
+		want := map[string]string{
+			"foo":   "bar",
+			"hello": "world",
+		}
+		got := cfg.Config.Labels
+		if d := cmp.Diff(got, want); d != "" {
+			t.Fatalf("Labels diff (-got,+want): %s", d)
+		}
+	})
 }
 
 func TestGoBuildIndex(t *testing.T) {
diff --git a/pkg/build/options.go b/pkg/build/options.go
index 413bd9929a..a8a765811b 100644
--- a/pkg/build/options.go
+++ b/pkg/build/options.go
@@ -58,6 +58,17 @@ func WithPlatforms(platforms string) Option {
 	}
 }
 
+// WithLabel is a functional option for adding labels to built images.
+func WithLabel(k, v string) Option {
+	return func(gbo *gobuildOpener) error {
+		if gbo.labels == nil {
+			gbo.labels = map[string]string{}
+		}
+		gbo.labels[k] = v
+		return nil
+	}
+}
+
 // withBuilder is a functional option for overriding the way go binaries
 // are built.
 func withBuilder(b builder) Option {
diff --git a/pkg/commands/options/build.go b/pkg/commands/options/build.go
index 67a282c866..afd6201800 100644
--- a/pkg/commands/options/build.go
+++ b/pkg/commands/options/build.go
@@ -25,6 +25,7 @@ type BuildOptions struct {
 	ConcurrentBuilds     int
 	DisableOptimizations bool
 	Platform             string
+	Labels               []string
 }
 
 func AddBuildOptions(cmd *cobra.Command, bo *BuildOptions) {
@@ -34,4 +35,6 @@ func AddBuildOptions(cmd *cobra.Command, bo *BuildOptions) {
 		"Disable optimizations when building Go code. Useful when you want to interactively debug the created container.")
 	cmd.Flags().StringVar(&bo.Platform, "platform", "",
 		"Which platform to use when pulling a multi-platform base. Format: all | <os>[/<arch>[/<variant>]][,platform]*")
+	cmd.Flags().StringSliceVar(&bo.Labels, "image-label", []string{},
+		"Which labels (key=value) to add to the image.")
 }
diff --git a/pkg/commands/resolver.go b/pkg/commands/resolver.go
index 765b9d7880..7e9acc1f8e 100644
--- a/pkg/commands/resolver.go
+++ b/pkg/commands/resolver.go
@@ -86,6 +86,13 @@ func gobuildOptions(bo *options.BuildOptions) ([]build.Option, error) {
 	if bo.DisableOptimizations {
 		opts = append(opts, build.WithDisabledOptimizations())
 	}
+	for _, lf := range bo.Labels {
+		parts := strings.SplitN(lf, "=", 2)
+		if len(parts) != 2 {
+			return nil, fmt.Errorf("invalid label flag: %s", lf)
+		}
+		opts = append(opts, build.WithLabel(parts[0], parts[1]))
+	}
 	return opts, nil
 }