Allow pulling partial layer chains from an image

Adds a new `WithLayerLimit` option to `llb.Image`
only pulls specified number of layers instead of
full image.

This can be used in combination with DiffOp/MergeOp
to pull any subset of layers from an image in any order.

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

client/client_test.go

@@ -161,6 +161,7 @@ func TestIntegration(t *testing.T) {
 		testUncompressedRegistryCacheImportExport,
 		testStargzLazyRegistryCacheImportExport,
 		testCallInfo,
+		testPullWithLayerLimit,
 	)
 	tests = append(tests, diffOpTestCases()...)
 	integration.Run(t, tests, mirrors)
@@ -5728,6 +5729,115 @@ func testBuildInfoNoExport(t *testing.T, sb integration.Sandbox) {
 	require.Equal(t, exbi.Sources[0].Ref, "docker.io/library/busybox:latest")
 }
 
+func testPullWithLayerLimit(t *testing.T, sb integration.Sandbox) {
+	integration.SkipIfDockerd(t, sb, "direct push")
+	requiresLinux(t)
+	c, err := New(sb.Context(), sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	st := llb.Scratch().
+		File(llb.Mkfile("/first", 0644, []byte("first"))).
+		File(llb.Mkfile("/second", 0644, []byte("second"))).
+		File(llb.Mkfile("/third", 0644, []byte("third")))
+
+	def, err := st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	registry, err := sb.NewRegistry()
+	if errors.Is(err, integration.ErrRequirements) {
+		t.Skip(err.Error())
+	}
+	require.NoError(t, err)
+
+	target := registry + "/buildkit/testlayers:latest"
+
+	_, err = c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type: ExporterImage,
+				Attrs: map[string]string{
+					"name": target,
+					"push": "true",
+				},
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	// pull 2 first layers
+	st = llb.Image(target, llb.WithLayerLimit(2)).
+		File(llb.Mkfile("/forth", 0644, []byte("forth")))
+
+	def, err = st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	destDir, err := os.MkdirTemp("", "buildkit")
+	require.NoError(t, err)
+	defer os.RemoveAll(destDir)
+
+	_, err = c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{{
+			Type:      ExporterLocal,
+			OutputDir: destDir,
+		}},
+	}, nil)
+	require.NoError(t, err)
+
+	dt, err := os.ReadFile(filepath.Join(destDir, "first"))
+	require.NoError(t, err)
+	require.Equal(t, string(dt), "first")
+
+	dt, err = os.ReadFile(filepath.Join(destDir, "second"))
+	require.NoError(t, err)
+	require.Equal(t, string(dt), "second")
+
+	_, err = os.ReadFile(filepath.Join(destDir, "third"))
+	require.Error(t, err)
+	require.True(t, errors.Is(err, os.ErrNotExist))
+
+	dt, err = os.ReadFile(filepath.Join(destDir, "forth"))
+	require.NoError(t, err)
+	require.Equal(t, string(dt), "forth")
+
+	// pull 3rd layer only
+	st = llb.Diff(
+		llb.Image(target, llb.WithLayerLimit(2)),
+		llb.Image(target)).
+		File(llb.Mkfile("/forth", 0644, []byte("forth")))
+
+	def, err = st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	destDir, err = os.MkdirTemp("", "buildkit")
+	require.NoError(t, err)
+	defer os.RemoveAll(destDir)
+
+	_, err = c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{{
+			Type:      ExporterLocal,
+			OutputDir: destDir,
+		}},
+	}, nil)
+	require.NoError(t, err)
+
+	_, err = os.ReadFile(filepath.Join(destDir, "first"))
+	require.Error(t, err)
+	require.True(t, errors.Is(err, os.ErrNotExist))
+
+	_, err = os.ReadFile(filepath.Join(destDir, "second"))
+	require.Error(t, err)
+	require.True(t, errors.Is(err, os.ErrNotExist))
+
+	dt, err = os.ReadFile(filepath.Join(destDir, "third"))
+	require.NoError(t, err)
+	require.Equal(t, string(dt), "third")
+
+	dt, err = os.ReadFile(filepath.Join(destDir, "forth"))
+	require.NoError(t, err)
+	require.Equal(t, string(dt), "forth")
+}
+
 func testCallInfo(t *testing.T, sb integration.Sandbox) {
 	c, err := New(sb.Context(), sb.Address())
 	require.NoError(t, err)

client/llb/resolver.go

@@ -23,6 +23,12 @@ func ResolveDigest(v bool) ImageOption {
 	})
 }
 
