✨ Normalize media type for layers during image rebasing

pkg/v1/mutate/rebase.go

@@ -16,6 +16,7 @@ package mutate
 
 import (
 	"fmt"
+	"github.com/google/go-containerregistry/pkg/v1/types"
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/empty"
@@ -96,14 +97,23 @@ func Rebase(orig, oldBase, newBase v1.Image) (v1.Image, error) {
 	if err != nil {
 		return nil, fmt.Errorf("could not get new base layers for new base: %w", err)
 	}
+
+	// Find the last DockerLayer or OCILayer type in the new base image
+	var newBaseLastLayerType = lastDockerOrOciLayerType(newBaseLayers)
+
+	// If no layer type was found in the new base, check the original image
+	if newBaseLastLayerType == "" {
+		newBaseLastLayerType = lastDockerOrOciLayerType(origLayers)
+	}
+
 	// Add new base layers.
-	rebasedImage, err = Append(rebasedImage, createAddendums(0, 0, newConfig.History, newBaseLayers)...)
+	rebasedImage, err = Append(rebasedImage, createAddendums(0, 0, newConfig.History, newBaseLayers, newBaseLastLayerType)...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to append new base image: %w", err)
 	}
 
 	// Add original layers above the old base.
-	rebasedImage, err = Append(rebasedImage, createAddendums(len(oldConfig.History), len(oldBaseLayers)+1, origConfig.History, origLayers)...)
+	rebasedImage, err = Append(rebasedImage, createAddendums(len(oldConfig.History), len(oldBaseLayers)+1, origConfig.History, origLayers, newBaseLastLayerType)...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to append original image: %w", err)
 	}
@@ -113,7 +123,7 @@ func Rebase(orig, oldBase, newBase v1.Image) (v1.Image, error) {
 
 // createAddendums makes a list of addendums from a history and layers starting from a specific history and layer
 // indexes.
-func createAddendums(startHistory, startLayer int, history []v1.History, layers []v1.Layer) []Addendum {
+func createAddendums(startHistory, startLayer int, history []v1.History, layers []v1.Layer, layerType types.MediaType) []Addendum {
 	var adds []Addendum
 	// History should be a superset of layers; empty layers (e.g. ENV statements) only exist in history.
 	// They cannot be iterated identically but must be walked independently, only advancing the iterator for layers
@@ -128,17 +138,65 @@ func createAddendums(startHistory, startLayer int, history []v1.History, layers
 		}
 		if historyIndex >= startHistory || layerIndex >= startLayer {
 			adds = append(adds, Addendum{
-				Layer:   layer,
+				Layer:   layerWithNormalizedMediaType(layer, layerType),
 				History: history[historyIndex],
 			})
 		}
 	}
 	// In the event history was malformed or non-existent, append the remaining layers.
 	for i := layerIndex; i < len(layers); i++ {
 		if i >= startLayer {
-			adds = append(adds, Addendum{Layer: layers[layerIndex]})
+			adds = append(adds, Addendum{
+				Layer: layerWithNormalizedMediaType(layers[layerIndex], layerType),
+			})
 		}
 	}
 
 	return adds
 }
+
+func lastDockerOrOciLayerType(fromLayers []v1.Layer) types.MediaType {
+	for i := len(fromLayers) - 1; i >= 0; i-- {
+		mediaType, err := fromLayers[i].MediaType()
+		if err == nil && isDockerOrOciLayerType(mediaType) {
+			return mediaType
+		}
+	}
+	return ""
+}
+
+func isDockerOrOciLayerType(mediaType types.MediaType) bool {
+	return mediaType == types.DockerLayer || mediaType == types.OCILayer
+}
+
+// mediaTypeOverridingLayer wraps a layer and overrides its MediaType method.
+type mediaTypeOverridingLayer struct {
+	v1.Layer
+	mediaType types.MediaType
+}
+
+// MediaType implements v1.Layer
+func (l *mediaTypeOverridingLayer) MediaType() (types.MediaType, error) {
+	var mt = l.mediaType
+	var err error = nil
+	if mt == "" {
+		mt, err = l.Layer.MediaType()
+	}
+	return mt, err
+}
+
+func layerWithNormalizedMediaType(layer v1.Layer, layerType types.MediaType) v1.Layer {
+	// Check the layer's media type
+	if layer != nil && layerType != "" {
+		mediaType, err := layer.MediaType()
+		// Only modify the media type if it's DockerLayer or OCILayer
+		if err == nil && isDockerOrOciLayerType(mediaType) {
+			// Wrap the layer to override its MediaType method
+			return &mediaTypeOverridingLayer{
+				Layer:     layer,
+				mediaType: layerType,
+			}
+		}
+	}
+	return layer
+}

pkg/v1/mutate/rebase_test.go

@@ -19,8 +19,10 @@ import (
 	"time"
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/empty"
 	"github.com/google/go-containerregistry/pkg/v1/mutate"
 	"github.com/google/go-containerregistry/pkg/v1/random"
+	"github.com/google/go-containerregistry/pkg/v1/types"
 )
 
 func layerDigests(t *testing.T, img v1.Image) []string {
@@ -177,3 +179,95 @@ func TestRebase(t *testing.T) {
 		t.Errorf("ConfigFile property OSVersion mismatch, got %q, want %q", rebasedConfig.OSVersion, newBaseConfig.OSVersion)
 	}
 }
+
+// TestRebaseLayerType tests that when rebasing an image, if the new base doesn't have
+// the same layer type as the image to rebase, the new base layer type is used.
+func TestRebaseLayerType(t *testing.T) {
+	// Create an old base image with DockerLayer type
+	oldBaseLayer, err := random.Layer(100, types.DockerLayer)
+	if err != nil {
+		t.Fatalf("random.Layer (oldBase): %v", err)
+	}
+	oldBase, err := mutate.AppendLayers(empty.Image, oldBaseLayer)
+	if err != nil {
+		t.Fatalf("mutate.AppendLayers (oldBase): %v", err)
+	}
+
+	// Create an image to rebase with OCILayer type
+	origLayer, err := random.Layer(100, types.DockerLayer)
+	if err != nil {
+		t.Fatalf("random.Layer (orig): %v", err)
+	}
+	orig, err := mutate.Append(oldBase, mutate.Addendum{
+		Layer: origLayer,
+		History: v1.History{
+			Author:    "me",
+			Created:   v1.Time{Time: time.Now()},
+			CreatedBy: "test",
+			Comment:   "this is a test",
+		},
+	})
+	if err != nil {
+		t.Fatalf("mutate.Append (orig): %v", err)
+	}
+
+	// Create a new base image with DockerLayer type
+	newBaseLayer, err := random.Layer(100, types.OCILayer)
+	if err != nil {
+		t.Fatalf("random.Layer (newBase): %v", err)
+	}
+	newBase, err := mutate.AppendLayers(empty.Image, newBaseLayer)
+	if err != nil {
+		t.Fatalf("mutate.AppendLayers (newBase): %v", err)
+	}
+
+	// Rebase the image onto the new base
+	rebased, err := mutate.Rebase(orig, oldBase, newBase)
+	if err != nil {
+		t.Fatalf("mutate.Rebase: %v", err)
+	}
+
+	// Get the layers of the rebased image
+	rebasedLayers, err := rebased.Layers()
+	if err != nil {
+		t.Fatalf("rebased.Layers: %v", err)
+	}
+
+	// Check that the rebased image has 2 layers
+	if len(rebasedLayers) != 2 {
+		t.Fatalf("rebased image has %d layers, expected 2", len(rebasedLayers))
+	}
+
+	// Check that the first layer (from the new base) has OCILayer type
+	firstLayerType, err := rebasedLayers[0].MediaType()
+	if err != nil {
+		t.Fatalf("rebasedLayers[0].MediaType: %v", err)
+	}
+	if firstLayerType != types.OCILayer {
+		t.Errorf("first layer has media type %v, expected %v", firstLayerType, types.OCILayer)
+	}
+
+	// Check that the second layer (from orig) has OCILayer type (should have been changed from DockerLayer)
+	secondLayerType, err := rebasedLayers[1].MediaType()
+	if err != nil {
+		t.Fatalf("rebasedLayers[1].MediaType: %v", err)
+	}
+
+	// Print more information about the layers to help diagnose the issue
+	t.Logf("Original layer media type: %v", types.DockerLayer)
+	t.Logf("New base layer media type: %v", types.OCILayer)
+	t.Logf("First rebased layer media type: %v", firstLayerType)
+	t.Logf("Second rebased layer media type: %v", secondLayerType)
+
+	// Check if the original layer is actually a DockerLayer
+	origLayerType, err := origLayer.MediaType()
+	if err != nil {
+		t.Logf("Error getting original layer media type: %v", err)
+	} else {
+		t.Logf("Original layer actual media type: %v", origLayerType)
+	}
+
+	if secondLayerType != types.OCILayer {
+		t.Errorf("second layer has media type %v, expected %v", secondLayerType, types.OCILayer)
+	}
+}