From 41f8d92709b7d7c96a210f2652f2862314f5ba7f Mon Sep 17 00:00:00 2001
From: Jason Hall <jasonhall@redhat.com>
Date: Fri, 11 Feb 2022 12:30:31 -0500
Subject: [PATCH] Add tarball.WithMediaType to specify layer media type (#1286)

---
 pkg/v1/tarball/layer.go      | 13 +++++++++++--
 pkg/v1/tarball/layer_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/pkg/v1/tarball/layer.go b/pkg/v1/tarball/layer.go
index 5ec1d5515..ac9e14c76 100644
--- a/pkg/v1/tarball/layer.go
+++ b/pkg/v1/tarball/layer.go
@@ -39,6 +39,7 @@ type layer struct {
 	compression        int
 	annotations        map[string]string
 	estgzopts          []estargz.Option
+	mediaType          types.MediaType
 }
 
 // Descriptor implements partial.withDescriptor.
@@ -51,7 +52,7 @@ func (l *layer) Descriptor() (*v1.Descriptor, error) {
 		Size:        l.size,
 		Digest:      digest,
 		Annotations: l.annotations,
-		MediaType:   types.DockerLayer,
+		MediaType:   l.mediaType,
 	}, nil
 }
 
@@ -82,7 +83,7 @@ func (l *layer) Size() (int64, error) {
 
 // MediaType implements v1.Layer
 func (l *layer) MediaType() (types.MediaType, error) {
-	return types.DockerLayer, nil
+	return l.mediaType, nil
 }
 
 // LayerOption applies options to layer
@@ -96,6 +97,13 @@ func WithCompressionLevel(level int) LayerOption {
 	}
 }
 
+// WithMediaType is a functional option for overriding the layer's media type.
+func WithMediaType(mt types.MediaType) LayerOption {
+	return func(l *layer) {
+		l.mediaType = mt
+	}
+}
+
 // WithCompressedCaching is a functional option that overrides the
 // logic for accessing the compressed bytes to memoize the result
 // and avoid expensive repeated gzips.
@@ -204,6 +212,7 @@ func LayerFromOpener(opener Opener, opts ...LayerOption) (v1.Layer, error) {
 	layer := &layer{
 		compression: gzip.BestSpeed,
 		annotations: make(map[string]string, 1),
+		mediaType:   types.DockerLayer,
 	}
 
 	if estgz := os.Getenv("GGCR_EXPERIMENT_ESTARGZ"); estgz == "1" {
diff --git a/pkg/v1/tarball/layer_test.go b/pkg/v1/tarball/layer_test.go
index e46d61811..839eccf8c 100644
--- a/pkg/v1/tarball/layer_test.go
+++ b/pkg/v1/tarball/layer_test.go
@@ -24,6 +24,7 @@ import (
 
 	"github.com/containerd/stargz-snapshotter/estargz"
 	"github.com/google/go-containerregistry/internal/compare"
+	"github.com/google/go-containerregistry/pkg/v1/types"
 	"github.com/google/go-containerregistry/pkg/v1/validate"
 )
 
@@ -225,6 +226,35 @@ func TestLayerFromOpenerReader(t *testing.T) {
 	}
 }
 
+func TestWithMediaType(t *testing.T) {
+	setupFixtures(t)
+	defer teardownFixtures(t)
+
+	l, err := LayerFromFile("testdata/content.tar")
+	if err != nil {
+		t.Fatalf("Unable to create layer from tar file: %v", err)
+	}
+	got, err := l.MediaType()
+	if err != nil {
+		t.Fatalf("MediaType: %v", err)
+	}
+	if want := types.DockerLayer; got != want {
+		t.Errorf("got %v, want %v", got, want)
+	}
+
+	l, err = LayerFromFile("testdata/content.tar", WithMediaType(types.OCILayer))
+	if err != nil {
+		t.Fatalf("Unable to create layer from tar file: %v", err)
+	}
+	got, err = l.MediaType()
+	if err != nil {
+		t.Fatalf("MediaType: %v", err)
+	}
+	if want := types.OCILayer; got != want {
+		t.Errorf("got %v, want %v", got, want)
+	}
+}
+
 func TestLayerFromReader(t *testing.T) {
 	setupFixtures(t)
 	defer teardownFixtures(t)