dockerfile2llb: filter unused paths for named contexts

frontend/dockerfile/dockerfile2llb/convert.go

@@ -153,11 +153,7 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 			if copt.Platform == nil {
 				copt.Platform = opt.TargetPlatform
 			}
-			st, img, err := opt.Client.NamedContext(ctx, name, copt)
-			if err != nil {
-				return nil, nil, err
-			}
-			return st, img, nil
+			return opt.Client.NamedContext(ctx, name, copt)
 		}
 		return nil, nil, nil
 	}
@@ -237,6 +233,7 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 			stage:          st,
 			deps:           make(map[*dispatchState]struct{}),
 			ctxPaths:       make(map[string]struct{}),
+			paths:          make(map[string]struct{}),
 			stageName:      st.Name,
 			prefixPlatform: opt.MultiPlatformRequested,
 			outline:        outline.clone(),
@@ -260,7 +257,11 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 		}
 
 		if st.Name != "" {
-			s, img, err := namedContext(ctx, st.Name, dockerui.ContextOpt{Platform: ds.platform, ResolveMode: opt.ImageResolveMode.String()})
+			s, img, err := namedContext(ctx, st.Name, dockerui.ContextOpt{
+				Platform:       ds.platform,
+				ResolveMode:    opt.ImageResolveMode.String(),
+				AsyncLocalOpts: ds.asyncLocalOpts,
+			})
 			if err != nil {
 				return nil, err
 			}
@@ -391,7 +392,11 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 					d.stage.BaseName = reference.TagNameOnly(ref).String()
 
 					var isScratch bool
-					st, img, err := namedContext(ctx, d.stage.BaseName, dockerui.ContextOpt{ResolveMode: opt.ImageResolveMode.String(), Platform: platform})
+					st, img, err := namedContext(ctx, d.stage.BaseName, dockerui.ContextOpt{
+						ResolveMode:    opt.ImageResolveMode.String(),
+						Platform:       platform,
+						AsyncLocalOpts: d.asyncLocalOpts,
+					})
 					if err != nil {
 						return err
 					}
@@ -492,11 +497,23 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 			d.state = d.base.state
 			d.platform = d.base.platform
 			d.image = clone(d.base.image)
+			// Utilize the same path index as our base image so we propagate
+			// the paths we use back to the base image.
+			d.paths = d.base.paths
+		}
+
+		// Ensure platform is set.
+		if d.platform == nil {
+			d.platform = &d.opt.targetPlatform
 		}
 
 		// make sure that PATH is always set
 		if _, ok := shell.BuildEnvs(d.image.Config.Env)["PATH"]; !ok {
-			d.image.Config.Env = append(d.image.Config.Env, "PATH="+system.DefaultPathEnv(d.platform.OS))
+			var osName string
+			if d.platform != nil {
+				osName = d.platform.OS
+			}
+			d.image.Config.Env = append(d.image.Config.Env, "PATH="+system.DefaultPathEnv(osName))
 		}
 
 		// initialize base metadata from image conf
@@ -565,17 +582,19 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 		}
 	}
 
+	// Ensure the entirety of the target state is marked as used.
+	// This is done after we've already evaluated every stage to ensure
+	// the paths attribute is set correctly.
+	target.paths["/"] = struct{}{}
+
 	if len(opt.Labels) != 0 && target.image.Config.Labels == nil {
 		target.image.Config.Labels = make(map[string]string, len(opt.Labels))
 	}
 	for k, v := range opt.Labels {
 		target.image.Config.Labels[k] = v
 	}
 
-	opts := []llb.LocalOption{}
-	if includePatterns := normalizeContextPaths(ctxPaths); includePatterns != nil {
-		opts = append(opts, llb.FollowPaths(includePatterns))
-	}
+	opts := filterPaths(ctxPaths)
 	bctx := opt.MainContext
 	if opt.Client != nil {
 		bctx, err = opt.Client.MainContext(ctx, opts...)
@@ -630,6 +649,7 @@ func toCommand(ic instructions.Command, allDispatchStates *dispatchStates) (comm
 					stn = &dispatchState{
 						stage:        instructions.Stage{BaseName: c.From, Location: ic.Location()},
 						deps:         make(map[*dispatchState]struct{}),
+						paths:        make(map[string]struct{}),
 						unregistered: true,
 					}
 				}
@@ -773,9 +793,19 @@ func dispatch(d *dispatchState, cmd command, opt dispatchOpt) error {
 			location:     c.Location(),
 			opt:          opt,
 		})
