From 5e5623e5a495bd502d8bc00f5ca5f78f814958d6 Mon Sep 17 00:00:00 2001
From: Priya Wadhwa <priyawadhwa@google.com>
Date: Mon, 20 Aug 2018 15:10:17 -0700
Subject: [PATCH] Fix bug in SaveStage function for multistage builds This
 change should fix the bug in #294, where kaniko wasn't recognizing that a
 stage would be used in a later build and so wasn't saving it as a tarball.

Each stage of the Dockerfile has a Name and a BaseName (FROM BaseName as
Name), but if a Name isn't specified then it's set to the same value as
BaseName. Our test cases weren't complete enough to catch this
distinction, which is why this bug occurred.

I added more test cases to the unit tests to make sure this fix works.
---
 pkg/dockerfile/dockerfile.go      |  6 ++++--
 pkg/dockerfile/dockerfile_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/pkg/dockerfile/dockerfile.go b/pkg/dockerfile/dockerfile.go
index d2b4de36a6..cf8ec99601 100644
--- a/pkg/dockerfile/dockerfile.go
+++ b/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) {
diff --git a/pkg/dockerfile/dockerfile_test.go b/pkg/dockerfile/dockerfile_test.go
index a91d53196a..d285f25a75 100644
--- a/pkg/dockerfile/dockerfile_test.go
+++ b/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) {