4 files changed, 23 insertions, 6 deletions
diff --git a/builder/builder.go b/builder/builder.go
index 8f33485250..d3521ddfbb 100644
--- a/builder/builder.go
+++ b/builder/builder.go
@@ -14,6 +14,7 @@ import (
 	containerpkg "github.com/docker/docker/container"
 	"github.com/docker/docker/image"
 	"github.com/docker/docker/layer"
+	"github.com/opencontainers/go-digest"
 )
 
 const (
@@ -45,7 +46,7 @@ type Backend interface {
 	// ContainerCreateWorkdir creates the workdir
 	ContainerCreateWorkdir(containerID string) error
 
-	CreateImage(config []byte, parent string) (Image, error)
+	CreateImage(ctx context.Context, config []byte, parent string, contentStoreDigest digest.Digest) (Image, error)
 
 	ImageCacheBuilder
 }
@@ -104,6 +105,7 @@ type ROLayer interface {
 	Release() error
 	NewRWLayer() (RWLayer, error)
 	DiffID() layer.DiffID
+	ContentStoreDigest() digest.Digest
 }
 
 // RWLayer is active layer that can be read/modified
diff --git a/builder/dockerfile/internals.go b/builder/dockerfile/internals.go
index 38d0aa706d..050deb1aad 100644
--- a/builder/dockerfile/internals.go
+++ b/builder/dockerfile/internals.go
@@ -63,7 +63,7 @@ func (b *Builder) commitContainer(ctx context.Context, dispatchState *dispatchSt
 	return err
 }
 
-func (b *Builder) exportImage(state *dispatchState, layer builder.RWLayer, parent builder.Image, runConfig *container.Config) error {
+func (b *Builder) exportImage(ctx context.Context, state *dispatchState, layer builder.RWLayer, parent builder.Image, runConfig *container.Config) error {
 	newLayer, err := layer.Commit()
 	if err != nil {
 		return err
@@ -98,7 +98,15 @@ func (b *Builder) exportImage(state *dispatchState, layer builder.RWLayer, paren
 		return errors.Wrap(err, "failed to encode image config")
 	}
 
-	exportedImage, err := b.docker.CreateImage(config, state.imageID)
+	// when writing the new image's manifest, we now need to pass in the new layer's digest.
+	// before the containerd store work this was unnecessary since we get the layer id
+	// from the image's RootFS ChainID -- see:
+	// https://github.com/moby/moby/blob/8cf66ed7322fa885ef99c4c044fa23e1727301dc/image/store.go#L162
+	// however, with the containerd store we can't do this. An alternative implementation here
+	// without changing the signature would be to get the layer digest by walking the content store
+	// and filtering the objects to find the layer with the DiffID we want, but that has performance
+	// implications that should be called out/investigated
+	exportedImage, err := b.docker.CreateImage(ctx, config, state.imageID, newLayer.ContentStoreDigest())
 	if err != nil {
 		return errors.Wrapf(err, "failed to export image")
 	}
@@ -170,7 +178,7 @@ func (b *Builder) performCopy(ctx context.Context, req dispatchRequest, inst cop
 			return errors.Wrapf(err, "failed to copy files")
 		}
 	}
-	return b.exportImage(state, rwLayer, imageMount.Image(), runConfigWithCommentCmd)
+	return b.exportImage(ctx, state, rwLayer, imageMount.Image(), runConfigWithCommentCmd)
 }
 
 func createDestInfo(workingDir string, inst copyInstruction, rwLayer builder.RWLayer, platform string) (copyInfo, error) {
diff --git a/builder/dockerfile/internals_test.go b/builder/dockerfile/internals_test.go
index 8145fac90d..f56c44b751 100644
--- a/builder/dockerfile/internals_test.go
+++ b/builder/dockerfile/internals_test.go
@@ -1,6 +1,7 @@
 package dockerfile // import "github.com/docker/docker/builder/dockerfile"
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"runtime"
@@ -193,6 +194,7 @@ type MockROLayer struct {
 	diffID layer.DiffID
 }
 
+func (l *MockROLayer) ContentStoreDigest() digest.Digest    { return "" }
 func (l *MockROLayer) Release() error                       { return nil }
 func (l *MockROLayer) NewRWLayer() (builder.RWLayer, error) { return nil, nil }
 func (l *MockROLayer) DiffID() layer.DiffID                 { return l.diffID }
@@ -217,6 +219,6 @@ func TestExportImage(t *testing.T) {
 		imageSources: getMockImageSource(nil, nil, nil),
 		docker:       getMockBuildBackend(),
 	}
-	err := b.exportImage(ds, layer, parentImage, runConfig)
+	err := b.exportImage(context.TODO(), ds, layer, parentImage, runConfig)
 	assert.NilError(t, err)
 }
diff --git a/builder/dockerfile/mockbackend_test.go b/builder/dockerfile/mockbackend_test.go
index 5b5419cfd6..a9e43e9bd1 100644
--- a/builder/dockerfile/mockbackend_test.go
+++ b/builder/dockerfile/mockbackend_test.go
@@ -13,6 +13,7 @@ import (
 	containerpkg "github.com/docker/docker/container"
 	"github.com/docker/docker/image"
 	"github.com/docker/docker/layer"
+	"github.com/opencontainers/go-digest"
 )
 
 // MockBackend implements the builder.Backend interface for unit testing
@@ -80,7 +81,7 @@ func (m *MockBackend) MakeImageCache(ctx context.Context, cacheFrom []string) (b
 	return nil, nil
 }
 
-func (m *MockBackend) CreateImage(config []byte, parent string) (builder.Image, error) {
+func (m *MockBackend) CreateImage(ctx context.Context, config []byte, parent string, layerDigest digest.Digest) (builder.Image, error) {
 	return &mockImage{id: "test"}, nil
 }
 
@@ -119,6 +120,10 @@ func (mic *mockImageCache) GetCache(parentID string, cfg *container.Config) (str
 
 type mockLayer struct{}
 
+func (l *mockLayer) ContentStoreDigest() digest.Digest {
+	return ""
+}
+
 func (l *mockLayer) Release() error {
 	return nil
 }