From 0615bb4b9de2541af1d7aa4e2022108b728374d7 Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Thu, 18 Jan 2024 23:30:51 -0800
Subject: [PATCH] dockerfile: validate order when linking stages

Dockerfile stages always needed to be in order (you can not refer
to a stage that is defined after your current command), but the
behavior for incorrect order was undefined instead of consistently
erroring.

Previously copying from a stage defined later would usually cause
an error about a missing source file that could have been quite
confusing. When the source path was "/" then the behavior could have
been worse and resulted in COPY passing without error, without any
files actually being copied (or only images from base image being
copied).

This validates the order when on COPY and mount dispatch, returns
proper error with correct source code location.

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 frontend/dockerfile/dockerfile2llb/convert.go | 15 +++++++-
 .../dockerfile2llb/convert_runmount.go        |  6 ++-
 frontend/dockerfile/dockerfile_test.go        | 37 +++++++++++++++++++
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/frontend/dockerfile/dockerfile2llb/convert.go b/frontend/dockerfile/dockerfile2llb/convert.go
index 4bb9d8903b88..a7375949e8bf 100644
--- a/frontend/dockerfile/dockerfile2llb/convert.go
+++ b/frontend/dockerfile/dockerfile2llb/convert.go
@@ -371,6 +371,9 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 				d.state = llb.Scratch()
 				d.image = emptyImage(platformOpt.targetPlatform)
 				d.platform = &platformOpt.targetPlatform
+				if d.unregistered {
+					d.noinit = true
+				}
 				continue
 			}
 			func(i int, d *dispatchState) {
@@ -379,6 +382,10 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 						if err != nil {
 							err = parser.WithLocation(err, d.stage.Location)
 						}
+						if d.unregistered {
+							// implicit stages don't need further dispatch
+							d.noinit = true
+						}
 					}()
 					origName := d.stage.BaseName
 					ref, err := reference.ParseNormalizedNamed(d.stage.BaseName)
@@ -492,6 +499,8 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 		if !isReachable(target, d) || d.noinit {
 			continue
 		}
+		// mark as initialized, used to determine states that have not been dispatched yet
+		d.noinit = true
 
 		if d.base != nil {
 			d.state = d.base.state
@@ -779,7 +788,11 @@ func dispatch(d *dispatchState, cmd command, opt dispatchOpt) error {
 	case *instructions.CopyCommand:
 		l := opt.buildContext
 		if len(cmd.sources) != 0 {
-			l = cmd.sources[0].state
+			src := cmd.sources[0]
+			if !src.noinit {
+				return errors.Errorf("cannot copy from stage %q, it needs to be defined before current stage %q", c.From, d.stageName)
+			}
+			l = src.state
 		}
 		err = dispatchCopy(d, copyConfig{
 			params:       c.SourcesAndDest,
diff --git a/frontend/dockerfile/dockerfile2llb/convert_runmount.go b/frontend/dockerfile/dockerfile2llb/convert_runmount.go
index 1f4597eb23af..ce59770384d3 100644
--- a/frontend/dockerfile/dockerfile2llb/convert_runmount.go
+++ b/frontend/dockerfile/dockerfile2llb/convert_runmount.go
@@ -70,7 +70,11 @@ func dispatchRunMounts(d *dispatchState, c *instructions.RunCommand, sources []*
 		}
 		st := opt.buildContext
 		if mount.From != "" {
-			st = sources[i].state
+			src := sources[i]
+			st = src.state
+			if !src.noinit {
+				return nil, errors.Errorf("cannot mount from stage %q to %q, stage needs to be defined before current command", mount.From, mount.Target)
+			}
 		}
 		var mountOpts []llb.MountOption
 		if mount.Type == instructions.MountTypeTmpfs {
diff --git a/frontend/dockerfile/dockerfile_test.go b/frontend/dockerfile/dockerfile_test.go
index 7b3a0a225d34..16199e766c98 100644
--- a/frontend/dockerfile/dockerfile_test.go
+++ b/frontend/dockerfile/dockerfile_test.go
@@ -156,6 +156,7 @@ var allTests = integration.TestFuncs(
 	testWorkdirUser,
 	testWorkdirExists,
 	testWorkdirCopyIgnoreRelative,
+	testOutOfOrderStage,
 	testCopyFollowAllSymlinks,
 	testDockerfileAddChownExpand,
 	testSourceDateEpochWithoutExporter,
@@ -953,6 +954,42 @@ RUN [ "$(stat -c "%U %G" /mydir)" == "user user" ]
 	require.NoError(t, err)
 }
 
+func testOutOfOrderStage(t *testing.T, sb integration.Sandbox) {
+	integration.SkipOnPlatform(t, "windows")
+	f := getFrontend(t, sb)
+
+	for _, src := range []string{"/", "/d2"} {
+		dockerfile := []byte(fmt.Sprintf(`
+FROM busybox AS target
+COPY --from=build %s /out
+
+FROM alpine AS build
+COPY /Dockerfile /d2
+
+FROM target
+`, src))
+
+		dir := integration.Tmpdir(
+			t,
+			fstest.CreateFile("Dockerfile", dockerfile, 0600),
+		)
+
+		c, err := client.New(sb.Context(), sb.Address())
+		require.NoError(t, err)
+		defer c.Close()
+
+		_, err = f.Solve(sb.Context(), c, client.SolveOpt{
+			LocalDirs: map[string]string{
+				dockerui.DefaultLocalNameDockerfile: dir,
+				dockerui.DefaultLocalNameContext:    dir,
+			},
+		}, nil)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "cannot copy from stage")
+		require.Contains(t, err.Error(), "needs to be defined before current stage")
+	}
+}
+
 func testWorkdirCopyIgnoreRelative(t *testing.T, sb integration.Sandbox) {
 	integration.SkipOnPlatform(t, "windows")
 	f := getFrontend(t, sb)