-		if err == nil && len(cmd.sources) == 0 {
-			for _, src := range c.SourcePaths {
-				d.ctxPaths[path.Join("/", filepath.ToSlash(src))] = struct{}{}
+		if err == nil {
+			if len(cmd.sources) == 0 {
+				for _, src := range c.SourcePaths {
+					d.ctxPaths[path.Join("/", filepath.ToSlash(src))] = struct{}{}
+				}
+			} else {
+				source := cmd.sources[0]
+				if source.paths == nil {
+					source.paths = make(map[string]struct{})
+				}
+				for _, src := range c.SourcePaths {
+					source.paths[path.Join("/", filepath.ToSlash(src))] = struct{}{}
+				}
 			}
 		}
 	default:
@@ -784,17 +814,20 @@ func dispatch(d *dispatchState, cmd command, opt dispatchOpt) error {
 }
 
 type dispatchState struct {
-	opt            dispatchOpt
-	state          llb.State
-	image          image.Image
-	platform       *ocispecs.Platform
-	stage          instructions.Stage
-	base           *dispatchState
-	noinit         bool
-	deps           map[*dispatchState]struct{}
-	buildArgs      []instructions.KeyValuePairOptional
-	commands       []command
-	ctxPaths       map[string]struct{}
+	opt       dispatchOpt
+	state     llb.State
+	image     image.Image
+	platform  *ocispecs.Platform
+	stage     instructions.Stage
+	base      *dispatchState
+	noinit    bool
+	deps      map[*dispatchState]struct{}
+	buildArgs []instructions.KeyValuePairOptional
+	commands  []command
+	// ctxPaths marks the paths this dispatchState uses from the build context.
+	ctxPaths map[string]struct{}
+	// paths marks the paths that are used by this dispatchState.
+	paths          map[string]struct{}
 	ignoreCache    bool
 	cmdSet         bool
 	unregistered   bool
@@ -808,6 +841,10 @@ type dispatchState struct {
 	scanContext    bool
 }
 
+func (ds *dispatchState) asyncLocalOpts() []llb.LocalOption {
+	return filterPaths(ds.paths)
+}
+
 type dispatchStates struct {
 	states       []*dispatchState
 	statesByName map[string]*dispatchState
@@ -890,6 +927,9 @@ func dispatchRun(d *dispatchState, c *instructions.RunCommand, proxy *llb.ProxyE
 
 	customname := c.String()
 
+	// Run command can potentially access any file. Mark the full filesystem as used.
+	d.paths["/"] = struct{}{}
+
 	var args []string = c.CmdLine
 	if len(c.Files) > 0 {
 		if len(args) != 1 || !c.PrependShell {
@@ -1603,6 +1643,11 @@ func hasCircularDependency(states []*dispatchState) (bool, *dispatchState) {
 }
 
 func normalizeContextPaths(paths map[string]struct{}) []string {
+	// Avoid a useless allocation if the set of paths is empty.
+	if len(paths) == 0 {
+		return nil
+	}
+
 	pathSlice := make([]string, 0, len(paths))
 	for p := range paths {
 		if p == "/" {
@@ -1617,6 +1662,15 @@ func normalizeContextPaths(paths map[string]struct{}) []string {
 	return pathSlice
 }
 
+// filterPaths returns the local options required to filter an llb.Local
+// to only the required paths.
+func filterPaths(paths map[string]struct{}) []llb.LocalOption {
+	if includePaths := normalizeContextPaths(paths); len(includePaths) > 0 {
+		return []llb.LocalOption{llb.FollowPaths(includePaths)}
+	}
+	return nil
+}
+
 func proxyEnvFromBuildArgs(args map[string]string) *llb.ProxyEnv {
 	pe := &llb.ProxyEnv{}
 	isNil := true

frontend/dockerfile/dockerfile2llb/convert_runmount.go

@@ -31,6 +31,7 @@ func detectRunMount(cmd *command, allDispatchStates *dispatchStates) bool {
 				stn = &dispatchState{
 					stage:        instructions.Stage{BaseName: from},
 					deps:         make(map[*dispatchState]struct{}),
+					paths:        make(map[string]struct{}),
 					unregistered: true,
 				}
 			}
@@ -136,6 +137,9 @@ func dispatchRunMounts(d *dispatchState, c *instructions.RunCommand, sources []*
 
 		if mount.From == "" {
 			d.ctxPaths[path.Join("/", filepath.ToSlash(mount.Source))] = struct{}{}
+		} else {
+			source := sources[i]
+			source.paths[path.Join("/", filepath.ToSlash(mount.Source))] = struct{}{}
 		}
 	}
 	return out, nil

frontend/dockerfile/dockerfile_test.go

@@ -8,19 +8,22 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"math"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"os/exec"
 	"path"
 	"path/filepath"
+	"regexp"
 	"runtime"
 	"sort"
 	"strings"
 	"testing"
 	"time"
 
 	v1 "github.com/moby/buildkit/cache/remotecache/v1"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/containerd/containerd"
 	"github.com/containerd/containerd/content"
@@ -130,6 +133,7 @@ var allTests = integration.TestFuncs(
 	testNamedOCILayoutContextExport,
 	testNamedInputContext,
 	testNamedMultiplatformInputContext,
+	testNamedFilteredContext,
 	testEmptyDestDir,
 	testCopyChownCreateDest,
 	testCopyThroughSymlinkContext,
@@ -6063,6 +6067,114 @@ COPY --from=build /foo /out /
 	require.Equal(t, "foo is bar-arm64\n", string(dt))
 }
 
+func testNamedFilteredContext(t *testing.T, sb integration.Sandbox) {
+	ctx := sb.Context()
+
+	c, err := client.New(ctx, sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	fooDir := integration.Tmpdir(t,
+		// small file
+		fstest.CreateFile("foo", []byte(`foo`), 0600),
+		// blank file that's just large
+		fstest.CreateFile("bar", make([]byte, 4096*1000), 0600),
+	)
+
+	f := getFrontend(t, sb)
+
+	runTest := func(t *testing.T, dockerfile []byte, target string, min, max int64) {
+		t.Run(target, func(t *testing.T) {
+			dir := integration.Tmpdir(
+				t,
+				fstest.CreateFile(dockerui.DefaultDockerfileName, dockerfile, 0600),
+			)
+
+			ch := make(chan *client.SolveStatus)
+
+			eg, ctx := errgroup.WithContext(sb.Context())
+			eg.Go(func() error {
+				_, err := f.Solve(ctx, c, client.SolveOpt{
+					FrontendAttrs: map[string]string{
+						"context:foo": "local:foo",
+						"target":      target,
+					},
+					LocalDirs: map[string]string{
+						dockerui.DefaultLocalNameDockerfile: dir,
+						dockerui.DefaultLocalNameContext:    dir,
+						"foo":                               fooDir,
+					},
+				}, ch)
+				return err
+			})
+
+			eg.Go(func() error {
+				transferred := make(map[string]int64)
+				re := regexp.MustCompile(`transferring (.+):`)
+				for ss := range ch {
+					for _, status := range ss.Statuses {
+						m := re.FindStringSubmatch(status.ID)
+						if m == nil {
+							continue
+						}
+
+						ctxName := m[1]
+						transferred[ctxName] = status.Current
+					}
+				}
+
+				if foo := transferred["foo"]; foo < min {
+					return errors.Errorf("not enough data was transferred, %d < %d", foo, min)
+				} else if foo > max {
+					return errors.Errorf("too much data was transferred, %d > %d", foo, max)
+				}
+				return nil
+			})
+
+			err := eg.Wait()
+			require.NoError(t, err)
+		})
+	}
+
+	dockerfileBase := []byte(`
+FROM scratch AS copy_from
+COPY --from=foo /foo /
+
+FROM alpine AS run_mount
+RUN --mount=from=foo,src=/foo,target=/in/foo cp /in/foo /foo
+
+FROM foo AS image_source
+COPY --from=alpine / /
+RUN cat /foo > /bar
+
+FROM scratch AS all
+COPY --link --from=copy_from /foo /foo.b
+COPY --link --from=run_mount /foo /foo.c
+COPY --link --from=image_source /bar /foo.d
+`)
+
+	t.Run("new", func(t *testing.T) {
+		runTest(t, dockerfileBase, "run_mount", 1, 1024)
+		runTest(t, dockerfileBase, "copy_from", 1, 1024)
+		runTest(t, dockerfileBase, "image_source", 4096*1000, math.MaxInt64)
+		runTest(t, dockerfileBase, "all", 4096*1000, math.MaxInt64)
+	})
+
+	dockerfileFull := append([]byte(`
+FROM scratch AS foo
+COPY <<EOF /foo
+test
+EOF
+`), dockerfileBase...)
+
+	t.Run("replace", func(t *testing.T) {
+		runTest(t, dockerfileFull, "run_mount", 1, 1024)
+		runTest(t, dockerfileFull, "copy_from", 1, 1024)
+		runTest(t, dockerfileFull, "image_source", 4096*1000, math.MaxInt64)
+		runTest(t, dockerfileFull, "all", 4096*1000, math.MaxInt64)
+	})
+}
+
 func testSourceDateEpochWithoutExporter(t *testing.T, sb integration.Sandbox) {
 	workers.CheckFeatureCompat(t, sb, workers.FeatureOCIExporter, workers.FeatureSourceDateEpoch)
 	f := getFrontend(t, sb)

frontend/dockerui/config.go

@@ -95,7 +95,7 @@ type Source struct {
 
 type ContextOpt struct {
 	NoDockerignore bool
-	LocalOpts      []llb.LocalOption
+	AsyncLocalOpts func() []llb.LocalOption
 	Platform       *ocispecs.Platform
 	ResolveMode    string
 	CaptureDigest  *digest.Digest
@@ -473,11 +473,8 @@ func (bc *Client) NamedContext(ctx context.Context, name string, opt ContextOpt)
 	}
 	pname := name + "::" + platforms.Format(platforms.Normalize(pp))
 	st, img, err := bc.namedContext(ctx, name, pname, opt)
-	if err != nil {
-		return nil, nil, err
-	}
-	if st != nil {
-		return st, img, nil
+	if err != nil || st != nil {
+		return st, img, err
 	}
 	return bc.namedContext(ctx, name, name, opt)
 }