From a3c62e698d902b35cba5f5649f863849293651f5 Mon Sep 17 00:00:00 2001
From: roman-kiselenko <shindu666@gmail.com>
Date: Tue, 5 Dec 2023 12:46:57 +0300
Subject: [PATCH] add label flags to ctr import, add integration tests

Signed-off-by: roman-kiselenko <roman.kiselenko.dev@gmail.com>
---
 client/import.go                  | 16 +++++++++++++++-
 cmd/ctr/commands/images/import.go | 12 +++++++++++-
 integration/client/import_test.go | 25 +++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/client/import.go b/client/import.go
index 79031f5c3..d852f3712 100644
--- a/client/import.go
+++ b/client/import.go
@@ -38,6 +38,7 @@ type importOpts struct {
 	compress        bool
 	discardLayers   bool
 	skipMissing     bool
+	imageLabels     map[string]string
 }
 
 // ImportOpt allows the caller to specify import specific options
@@ -52,6 +53,14 @@ func WithImageRefTranslator(f func(string) string) ImportOpt {
 	}
 }
 
+// WithImageLabels are the image labels to apply to a new image
+func WithImageLabels(labels map[string]string) ImportOpt {
+	return func(c *importOpts) error {
+		c.imageLabels = labels
+		return nil
+	}
+}
+
 // WithDigestRef is used to create digest images for each
 // manifest in the index.
 func WithDigestRef(f func(digest.Digest) string) ImportOpt {
@@ -223,7 +232,12 @@ func (c *Client) Import(ctx context.Context, reader io.Reader, opts ...ImportOpt
 	}
 
 	for i := range imgs {
-		img, err := is.Update(ctx, imgs[i], "target")
+		fieldsPath := []string{"target"}
+		if iopts.imageLabels != nil {
+			fieldsPath = append(fieldsPath, "labels")
+			imgs[i].Labels = iopts.imageLabels
+		}
+		img, err := is.Update(ctx, imgs[i], fieldsPath...)
 		if err != nil {
 			if !errdefs.IsNotFound(err) {
 				return nil, err
diff --git a/cmd/ctr/commands/images/import.go b/cmd/ctr/commands/images/import.go
index 724964b54..bb663904a 100644
--- a/cmd/ctr/commands/images/import.go
+++ b/cmd/ctr/commands/images/import.go
@@ -98,7 +98,7 @@ If foobar.tar contains an OCI ref named "latest" and anonymous ref "sha256:deadb
 			Name:  "discard-unpacked-layers",
 			Usage: "Allow the garbage collector to clean layers up from the content store after unpacking, cannot be used with --no-unpack, false by default",
 		},
-	}, commands.SnapshotterFlags...),
+	}, append(commands.SnapshotterFlags, commands.LabelFlag)...),
 
 	Action: func(context *cli.Context) error {
 		var (
@@ -123,6 +123,11 @@ If foobar.tar contains an OCI ref named "latest" and anonymous ref "sha256:deadb
 				overwrite = true
 			}
 
+			labels := context.StringSlice("label")
+			if len(labels) > 0 {
+				opts = append(opts, image.WithImageLabels(commands.LabelArgs(labels)))
+			}
+
 			if context.Bool("digests") {
 				opts = append(opts, image.WithDigestRef(prefix, overwrite, !context.Bool("skip-digest-for-named")))
 			} else {
@@ -237,6 +242,11 @@ If foobar.tar contains an OCI ref named "latest" and anonymous ref "sha256:deadb
 			opts = append(opts, containerd.WithDiscardUnpackedLayers())
 		}
 
+		labels := context.StringSlice("label")
+		if len(labels) > 0 {
+			opts = append(opts, containerd.WithImageLabels(commands.LabelArgs(labels)))
+		}
+
 		ctx, done, err := client.WithLease(ctx)
 		if err != nil {
 			return err
diff --git a/integration/client/import_test.go b/integration/client/import_test.go
index e228d2039..b52cb16ab 100644
--- a/integration/client/import_test.go
+++ b/integration/client/import_test.go
@@ -173,6 +173,8 @@ func TestImport(t *testing.T) {
 
 	m1, d3, expManifest := createManifest(c1, [][]byte{b1})
 
+	importLabels := map[string]string{"foo": "bar"}
+
 	c2, _ := createConfig(runtime.GOOS, runtime.GOARCH, "test2")
 	m2, d5, _ := createManifest(c2, [][]byte{{1, 2, 3, 4, 5}})
 
@@ -315,6 +317,29 @@ func TestImport(t *testing.T) {
 				checkManifest(ctx, t, client.ContentStore(), imgs[0].Target, expManifest)
 			},
 		},
+		{
+			Name: "OCI-Labels",
+			Writer: tartest.TarAll(
+				tc.Dir(ocispec.ImageBlobsDir, 0o755),
+				tc.Dir(ocispec.ImageBlobsDir+"/sha256", 0o755),
+				tc.File(ocispec.ImageBlobsDir+"/sha256/"+d1.Encoded(), b1, 0o644),
+				tc.File(ocispec.ImageBlobsDir+"/sha256/"+d2.Encoded(), c1, 0o644),
+				tc.File(ocispec.ImageBlobsDir+"/sha256/"+d3.Encoded(), m1, 0o644),
+				tc.File(ocispec.ImageIndexFile, createIndex(m1, "latest", "docker.io/lib/img:ok"), 0o644),
+				tc.File(ocispec.ImageLayoutFile, []byte(`{"imageLayoutVersion":"`+ocispec.ImageLayoutVersion+`"}`), 0o644),
+			),
+			Check: func(ctx context.Context, t *testing.T, _ *Client, imgs []images.Image) {
+				for i := range imgs {
+					if !reflect.DeepEqual(imgs[i].Labels, importLabels) {
+						t.Fatalf("DeepEqual on labels failed img.Labels: %+v expected: %+v", imgs[i].Labels, importLabels)
+					}
+				}
+			},
+			Opts: []ImportOpt{
+				WithImageLabels(importLabels),
+				WithImageRefTranslator(archive.AddRefPrefix("localhost:5000/myimage")),
+			},
+		},
 		{
 			Name: "OCIPrefixName",
 			Writer: tartest.TarAll(