feat: discard patch layer

go.mod

@@ -119,6 +119,7 @@ require (
 	github.com/spf13/afero v1.11.0 // indirect
 	github.com/spf13/cast v1.6.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	github.com/subosito/gotenv v1.6.0 // indirect
 	github.com/theupdateframework/notary v0.7.0 // indirect
 	github.com/tonistiigi/fsutil v0.0.0-20240424095704-91a3fc46842c // indirect

go.sum

@@ -422,6 +422,7 @@ github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/objx v0.2.0/go.mod h1:qt09Ya8vawLte6SNmTgCsAVtYtaKzEcn8ATUoHMkEqE=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
+github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
 github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=

mocks/mock_gwclient.go

@@ -0,0 +1,141 @@
+package mocks
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/moby/buildkit/client/llb"
+	"github.com/moby/buildkit/client/llb/sourceresolver"
+	gwclient "github.com/moby/buildkit/frontend/gateway/client"
+	"github.com/moby/buildkit/solver/pb"
+	"github.com/opencontainers/go-digest"
+	"github.com/stretchr/testify/mock"
+)
+
+// Mock for gwclient.Client.
+type MockGWClient struct {
+	mock.Mock
+}
+
+func (m *MockGWClient) ResolveSourceMetadata(ctx context.Context, op *pb.SourceOp, opt sourceresolver.Opt) (*sourceresolver.MetaResponse, error) {
+	args := m.Called(ctx, op, opt)
+
+	metaResponse, ok := args.Get(0).(*sourceresolver.MetaResponse)
+	if !ok {
+		return nil, fmt.Errorf("type assertion to *sourceresolver.MetaResponse failed")
+	}
+
+	return metaResponse, args.Error(1)
+}
+
+//nolint:gocritic
+func (m *MockGWClient) Solve(ctx context.Context, req gwclient.SolveRequest) (*gwclient.Result, error) {
+	args := m.Called(ctx, req)
+
+	result, ok := args.Get(0).(*gwclient.Result)
+	if !ok {
+		return nil, fmt.Errorf("type assertion to *gwclient.Result failed")
+	}
+
+	return result, args.Error(1)
+}
+
+func (m *MockGWClient) ResolveImageConfig(ctx context.Context, ref string, opt sourceresolver.Opt) (string, digest.Digest, []byte, error) {
+	args := m.Called(ctx, ref, opt)
+
+	digestResult, ok1 := args.Get(1).(digest.Digest)
+	if !ok1 {
+		return "", digest.Digest(""), nil, fmt.Errorf("type assertion to digest.Digest failed")
+	}
+
+	byteResult, ok2 := args.Get(2).([]byte)
+	if !ok2 {
+		return "", digest.Digest(""), nil, fmt.Errorf("type assertion to []byte failed")
+	}
+
+	return args.String(0), digestResult, byteResult, args.Error(3)
+}
+
+func (m *MockGWClient) BuildOpts() gwclient.BuildOpts {
+	args := m.Called()
+
+	buildOpts, ok := args.Get(0).(gwclient.BuildOpts)
+	if !ok {
+		panic("type assertion to gwclient.BuildOpts failed")
+	}
+
+	return buildOpts
+}
+
+func (m *MockGWClient) Inputs(ctx context.Context) (map[string]llb.State, error) {
+	args := m.Called(ctx)
+
+	stateMap, ok := args.Get(0).(map[string]llb.State)
+	if !ok {
+		return nil, fmt.Errorf("type assertion to map[string]llb.State failed")
+	}
+
+	return stateMap, args.Error(1)
+}
+
+//nolint:gocritic
+func (m *MockGWClient) NewContainer(ctx context.Context, req gwclient.NewContainerRequest) (gwclient.Container, error) {
+	args := m.Called(ctx, req)
+
+	container, ok := args.Get(0).(gwclient.Container)
+	if !ok {
+		return nil, fmt.Errorf("type assertion to gwclient.Container failed")
+	}
+
+	return container, args.Error(1)
+}
+
+//nolint:gocritic
+func (m *MockGWClient) Warn(ctx context.Context, dgst digest.Digest, msg string, opts gwclient.WarnOpts) error {
+	args := m.Called(ctx, dgst, msg, opts)
+
+	warnErr, ok := args.Get(0).(error)
+	if !ok {
+		return fmt.Errorf("type assertion to error failed")
+	}
+
+	return warnErr
+}
+
+// MockReference is a mock of the Reference interface.
+type MockReference struct {
+	mock.Mock
+}
+
+func (m *MockReference) ReadFile(ctx context.Context, req gwclient.ReadRequest) ([]byte, error) {
+	args := m.Called(ctx, req)
+
+	byteResult, ok := args.Get(0).([]byte)
+	if !ok {
+		return nil, fmt.Errorf("type assertion to []byte failed")
+	}
+
+	return byteResult, args.Error(1)
+}
+
+func (m *MockReference) ToState() (llb.State, error) {
+	args := m.Called()
+
+	state, ok := args.Get(0).(llb.State)
+	if !ok {
+		return state, fmt.Errorf("type assertion to llb.State failed")
+	}
+
+	return state, args.Error(1)
+}
+
+func (m *MockReference) Evaluate(ctx context.Context) error {
+	args := m.Called(ctx)
+
+	evalErr, ok := args.Get(0).(error)
+	if !ok {
+		return fmt.Errorf("type assertion to error failed")
+	}
+
+	return evalErr
+}

