From d6eb581f836c2387ba936d4e53dc978d22cc4440 Mon Sep 17 00:00:00 2001
From: Natalie Arellano <narellano@vmware.com>
Date: Tue, 2 Apr 2024 15:27:04 -0400
Subject: [PATCH] Fix media types when there is a base image

Signed-off-by: Natalie Arellano <narellano@vmware.com>
---
 layout/layout_test.go |  2 +-
 local/v1_facade.go    |  2 +-
 remote/new.go         |  5 +++++
 remote/remote_test.go | 31 +++++++++++++++++++++++++++++++
 4 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/layout/layout_test.go b/layout/layout_test.go
index 973cb9a4..63e4832b 100644
--- a/layout/layout_test.go
+++ b/layout/layout_test.go
@@ -284,7 +284,7 @@ func testImage(t *testing.T, when spec.G, it spec.S) {
 				h.AssertDockerMediaTypes(t, img) // after saving
 			})
 
-			when("using a sparse image", func() {
+			when("using a (sparse) base image", func() {
 				it("sets the requested media types", func() {
 					img, err := layout.NewImage(
 						imagePath,
diff --git a/local/v1_facade.go b/local/v1_facade.go
index 147a135c..b0159bd4 100644
--- a/local/v1_facade.go
+++ b/local/v1_facade.go
@@ -41,7 +41,7 @@ func newV1ImageFacadeFromInspect(dockerInspect types.ImageInspect, history []ima
 		Variant:       dockerInspect.Variant,   // FIXME: this should come from options.Platform
 	}
 	layersToSet := newEmptyLayerListFrom(configFile, dockerInspect.ID, withStore, downloadLayersOnAccess)
-	return imageFrom(layersToSet, configFile, imgutil.DockerTypes)
+	return imageFrom(layersToSet, configFile, imgutil.DockerTypes) // FIXME: this should be configurable with options.MediaTypes
 }
 
 func imageFrom(layers []v1.Layer, configFile *v1.ConfigFile, requestedTypes imgutil.MediaTypes) (v1.Image, error) {
diff --git a/remote/new.go b/remote/new.go
index f84ad0c2..af9ae20b 100644
--- a/remote/new.go
+++ b/remote/new.go
@@ -38,6 +38,10 @@ func NewImage(repoName string, keychain authn.Keychain, ops ...func(*imgutil.Ima
 	if err != nil {
 		return nil, err
 	}
+	options.MediaTypes = imgutil.GetPreferredMediaTypes(*options)
+	if options.BaseImage != nil {
+		options.BaseImage, _, err = imgutil.EnsureMediaTypesAndLayers(options.BaseImage, options.MediaTypes, imgutil.PreserveLayers)
+	}
 
 	cnbImage, err := imgutil.NewCNBImage(*options)
 	if err != nil {
@@ -159,6 +163,7 @@ func emptyImage(platform imgutil.Platform) (v1.Image, error) {
 // NewV1Image returns a new v1.Image.
 // It exists to provide library users (such as pack) an easy way to construct a v1.Image with configurable options
 // (such as platform and insecure registry).
+// FIXME: this function can be deprecated in favor of remote.NewImage as this now also implements the v1.Image interface
 func NewV1Image(baseImageRepoName string, keychain authn.Keychain, ops ...func(*imgutil.ImageOptions)) (v1.Image, error) {
 	options := &imgutil.ImageOptions{}
 	for _, op := range ops {
diff --git a/remote/remote_test.go b/remote/remote_test.go
index c0fcb21a..2e2338e2 100644
--- a/remote/remote_test.go
+++ b/remote/remote_test.go
@@ -577,6 +577,37 @@ func testImage(t *testing.T, when spec.G, it spec.S) {
 				h.AssertNil(t, img.Save())
 				h.AssertOCIMediaTypes(t, img.UnderlyingImage()) // after saving
 			})
+
+			when("using a base image", func() {
+				it("sets the requested media types", func() {
+					baseImageName := newTestImageName()
+					baseImage, err := remote.NewImage(
+						baseImageName,
+						authn.DefaultKeychain,
+						remote.WithMediaTypes(imgutil.DockerTypes),
+					)
+					h.AssertNil(t, err)
+					h.AssertNil(t, baseImage.Save())
+
+					img, err := remote.NewImage(
+						newTestImageName(),
+						authn.DefaultKeychain,
+						remote.WithMediaTypes(imgutil.OCITypes),
+						remote.FromBaseImage(baseImageName),
+					)
+					h.AssertNil(t, err)
+					h.AssertOCIMediaTypes(t, img.UnderlyingImage()) // before saving
+					// add a random layer
+					newLayerPath, err := h.CreateSingleFileLayerTar("/new-layer.txt", "new-layer", "linux")
+					h.AssertNil(t, err)
+					defer os.Remove(newLayerPath)
+					err = img.AddLayer(newLayerPath)
+					h.AssertNil(t, err)
+					h.AssertOCIMediaTypes(t, img.UnderlyingImage()) // after adding a layer
+					h.AssertNil(t, img.Save())
+					h.AssertOCIMediaTypes(t, img.UnderlyingImage()) // after saving
+				})
+			})
 		})
 
 		when("#WithConfig", func() {