Fix bug in SaveStage function for multistage builds

pkg/dockerfile/dockerfile.go

@@ -116,8 +116,10 @@ func SaveStage(index int, stages []instructions.Stage) bool {
 		if stageIndex <= index {
 			continue
 		}
-		if stage.Name == stages[index].BaseName {
-			return true
+		if stage.BaseName == stages[index].Name {
+			if stage.BaseName != "" {
+				return true
+			}
 		}
 		for _, cmd := range stage.Commands {
 			switch c := cmd.(type) {

pkg/dockerfile/dockerfile_test.go

@@ -113,6 +113,19 @@ func Test_SaveStage(t *testing.T) {
 
 		FROM scratch
 		COPY --from=second /hi2 /hi3
+
+		FROM ubuntu:16.04 AS base
+		ENV DEBIAN_FRONTEND noninteractive
+		ENV LC_ALL C.UTF-8
+
+		FROM base AS development
+		ENV PS1 " 🐳 \[\033[1;36m\]\W\[\033[0;35m\] # \[\033[0m\]"
+
+		FROM development AS test
+		ENV ORG_ENV UnitTest
+
+		FROM base AS production
+		COPY . /code
 		`,
 	}
 	if err := testutil.SetupFiles(tempDir, files); err != nil {
@@ -142,6 +155,21 @@ func Test_SaveStage(t *testing.T) {
 			index:    2,
 			expected: false,
 		},
+		{
+			name:     "reference current stage in next stage",
+			index:    4,
+			expected: true,
+		},
+		{
+			name:     "from prebuilt stage, and reference current stage in next stage",
+			index:    5,
+			expected: true,
+		},
+		{
+			name:     "final stage",
+			index:    6,
+			expected: false,
+		},
 	}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {