From f33e5070180dcad0450b84ef5dc68e218e22de84 Mon Sep 17 00:00:00 2001
From: Priya Wadhwa <priyawadhwa@google.com>
Date: Fri, 16 Mar 2018 12:16:07 -0700
Subject: [PATCH] Replaced shlex with dockerfile/shell package

---
 .../dockerfiles/Dockerfile_test_env           |  4 ++++
 integration_tests/dockerfiles/test_env.yaml   |  8 +++++++
 pkg/commands/env.go                           | 22 ++++++++++++++-----
 pkg/dockerfile/dockerfile.go                  |  3 +++
 4 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/integration_tests/dockerfiles/Dockerfile_test_env b/integration_tests/dockerfiles/Dockerfile_test_env
index 50de509cb..ce37081cb 100644
--- a/integration_tests/dockerfiles/Dockerfile_test_env
+++ b/integration_tests/dockerfiles/Dockerfile_test_env
@@ -13,3 +13,7 @@ ENV test2='a"b"c'
 ENV test3=a\ b name2=b\ c
 ENV test4="a\"b"
 ENV test5='a\"b'
+ENV test6="a\'b"
+ENV atomic=one
+ENV atomic=two newatomic=$atomic
+ENV newenv=$doesntexist/newenv
diff --git a/integration_tests/dockerfiles/test_env.yaml b/integration_tests/dockerfiles/test_env.yaml
index 435bc3ad3..c23c7a516 100644
--- a/integration_tests/dockerfiles/test_env.yaml
+++ b/integration_tests/dockerfiles/test_env.yaml
@@ -31,3 +31,11 @@ metadataTest:
       value: a"b
     - key: test5
       value: a\"b
+    - key: test6
+      value: a\'b
+    - key: atomic
+      value: two
+    - key: newatomic
+      value: one
+    - key: newenv
+      value: /newenv
diff --git a/pkg/commands/env.go b/pkg/commands/env.go
index 2f92339de..400fbdc6e 100644
--- a/pkg/commands/env.go
+++ b/pkg/commands/env.go
@@ -17,9 +17,10 @@ limitations under the License.
 package commands
 
 import (
+	"github.com/GoogleCloudPlatform/k8s-container-builder/pkg/dockerfile"
 	"github.com/containers/image/manifest"
 	"github.com/docker/docker/builder/dockerfile/instructions"
-	"github.com/google/shlex"
+	"github.com/docker/docker/builder/dockerfile/shell"
 	"github.com/sirupsen/logrus"
 	"os"
 	"strings"
@@ -33,14 +34,23 @@ func (e *EnvCommand) ExecuteCommand(config *manifest.Schema2Config) error {
 	logrus.Info("cmd: ENV")
 	newEnvs := e.cmd.Env
 	for index, pair := range newEnvs {
-		newVal := os.Expand(pair.Value, os.Getenv)
-		tokens, _ := shlex.Split(newVal)
+		// The dockerfile/shell package handles processing env values
+		// It handles escape characters and supports expansion from the config.Env array
+		// Shlex handles some of the following use cases (these and more are tested in integration tests)
+		// ""a'b'c"" -> "a'b'c"
+		// "Rex\ The\ Dog \" -> "Rex The Dog"
+		// "a\"b" -> "a"b"
+		shlex := shell.NewLex(dockerfile.EscapeToken)
+		expandedValue, err := shlex.ProcessWord(pair.Value, config.Env)
+		if err != nil {
+			return err
+		}
 		newEnvs[index] = instructions.KeyValuePair{
 			Key:   pair.Key,
-			Value: strings.Join(tokens, " "),
+			Value: expandedValue,
 		}
-		logrus.Infof("Setting environment variable %s=%s", pair.Key, pair.Value)
-		if err := os.Setenv(pair.Key, pair.Value); err != nil {
+		logrus.Infof("Setting environment variable %s=%s", pair.Key, expandedValue)
+		if err := os.Setenv(pair.Key, expandedValue); err != nil {
 			return err
 		}
 	}
diff --git a/pkg/dockerfile/dockerfile.go b/pkg/dockerfile/dockerfile.go
index 948cffa50..c96a1724c 100644
--- a/pkg/dockerfile/dockerfile.go
+++ b/pkg/dockerfile/dockerfile.go
@@ -23,12 +23,15 @@ import (
 	"github.com/docker/docker/builder/dockerfile/parser"
 )
 
+var EscapeToken rune
+
 // Parse parses the contents of a Dockerfile and returns a list of commands
 func Parse(b []byte) ([]instructions.Stage, error) {
 	p, err := parser.Parse(bytes.NewReader(b))
 	if err != nil {
 		return nil, err
 	}
+	EscapeToken = p.EscapeToken
 	stages, _, err := instructions.Parse(p.AST)
 	if err != nil {
 		return nil, err