dockerfile: expose TARGETSTAGE as builtin argument

Allows capturing what build stage is currently being
built for patterns with better base stage reuse.

Signed-off-by: Tonis Tiigi <[email protected]>

Author: tonistiigi

frontend/dockerfile/dockerfile2llb/convert.go

@@ -227,10 +227,6 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 		opt.LLBCaps = &caps
 	}
 
-	platformOpt := buildPlatformOpt(&opt)
-
-	globalArgs := platformArgs(platformOpt, opt.BuildArgs)
-
 	dockerfile, err := parser.Parse(bytes.NewReader(dt))
 	if err != nil {
 		return nil, err
@@ -261,6 +257,13 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 	validateStageNames(stages, lint)
 	validateCommandCasing(stages, lint)
 
+	platformOpt := buildPlatformOpt(&opt)
+	targetName := opt.Target
+	if targetName == "" {
+		targetName = stages[len(stages)-1].Name
+	}
+	globalArgs := defaultArgs(platformOpt, opt.BuildArgs, targetName)
+
 	shlex := shell.NewLex(dockerfile.EscapeToken)
 	outline := newOutlineCapture()
 

frontend/dockerfile/dockerfile2llb/platform.go

@@ -36,9 +36,12 @@ func buildPlatformOpt(opt *ConvertOpt) *platformOpt {
 	}
 }
 
-func platformArgs(po *platformOpt, overrides map[string]string) *llb.EnvList {
+func defaultArgs(po *platformOpt, overrides map[string]string, target string) *llb.EnvList {
 	bp := po.buildPlatforms[0]
 	tp := po.targetPlatform
+	if target == "" {
+		target = "default"
+	}
 	s := [...][2]string{
 		{"BUILDPLATFORM", platforms.Format(bp)},
 		{"BUILDOS", bp.OS},
@@ -48,6 +51,7 @@ func platformArgs(po *platformOpt, overrides map[string]string) *llb.EnvList {
 		{"TARGETOS", tp.OS},
 		{"TARGETARCH", tp.Architecture},
 		{"TARGETVARIANT", tp.Variant},
+		{"TARGETSTAGE", target},
 	}
 	env := &llb.EnvList{}
 	for _, kv := range s {

frontend/dockerfile/dockerfile_test.go

@@ -207,6 +207,7 @@ var allTests = integration.TestFuncs(
 	testHistoryFinalizeTrace,
 	testEmptyStages,
 	testLocalCustomSessionID,
+	testTargetStageNameArg,
 )
 
 // Tests that depend on the `security.*` entitlements
@@ -1791,6 +1792,121 @@ COPY Dockerfile .
 	}
 }
 
+func testTargetStageNameArg(t *testing.T, sb integration.Sandbox) {
+	integration.SkipOnPlatform(t, "windows")
+	f := getFrontend(t, sb)
+
+	dockerfile := []byte(`
+FROM alpine AS base
+WORKDIR /out
+RUN echo -n "value:$TARGETSTAGE" > /out/first
+ARG TARGETSTAGE
+RUN echo -n "value:$TARGETSTAGE" > /out/second
+
+FROM scratch AS foo
+COPY --from=base /out/ /
+
+FROM scratch
+COPY --from=base /out/ /
+`)
+
+	dir := integration.Tmpdir(
+		t,
+		fstest.CreateFile("Dockerfile", dockerfile, 0600),
+	)
+
+	destDir := t.TempDir()
+
+	c, err := client.New(sb.Context(), sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	_, err = f.Solve(sb.Context(), c, client.SolveOpt{
+		LocalMounts: map[string]fsutil.FS{
+			dockerui.DefaultLocalNameDockerfile: dir,
+			dockerui.DefaultLocalNameContext:    dir,
+		},
+		FrontendAttrs: map[string]string{
+			"target": "foo",
+		},
+		Exports: []client.ExportEntry{
+			{
+				Type:      client.ExporterLocal,
+				OutputDir: destDir,
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	dt, err := os.ReadFile(filepath.Join(destDir, "first"))
+	require.NoError(t, err)
+	require.Equal(t, "value:", string(dt))
+
+	dt, err = os.ReadFile(filepath.Join(destDir, "second"))
+	require.NoError(t, err)
+	require.Equal(t, "value:foo", string(dt))
+
+	destDir = t.TempDir()
+
+	_, err = f.Solve(sb.Context(), c, client.SolveOpt{
+		LocalMounts: map[string]fsutil.FS{
+			dockerui.DefaultLocalNameDockerfile: dir,
+			dockerui.DefaultLocalNameContext:    dir,
+		},
+		Exports: []client.ExportEntry{
+			{
+				Type:      client.ExporterLocal,
+				OutputDir: destDir,
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	dt, err = os.ReadFile(filepath.Join(destDir, "first"))
+	require.NoError(t, err)
+	require.Equal(t, "value:", string(dt))
+
+	dt, err = os.ReadFile(filepath.Join(destDir, "second"))
+	require.NoError(t, err)
+	require.Equal(t, "value:default", string(dt))
+
+	// stage name defined in Dockerfile but not passed in request
+	dockerfile = append(dockerfile, []byte(`
+	
+	FROM scratch AS final
+	COPY --from=base /out/ /
+	`)...)
+
+	dir = integration.Tmpdir(
+		t,
+		fstest.CreateFile("Dockerfile", dockerfile, 0600),
+	)
+
+	destDir = t.TempDir()
+
+	_, err = f.Solve(sb.Context(), c, client.SolveOpt{
+		LocalMounts: map[string]fsutil.FS{
+			dockerui.DefaultLocalNameDockerfile: dir,
+			dockerui.DefaultLocalNameContext:    dir,
+		},
+		Exports: []client.ExportEntry{
+			{
+				Type:      client.ExporterLocal,
+				OutputDir: destDir,
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	dt, err = os.ReadFile(filepath.Join(destDir, "first"))
+	require.NoError(t, err)
+	require.Equal(t, "value:", string(dt))
+
+	dt, err = os.ReadFile(filepath.Join(destDir, "second"))
+	require.NoError(t, err)
+	require.Equal(t, "value:final", string(dt))
+}
+
 func testExportMultiPlatform(t *testing.T, sb integration.Sandbox) {
 	integration.SkipOnPlatform(t, "windows")
 	workers.CheckFeatureCompat(t, sb, workers.FeatureOCIExporter, workers.FeatureMultiPlatform)