Refactor image importer

Allow customization of reference creation.
Add option for digest references.

Signed-off-by: Derek McGowan <derek@mcgstyle.net>

import.go

@@ -18,35 +18,75 @@ package containerd
 
 import (
 	"context"
+	"encoding/json"
 	"io"
 
+	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/images"
+	"github.com/containerd/containerd/images/oci"
+	digest "github.com/opencontainers/go-digest"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
 type importOpts struct {
+	indexName string
+	imageRefT func(string) string
+	dgstRefT  func(digest.Digest) string
+	importer  images.Importer
 }
 
 // ImportOpt allows the caller to specify import specific options
-type ImportOpt func(c *importOpts) error
+type ImportOpt func(*importOpts) error
 
-func resolveImportOpt(opts ...ImportOpt) (importOpts, error) {
-	var iopts importOpts
-	for _, o := range opts {
-		if err := o(&iopts); err != nil {
-			return iopts, err
-		}
+// WithImageRefTranslator is used to translate the index reference
+// to an image reference for the image store.
+func WithImageRefTranslator(f func(string) string) ImportOpt {
+	return func(c *importOpts) error {
+		c.imageRefT = f
+		return nil
+	}
+}
+
+// WithDigestRef is used to create digest images for each
+// manifest in the index.
+func WithDigestRef(f func(digest.Digest) string) ImportOpt {
+	return func(c *importOpts) error {
+		c.dgstRefT = f
+		return nil
+	}
+}
+
+// WithIndexName creates a tag pointing to the imported index
+func WithIndexName(name string) ImportOpt {
+	return func(c *importOpts) error {
+		c.indexName = name
+		return nil
+	}
+}
+
+// WithImporter sets the importer to use for converting
+// the read stream into an OCI Index.
+func WithImporter(importer images.Importer) ImportOpt {
+	return func(c *importOpts) error {
+		c.importer = importer
+		return nil
 	}
-	return iopts, nil
 }
 
 // Import imports an image from a Tar stream using reader.
 // Caller needs to specify importer. Future version may use oci.v1 as the default.
 // Note that unreferrenced blobs may be imported to the content store as well.
-func (c *Client) Import(ctx context.Context, importer images.Importer, reader io.Reader, opts ...ImportOpt) ([]Image, error) {
-	_, err := resolveImportOpt(opts...) // unused now
-	if err != nil {
-		return nil, err
+func (c *Client) Import(ctx context.Context, reader io.Reader, opts ...ImportOpt) ([]images.Image, error) {
+	var iopts importOpts
+	for _, o := range opts {
+		if err := o(&iopts); err != nil {
+			return nil, err
+		}
+	}
+
+	if iopts.importer == nil {
+		iopts.importer = &oci.V1Importer{}
 	}
 
 	ctx, done, err := c.WithLease(ctx)
@@ -55,31 +95,86 @@ func (c *Client) Import(ctx context.Context, importer images.Importer, reader io
 	}
 	defer done(ctx)
 
-	imgrecs, err := importer.Import(ctx, c.ContentStore(), reader)
+	index, err := iopts.importer.Import(ctx, c.ContentStore(), reader)
 	if err != nil {
-		// is.Update() is not called on error
 		return nil, err
 	}
 
-	is := c.ImageService()
-	var images []Image
-	for _, imgrec := range imgrecs {
-		if updated, err := is.Update(ctx, imgrec, "target"); err != nil {
+	var (
+		imgs []images.Image
+		cs   = c.ContentStore()
+		is   = c.ImageService()
+	)
+
+	if iopts.indexName != "" {
+		imgs = append(imgs, images.Image{
+			Name:   iopts.indexName,
+			Target: index,
+		})
+	}
+
+	var handler images.HandlerFunc
+	handler = func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		// Only save images at top level
+		if desc.Digest != index.Digest {
+			return images.Children(ctx, cs, desc)
+		}
+
+		p, err := content.ReadBlob(ctx, cs, desc)
+		if err != nil {
+			return nil, err
+		}
+
+		var idx ocispec.Index
+		if err := json.Unmarshal(p, &idx); err != nil {
+			return nil, err
+		}
+
+		for _, m := range idx.Manifests {
+			if ref := m.Annotations[ocispec.AnnotationRefName]; ref != "" {
+				if iopts.imageRefT != nil {
+					ref = iopts.imageRefT(ref)
+				}
+				if ref != "" {
+					imgs = append(imgs, images.Image{
+						Name:   ref,
+						Target: m,
+					})
+				}
+			}
+			if iopts.dgstRefT != nil {
+				ref := iopts.dgstRefT(m.Digest)
+				if ref != "" {
+					imgs = append(imgs, images.Image{
+						Name:   ref,
+						Target: m,
+					})
+				}
+			}
+		}
+
+		return idx.Manifests, nil
+	}
+
+	handler = images.SetChildrenLabels(cs, handler)
+	if err := images.Walk(ctx, handler, index); err != nil {
+		return nil, err
+	}
+
+	for i := range imgs {
+		img, err := is.Update(ctx, imgs[i], "target")
+		if err != nil {
 			if !errdefs.IsNotFound(err) {
 				return nil, err
 			}
 
-			created, err := is.Create(ctx, imgrec)
+			img, err = is.Create(ctx, imgs[i])
 			if err != nil {
 				return nil, err
 			}
-
-			imgrec = created
-		} else {
-			imgrec = updated
 		}
-
-		images = append(images, NewImage(c, imgrec))
+		imgs[i] = img
 	}
-	return images, nil
+
+	return imgs, nil
 }