From dd87a7c5a8e6e818c8adcb914b499d4c34837bd2 Mon Sep 17 00:00:00 2001
From: Austin Vazquez <macedonv@amazon.com>
Date: Tue, 15 Oct 2024 19:49:59 +0000
Subject: [PATCH] Refactor load from archive functionality

Signed-off-by: Austin Vazquez <macedonv@amazon.com>
---
 cmd/nerdctl/image/image_load.go |  3 +-
 pkg/imgutil/load/load.go        | 92 ++++++++++++++++++++-------------
 2 files changed, 57 insertions(+), 38 deletions(-)

diff --git a/cmd/nerdctl/image/image_load.go b/cmd/nerdctl/image/image_load.go
index a80918f66f5..3ff8b18a892 100644
--- a/cmd/nerdctl/image/image_load.go
+++ b/cmd/nerdctl/image/image_load.go
@@ -94,5 +94,6 @@ func loadAction(cmd *cobra.Command, _ []string) error {
 	}
 	defer cancel()
 
-	return load.Load(ctx, client, options)
+	_, err = load.FromArchive(ctx, client, options)
+	return err
 }
diff --git a/pkg/imgutil/load/load.go b/pkg/imgutil/load/load.go
index 2673e442597..c1d8c239b0c 100644
--- a/pkg/imgutil/load/load.go
+++ b/pkg/imgutil/load/load.go
@@ -34,24 +34,12 @@ import (
 	"github.com/containerd/nerdctl/v2/pkg/platformutil"
 )
 
-type readCounter struct {
-	io.Reader
-	N int
-}
-
-func (r *readCounter) Read(p []byte) (int, error) {
-	n, err := r.Reader.Read(p)
-	if n > 0 {
-		r.N += n
-	}
-	return n, err
-}
-
-func Load(ctx context.Context, client *containerd.Client, options types.ImageLoadOptions) error {
+// FromArchive will load and unpack the images from the provided tar archive specified in image load options.
+func FromArchive(ctx context.Context, client *containerd.Client, options types.ImageLoadOptions) ([]images.Image, error) {
 	if options.Input != "" {
 		f, err := os.Open(options.Input)
 		if err != nil {
-			return err
+			return []images.Image{}, err
 		}
 		defer f.Close()
 		options.Stdin = f
@@ -59,55 +47,85 @@ func Load(ctx context.Context, client *containerd.Client, options types.ImageLoa
 		// check if stdin is empty.
 		stdinStat, err := os.Stdin.Stat()
 		if err != nil {
-			return err
+			return []images.Image{}, err
 		}
 		if stdinStat.Size() == 0 && (stdinStat.Mode()&os.ModeNamedPipe) == 0 {
-			return errors.New("stdin is empty and input flag is not specified")
+			return []images.Image{}, errors.New("stdin is empty and input flag is not specified")
 		}
 	}
 	decompressor, err := compression.DecompressStream(options.Stdin)
 	if err != nil {
-		return err
+		return []images.Image{}, err
 	}
 	platMC, err := platformutil.NewMatchComparer(options.AllPlatforms, options.Platform)
 	if err != nil {
-		return err
+		return []images.Image{}, err
+	}
+	imgs, err := importImages(ctx, client, decompressor, options.GOptions.Snapshotter, platMC)
+	if err != nil {
+		return []images.Image{}, err
+	}
+	unpackedImages := make([]images.Image, 0, len(imgs))
+	for _, img := range imgs {
+		err := unpackImage(ctx, client, img, platMC, options)
+		if err != nil {
+			return unpackedImages, fmt.Errorf("error unpacking image (%s): %w", img.Name, err)
+		}
+		unpackedImages = append(unpackedImages, img)
+	}
+	return unpackedImages, nil
+}
+
+type readCounter struct {
+	io.Reader
+	N int
+}
+
+func (r *readCounter) Read(p []byte) (int, error) {
+	n, err := r.Reader.Read(p)
+	if n > 0 {
+		r.N += n
 	}
-	return loadImage(ctx, client, decompressor, platMC, options)
+	return n, err
 }
 
-func loadImage(ctx context.Context, client *containerd.Client, in io.Reader, platMC platforms.MatchComparer, options types.ImageLoadOptions) error {
+func importImages(ctx context.Context, client *containerd.Client, in io.Reader, snapshotter string, platformMC platforms.MatchComparer) ([]images.Image, error) {
 	// In addition to passing WithImagePlatform() to client.Import(), we also need to pass WithDefaultPlatform() to NewClient().
 	// Otherwise unpacking may fail.
 	r := &readCounter{Reader: in}
-	imgs, err := client.Import(ctx, r, containerd.WithDigestRef(archive.DigestTranslator(options.GOptions.Snapshotter)), containerd.WithSkipDigestRef(func(name string) bool { return name != "" }), containerd.WithImportPlatform(platMC))
+	imgs, err := client.Import(ctx, r,
+		containerd.WithDigestRef(archive.DigestTranslator(snapshotter)),
+		containerd.WithSkipDigestRef(func(name string) bool { return name != "" }),
+		containerd.WithImportPlatform(platformMC),
+	)
 	if err != nil {
 		if r.N == 0 {
 			// Avoid confusing "unrecognized image format"
-			return errors.New("no image was built")
+			return []images.Image{}, errors.New("no image was built")
 		}
 		if errors.Is(err, images.ErrEmptyWalk) {
 			err = fmt.Errorf("%w (Hint: set `--platform=PLATFORM` or `--all-platforms`)", err)
 		}
-		return err
+		return []images.Image{}, err
 	}
-	for _, img := range imgs {
-		image := containerd.NewImageWithPlatform(client, img, platMC)
+	return imgs, nil
+}
 
-		// TODO: Show unpack status
-		if !options.Quiet {
-			fmt.Fprintf(options.Stdout, "unpacking %s (%s)...\n", img.Name, img.Target.Digest)
-		}
+func unpackImage(ctx context.Context, client *containerd.Client, model images.Image, platform platforms.MatchComparer, options types.ImageLoadOptions) error {
+	image := containerd.NewImageWithPlatform(client, model, platform)
 
-		err = image.Unpack(ctx, options.GOptions.Snapshotter)
-		if err != nil {
-			return err
-		}
+	if !options.Quiet {
+		fmt.Fprintf(options.Stdout, "unpacking %s (%s)...\n", model.Name, model.Target.Digest)
+	}
 
-		// Loaded message is shown even when quiet.
-		repo, tag := imgutil.ParseRepoTag(img.Name)
-		fmt.Fprintf(options.Stdout, "Loaded image: %s:%s\n", repo, tag)
+	err := image.Unpack(ctx, options.GOptions.Snapshotter)
+	if err != nil {
+		return err
 	}
 
+	// Loaded message is shown even when quiet.
+	repo, tag := imgutil.ParseRepoTag(model.Name)
+	fmt.Fprintf(options.Stdout, "Loaded image: %s:%s\n", repo, tag)
+
 	return nil
 }