feat: Add zstd compression to crane append and mutate

ReillyBrogan · ReillyBrogan · commit 188e23b4ea3d · 2023-09-28T10:49:12.000-05:00
diff --git a/cmd/crane/cmd/append.go b/cmd/crane/cmd/append.go
@@ -17,6 +17,7 @@ package cmd
 import (
 	"fmt"
 
+	"github.com/google/go-containerregistry/pkg/compression"
 	"github.com/google/go-containerregistry/pkg/crane"
 	"github.com/google/go-containerregistry/pkg/logs"
 	"github.com/google/go-containerregistry/pkg/name"
@@ -33,6 +34,7 @@ func NewCmdAppend(options *[]crane.Option) *cobra.Command {
 	var baseRef, newTag, outFile string
 	var newLayers []string
 	var annotate, ociEmptyBase bool
+	var comp = compression.GZip
 
 	appendCmd := &cobra.Command{
 		Use:   "append",
@@ -63,7 +65,7 @@ container image.`,
 				}
 			}
 
-			img, err := crane.Append(base, newLayers...)
+			img, err := crane.AppendWithCompression(base, comp, newLayers...)
 			if err != nil {
 				return fmt.Errorf("appending %v: %w", newLayers, err)
 			}
@@ -111,6 +113,7 @@ container image.`,
 	appendCmd.Flags().StringVarP(&baseRef, "base", "b", "", "Name of base image to append to")
 	appendCmd.Flags().StringVarP(&newTag, "new_tag", "t", "", "Tag to apply to resulting image")
 	appendCmd.Flags().StringSliceVarP(&newLayers, "new_layer", "f", []string{}, "Path to tarball to append to image")
+	appendCmd.Flags().VarP(&comp, "compression", "c", "Compression to use for new layers")
 	appendCmd.Flags().StringVarP(&outFile, "output", "o", "", "Path to new tarball of resulting image")
 	appendCmd.Flags().BoolVar(&annotate, "set-base-image-annotations", false, "If true, annotate the resulting image as being based on the base image")
 	appendCmd.Flags().BoolVar(&ociEmptyBase, "oci-empty-base", false, "If true, empty base image will have OCI media types instead of Docker")
diff --git a/cmd/crane/cmd/mutate.go b/cmd/crane/cmd/mutate.go
@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/google/go-containerregistry/pkg/compression"
 	"github.com/google/go-containerregistry/pkg/crane"
 	"github.com/google/go-containerregistry/pkg/name"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
@@ -30,6 +31,7 @@ import (
 func NewCmdMutate(options *[]crane.Option) *cobra.Command {
 	var labels map[string]string
 	var annotations map[string]string
+	var comp = compression.GZip
 	var envVars keyToValue
 	var entrypoint, cmd []string
 	var newLayers []string
@@ -67,7 +69,7 @@ func NewCmdMutate(options *[]crane.Option) *cobra.Command {
 				return fmt.Errorf("pulling %s: %w", ref, err)
 			}
 			if len(newLayers) != 0 {
-				img, err = crane.Append(img, newLayers...)
+				img, err = crane.AppendWithCompression(img, comp, newLayers...)
 				if err != nil {
 					return fmt.Errorf("appending %v: %w", newLayers, err)
 				}
@@ -174,6 +176,7 @@ func NewCmdMutate(options *[]crane.Option) *cobra.Command {
 	mutateCmd.Flags().StringToStringVarP(&annotations, "annotation", "a", nil, "New annotations to add")
 	mutateCmd.Flags().StringToStringVarP(&labels, "label", "l", nil, "New labels to add")
 	mutateCmd.Flags().VarP(&envVars, "env", "e", "New envvar to add")
+	mutateCmd.Flags().VarP(&comp, "compression", "c", "Compression to use for new layers")
 	mutateCmd.Flags().StringSliceVar(&entrypoint, "entrypoint", nil, "New entrypoint to set")
 	mutateCmd.Flags().StringSliceVar(&cmd, "cmd", nil, "New cmd to set")
 	mutateCmd.Flags().StringVar(&newRepo, "repo", "", "Repository to push the mutated image to. If provided, push by digest to this repository.")
diff --git a/cmd/crane/doc/crane_append.md b/cmd/crane/doc/crane_append.md
diff --git a/cmd/crane/doc/crane_mutate.md b/cmd/crane/doc/crane_mutate.md
diff --git a/pkg/compression/compression.go b/pkg/compression/compression.go
@@ -15,6 +15,8 @@
 // Package compression abstracts over gzip and zstd.
 package compression
 
+import "errors"
+
 // Compression is an enumeration of the supported compression algorithms
 type Compression string
 
@@ -24,3 +26,25 @@ const (
 	GZip Compression = "gzip"
 	ZStd Compression = "zstd"
 )
+
+// Used by fmt.Print and Cobra in help text
+func (e *Compression) String() string {
+	return string(*e)
+}
+
+func (e *Compression) Set(v string) error {
+	switch v {
+	case "none", "gzip", "zstd":
+		*e = Compression(v)
+		return nil
+	default:
+		return errors.New(`must be one of "none", "gzip, or "zstd"`)
+	}
+}
+
+// Used in Cobra help text
+func (e *Compression) Type() string {
+	return "Compression"
+}
+
+var ErrZStdNonOci = errors.New("ZSTD compression can only be used with an OCI base image")
diff --git a/pkg/crane/append.go b/pkg/crane/append.go
@@ -19,6 +19,7 @@ import (
 	"os"
 
 	"github.com/google/go-containerregistry/internal/windows"
+	"github.com/google/go-containerregistry/pkg/compression"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/mutate"
 	"github.com/google/go-containerregistry/pkg/v1/stream"
@@ -40,6 +41,16 @@ func isWindows(img v1.Image) (bool, error) {
 // "windows"), the contents of the tarballs will be modified to be suitable for
 // a Windows container image.`,
 func Append(base v1.Image, paths ...string) (v1.Image, error) {
+	return AppendWithCompression(base, compression.GZip, paths...)
+}
+
+// Append reads a layer from path and appends it the the v1.Image base
+// using the specified compression
+//
+// If the base image is a Windows base image (i.e., its config.OS is
+// "windows"), the contents of the tarballs will be modified to be suitable for
+// a Windows container image.`,
+func AppendWithCompression(base v1.Image, comp compression.Compression, paths ...string) (v1.Image, error) {
 	if base == nil {
 		return nil, fmt.Errorf("invalid argument: base")
 	}
@@ -61,6 +72,12 @@ func Append(base v1.Image, paths ...string) (v1.Image, error) {
 		layerType = types.OCILayer
 	}
 
+	if comp == compression.ZStd && layerType == types.OCILayer {
+		layerType = types.OCILayerZStd
+	} else if comp == compression.ZStd {
+		return nil, compression.ErrZStdNonOci
+	}
+
 	layers := make([]v1.Layer, 0, len(paths))
 	for _, path := range paths {
 		layer, err := getLayer(path, layerType)
@@ -90,7 +107,11 @@ func getLayer(path string, layerType types.MediaType) (v1.Layer, error) {
 		return stream.NewLayer(f, stream.WithMediaType(layerType)), nil
 	}
 
-	return tarball.LayerFromFile(path, tarball.WithMediaType(layerType))
+	if layerType == types.OCILayerZStd {
+		return tarball.LayerFromFile(path, tarball.WithMediaType(layerType), tarball.WithCompression(compression.ZStd))
+	} else {
+		return tarball.LayerFromFile(path, tarball.WithMediaType(layerType))
+	}
 }
 
 // If we're dealing with a named pipe, trying to open it multiple times will
diff --git a/pkg/crane/append_test.go b/pkg/crane/append_test.go
@@ -15,8 +15,10 @@
 package crane_test
 
 import (
+	"errors"
 	"testing"
 
+	"github.com/google/go-containerregistry/pkg/compression"
 	"github.com/google/go-containerregistry/pkg/crane"
 	"github.com/google/go-containerregistry/pkg/v1/empty"
 	"github.com/google/go-containerregistry/pkg/v1/mutate"
@@ -25,7 +27,7 @@ import (
 
 func TestAppendWithOCIBaseImage(t *testing.T) {
 	base := mutate.MediaType(empty.Image, types.OCIManifestSchema1)
-	img, err := crane.Append(base, "testdata/content.tar")
+	img, err := crane.AppendWithCompression(base, compression.GZip, "testdata/content.tar")
 
 	if err != nil {
 		t.Fatalf("crane.Append(): %v", err)
@@ -48,8 +50,33 @@ func TestAppendWithOCIBaseImage(t *testing.T) {
 	}
 }
 
+func TestAppendWithOCIBaseImageZstd(t *testing.T) {
+	base := mutate.MediaType(empty.Image, types.OCIManifestSchema1)
+	img, err := crane.AppendWithCompression(base, compression.ZStd, "testdata/content.tar")
+
+	if err != nil {
+		t.Fatalf("crane.Append(): %v", err)
+	}
+
+	layers, err := img.Layers()
+
+	if err != nil {
+		t.Fatalf("img.Layers(): %v", err)
+	}
+
+	mediaType, err := layers[0].MediaType()
+
+	if err != nil {
+		t.Fatalf("layers[0].MediaType(): %v", err)
+	}
+
+	if got, want := mediaType, types.OCILayerZStd; got != want {
+		t.Errorf("MediaType(): want %q, got %q", want, got)
+	}
+}
+
 func TestAppendWithDockerBaseImage(t *testing.T) {
-	img, err := crane.Append(empty.Image, "testdata/content.tar")
+	img, err := crane.AppendWithCompression(empty.Image, compression.GZip, "testdata/content.tar")
 
 	if err != nil {
 		t.Fatalf("crane.Append(): %v", err)
@@ -71,3 +98,15 @@ func TestAppendWithDockerBaseImage(t *testing.T) {
 		t.Errorf("MediaType(): want %q, got %q", want, got)
 	}
 }
+
+func TestAppendWithDockerBaseImageZstd(t *testing.T) {
+	img, err := crane.AppendWithCompression(empty.Image, compression.ZStd, "testdata/content.tar")
+
+	if img != nil {
+		t.Fatalf("Unexpected success")
+	}
+
+	if !errors.Is(err, compression.ErrZStdNonOci) {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+}
diff --git a/pkg/crane/crane_test.go b/pkg/crane/crane_test.go
@@ -30,6 +30,7 @@ import (
 
 	"github.com/google/go-containerregistry/internal/compare"
 	"github.com/google/go-containerregistry/pkg/authn"
+	"github.com/google/go-containerregistry/pkg/compression"
 	"github.com/google/go-containerregistry/pkg/crane"
 	"github.com/google/go-containerregistry/pkg/name"
 	"github.com/google/go-containerregistry/pkg/registry"
@@ -439,7 +440,7 @@ func TestCraneFilesystem(t *testing.T) {
 	tw.Flush()
 	tw.Close()
 
-	img, err = crane.Append(img, tmp.Name())
+	img, err = crane.AppendWithCompression(img, compression.GZip, tmp.Name())
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -500,7 +501,7 @@ func TestStreamingAppend(t *testing.T) {
 
 	os.Stdin = tmp
 
-	img, err := crane.Append(empty.Image, "-")
+	img, err := crane.AppendWithCompression(empty.Image, compression.GZip, "-")
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/pkg/v1/tarball/write.go b/pkg/v1/tarball/write.go
@@ -28,6 +28,7 @@ import (
 	"github.com/google/go-containerregistry/pkg/name"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/partial"
+	"github.com/google/go-containerregistry/pkg/v1/types"
 )
 
 // WriteToFile writes in the compressed format to a tarball, on disk.
@@ -182,9 +183,18 @@ func writeImagesToTar(imageToTags map[v1.Image][]string, m []byte, size int64, w
 			// Drop the algorithm prefix, e.g. "sha256:"
 			hex := d.Hex
 
-			// gunzip expects certain file extensions:
-			// https://www.gnu.org/software/gzip/manual/html_node/Overview.html
-			layerFiles[i] = fmt.Sprintf("%s.tar.gz", hex)
+			mt, err := l.MediaType()
+			if err != nil {
+				return sendProgressWriterReturn(pw, err)
+			}
+			// OCILayerZStd is the only layer type that currently supports ZSTD-compression
+			if mt == types.OCILayerZStd {
+				layerFiles[i] = fmt.Sprintf("%s.tar.zst", hex)
+			} else {
+				// gunzip expects certain file extensions:
+				// https://www.gnu.org/software/gzip/manual/html_node/Overview.html
+				layerFiles[i] = fmt.Sprintf("%s.tar.gz", hex)
+			}
 
 			if _, ok := seenLayerDigests[hex]; ok {
 				continue