+func WithLayerLimit(l int) ImageOption {
+	return imageOptionFunc(func(ii *ImageInfo) {
+		ii.layerLimit = &l
+	})
+}
+
 // ImageMetaResolver can resolve image config metadata from a reference
 type ImageMetaResolver interface {
 	ResolveImageConfig(ctx context.Context, ref string, opt ResolveImageConfigOpt) (digest.Digest, []byte, error)

client/llb/source.go

@@ -116,6 +116,11 @@ func Image(ref string, opts ...ImageOption) State {
 		attrs[pb.AttrImageRecordType] = info.RecordType
 	}
 
+	if ll := info.layerLimit; ll != nil {
+		attrs[pb.AttrImageLayerLimit] = strconv.FormatInt(int64(*ll), 10)
+		addCap(&info.Constraints, pb.CapSourceImageLayerLimit)
+	}
+
 	src := NewSource("docker-image://"+ref, attrs, info.Constraints) // controversial
 	if err != nil {
 		src.err = err
@@ -204,6 +209,7 @@ type ImageInfo struct {
 	metaResolver  ImageMetaResolver
 	resolveDigest bool
 	resolveMode   ResolveMode
+	layerLimit    *int
 	RecordType    string
 }
 

solver/pb/attr.go

@@ -26,6 +26,7 @@ const AttrImageResolveModeDefault = "default"
 const AttrImageResolveModeForcePull = "pull"
 const AttrImageResolveModePreferLocal = "local"
 const AttrImageRecordType = "image.recordtype"
+const AttrImageLayerLimit = "image.layerlimit"
 
 const AttrLocalDiffer = "local.differ"
 const AttrLocalDifferNone = "none"

solver/pb/caps.go

@@ -9,8 +9,10 @@ var Caps apicaps.CapList
 // considered immutable. After a capability is marked stable it should not be disabled.
 
 const (
-	CapSourceImage                apicaps.CapID = "source.image"
-	CapSourceImageResolveMode     apicaps.CapID = "source.image.resolvemode"
+	CapSourceImage            apicaps.CapID = "source.image"
+	CapSourceImageResolveMode apicaps.CapID = "source.image.resolvemode"
+	CapSourceImageLayerLimit  apicaps.CapID = "source.image.layerlimit"
+
 	CapSourceLocal                apicaps.CapID = "source.local"
 	CapSourceLocalUnique          apicaps.CapID = "source.local.unique"
 	CapSourceLocalSessionID       apicaps.CapID = "source.local.sessionid"
@@ -86,6 +88,12 @@ func init() {
 		Status:  apicaps.CapStatusExperimental,
 	})
 
+	Caps.Init(apicaps.Cap{
+		ID:      CapSourceImageLayerLimit,
+		Enabled: true,
+		Status:  apicaps.CapStatusExperimental,
+	})
+
 	Caps.Init(apicaps.Cap{
 		ID:      CapSourceLocal,
 		Enabled: true,

source/containerimage/pull.go

@@ -152,17 +152,19 @@ type puller struct {
 	*pull.Puller
 }
 
-func mainManifestKey(ctx context.Context, desc ocispecs.Descriptor, platform ocispecs.Platform) (digest.Digest, error) {
+func mainManifestKey(ctx context.Context, desc ocispecs.Descriptor, platform ocispecs.Platform, layerLimit *int) (digest.Digest, error) {
 	dt, err := json.Marshal(struct {
 		Digest  digest.Digest
 		OS      string
 		Arch    string
 		Variant string `json:",omitempty"`
+		Limit   *int   `json:",omitempty"`
 	}{
 		Digest:  desc.Digest,
 		OS:      platform.OS,
 		Arch:    platform.Architecture,
 		Variant: platform.Variant,
+		Limit:   layerLimit,
 	})
 	if err != nil {
 		return "", err
@@ -178,7 +180,7 @@ func (p *puller) CacheKey(ctx context.Context, g session.Group, index int) (cach
 	progressFactory := progress.FromContext(ctx)
 
 	_, err = p.g.Do(ctx, "", func(ctx context.Context) (_ interface{}, err error) {
-		if p.cacheKeyErr != nil || p.cacheKeyDone == true {
+		if p.cacheKeyErr != nil || p.cacheKeyDone {
 			return nil, p.cacheKeyErr
 		}
 		defer func() {
@@ -203,6 +205,13 @@ func (p *puller) CacheKey(ctx context.Context, g session.Group, index int) (cach
 			return nil, err
 		}
 
+		if ll := p.id.LayerLimit; ll != nil {
+			if *ll > len(p.manifest.Descriptors) {
+				return nil, errors.Errorf("layer limit %d is greater than the number of layers in the image %d", *ll, len(p.manifest.Descriptors))
+			}
+			p.manifest.Descriptors = p.manifest.Descriptors[:*ll]
+		}
+
 		if len(p.manifest.Descriptors) > 0 {
 			progressController := &controller.Controller{
 				WriterFactory: progressFactory,
@@ -233,7 +242,7 @@ func (p *puller) CacheKey(ctx context.Context, g session.Group, index int) (cach
 		}
 
 		desc := p.manifest.MainManifestDesc
-		k, err := mainManifestKey(ctx, desc, p.Platform)
+		k, err := mainManifestKey(ctx, desc, p.Platform, p.id.LayerLimit)
 		if err != nil {
 			return nil, err
 		}
@@ -243,7 +252,11 @@ func (p *puller) CacheKey(ctx context.Context, g session.Group, index int) (cach
 		if err != nil {
 			return nil, err
 		}
-		p.configKey = cacheKeyFromConfig(dt).String()
+		ck, err := cacheKeyFromConfig(dt, p.id.LayerLimit)
+		if err != nil {
+			return nil, err
+		}
+		p.configKey = ck.String()
 		p.cacheKeyDone = true
 		return nil, nil
 	})
@@ -336,16 +349,27 @@ func (p *puller) Snapshot(ctx context.Context, g session.Group) (ir cache.Immuta
 
 // cacheKeyFromConfig returns a stable digest from image config. If image config
 // is a known oci image we will use chainID of layers.
-func cacheKeyFromConfig(dt []byte) digest.Digest {
+func cacheKeyFromConfig(dt []byte, layerLimit *int) (digest.Digest, error) {
 	var img ocispecs.Image
 	err := json.Unmarshal(dt, &img)
 	if err != nil {
-		return digest.FromBytes(dt)
+		if layerLimit != nil {
+			return "", errors.Wrap(err, "failed to parse image config")
+		}
+		return digest.FromBytes(dt), nil // digest of config
+	}
+	if layerLimit != nil {
+		l := *layerLimit
+		if len(img.RootFS.DiffIDs) < l {
+			return "", errors.Errorf("image has %d layers, limit is %d", len(img.RootFS.DiffIDs), l)
+		}
+		img.RootFS.DiffIDs = img.RootFS.DiffIDs[:l]
 	}
 	if img.RootFS.Type != "layers" || len(img.RootFS.DiffIDs) == 0 {
-		return ""
+		return "", nil
 	}
-	return identity.ChainID(img.RootFS.DiffIDs)
+
+	return identity.ChainID(img.RootFS.DiffIDs), nil
 }
 
 func oneOffProgress(ctx context.Context, id string) func(err error) error {

source/identifier.go

@@ -85,6 +85,12 @@ func FromLLB(op *pb.Op_Source, platform *pb.Platform) (Identifier, error) {
 					return nil, err
 				}
 				id.RecordType = rt
+			case pb.AttrImageLayerLimit:
+				l, err := strconv.Atoi(v)
+				if err != nil {
+					return nil, errors.Wrapf(err, "invalid layer limit %s", v)
+				}
+				id.LayerLimit = &l
 			}
 		}
 	}
@@ -190,6 +196,7 @@ type ImageIdentifier struct {
 	Platform    *ocispecs.Platform
 	ResolveMode ResolveMode
 	RecordType  client.UsageRecordType
+	LayerLimit  *int
 }
 
 func NewImageIdentifier(str string) (*ImageIdentifier, error) {