pkg/buildkit/buildkit.go

@@ -3,6 +3,8 @@
 import (
 	"bytes"
 	"context"
+	"encoding/json"
+	"fmt"
 
 	"github.com/containerd/platforms"
 	"github.com/moby/buildkit/client/llb"
@@ -12,11 +14,13 @@
 )
 
 type Config struct {
-	ImageName  string
-	Client     gwclient.Client
-	ConfigData []byte
-	Platform   *ispec.Platform
-	ImageState llb.State
+	ImageName         string
+	Client            gwclient.Client
+	ConfigData        []byte
+	PatchedConfigData []byte
+	Platform          *ispec.Platform
+	ImageState        llb.State
+	PatchedImageState llb.State
 }
 
 type Opts struct {
@@ -26,14 +30,14 @@
 	KeyPath    string
 }
 
-func InitializeBuildkitConfig(ctx context.Context, c gwclient.Client, image string) (*Config, error) {
+func InitializeBuildkitConfig(ctx context.Context, c gwclient.Client, userImage string) (*Config, error) {
 	// Initialize buildkit config for the target image
 	config := Config{
-		ImageName: image,
+		ImageName: userImage,
 	}
 
 	// Resolve and pull the config for the target image
-	_, _, configData, err := c.ResolveImageConfig(ctx, image, sourceresolver.Opt{
+	_, _, configData, err := c.ResolveImageConfig(ctx, userImage, sourceresolver.Opt{
 		ImageOpt: &sourceresolver.ResolveImageOpt{
 			ResolveMode: llb.ResolveModePreferLocal.String(),
 		},
@@ -42,23 +46,106 @@
 		return nil, err
 	}
 
-	config.ConfigData = configData
+	var baseImage string
+	config.ConfigData, config.PatchedConfigData, baseImage, err = updateImageConfigData(ctx, c, configData, userImage)
+	if err != nil {
+		return nil, err
+	}
 
 	// Load the target image state with the resolved image config in case environment variable settings
 	// are necessary for running apps in the target image for updates
-	config.ImageState, err = llb.Image(image,
+	config.ImageState, err = llb.Image(baseImage,
 		llb.ResolveModePreferLocal,
 		llb.WithMetaResolver(c),
 	).WithImageConfig(config.ConfigData)
 	if err != nil {
 		return nil, err
 	}
 
+	// Only set PatchedImageState if the user supplied a patched image
+	// An image is deemed to be a patched image if it contains one of two metadata values
+	// BaseImage or ispec.AnnotationBaseImageName
+	if config.PatchedConfigData != nil {
+		config.PatchedImageState, err = llb.Image(userImage,
+			llb.ResolveModePreferLocal,
+			llb.WithMetaResolver(c),
+		).WithImageConfig(config.PatchedConfigData)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	config.Client = c
 
 	return &config, nil
 }
 
+func updateImageConfigData(ctx context.Context, c gwclient.Client, configData []byte, image string) ([]byte, []byte, string, error) {
+	baseImage, userImageConfig, err := setupLabels(image, configData)
+	if err != nil {
+		return nil, nil, "", err
+	}
+
+	if baseImage == "" {
+		configData = userImageConfig
+	} else {
+		patchedImageConfig := userImageConfig
+		_, _, baseImageConfig, err := c.ResolveImageConfig(ctx, baseImage, sourceresolver.Opt{
+			ImageOpt: &sourceresolver.ResolveImageOpt{
+				ResolveMode: llb.ResolveModePreferLocal.String(),
+			},
+		})
+		if err != nil {
+			return nil, nil, "", err
+		}
+
+		_, baseImageWithLabels, _ := setupLabels(baseImage, baseImageConfig)
+		configData = baseImageWithLabels
+
+		return configData, patchedImageConfig, baseImage, nil
+	}
+
+	return configData, nil, image, nil
+}
+
+func setupLabels(image string, configData []byte) (string, []byte, error) {
+	imageConfig := make(map[string]interface{})
+	err := json.Unmarshal(configData, &imageConfig)
+	if err != nil {
+		return "", nil, err
+	}
+
+	configMap, ok := imageConfig["config"].(map[string]interface{})
+	if !ok {
+		err := fmt.Errorf("type assertion to map[string]interface{} failed")
+		return "", nil, err
+	}
+
+	var baseImage string
+	labels := configMap["labels"]
+	if labels == nil {
+		configMap["labels"] = make(map[string]interface{})
+	}
+	labelsMap, ok := configMap["labels"].(map[string]interface{})
+	if !ok {
+		err := fmt.Errorf("type assertion to map[string]interface{} failed")
+		return "", nil, err
+	}
+	if baseImageValue := labelsMap["BaseImage"]; baseImageValue != nil {
+		baseImage, ok = baseImageValue.(string)
+		if !ok {
+			err := fmt.Errorf("type assertion to string failed")
+			return "", nil, err
+		}
+	} else {
+		labelsMap["BaseImage"] = image
+	}
+
+	imageWithLabels, _ := json.Marshal(imageConfig)
+
+	return baseImage, imageWithLabels, nil
+}
+
 // Extracts the bytes of the file denoted by `path` from the state `st`.
 func ExtractFileFromState(ctx context.Context, c gwclient.Client, st *llb.State, path string) ([]byte, error) {
 	// since platform is obtained from host, override it in the case of Darwin