Merge pull request #505 from yui-knk/use_string_map

Add `ProcessWordWithMap` to `Lex`

frontend/dockerfile/shell/lex.go

@@ -28,7 +28,7 @@ func NewLex(escapeToken rune) *Lex {
 // ProcessWord will use the 'env' list of environment variables,
 // and replace any env var references in 'word'.
 func (s *Lex) ProcessWord(word string, env []string) (string, error) {
-	word, _, err := s.process(word, env)
+	word, _, err := s.process(word, buildEnvs(env))
 	return word, err
 }
 
@@ -40,11 +40,18 @@ func (s *Lex) ProcessWord(word string, env []string) (string, error) {
 // Note, each one is trimmed to remove leading and trailing spaces (unless
 // they are quoted", but ProcessWord retains spaces between words.
 func (s *Lex) ProcessWords(word string, env []string) ([]string, error) {
-	_, words, err := s.process(word, env)
+	_, words, err := s.process(word, buildEnvs(env))
 	return words, err
 }
 
-func (s *Lex) process(word string, env []string) (string, []string, error) {
+// ProcessWordWithMap will use the 'env' list of environment variables,
+// and replace any env var references in 'word'.
+func (s *Lex) ProcessWordWithMap(word string, env map[string]string) (string, error) {
+	word, _, err := s.process(word, env)
+	return word, err
+}
+
+func (s *Lex) process(word string, env map[string]string) (string, []string, error) {
 	sw := &shellWord{
 		envs:        env,
 		escapeToken: s.escapeToken,
@@ -55,7 +62,7 @@ func (s *Lex) process(word string, env []string) (string, []string, error) {
 
 type shellWord struct {
 	scanner     scanner.Scanner
-	envs        []string
+	envs        map[string]string
 	escapeToken rune
 }
 
@@ -353,21 +360,33 @@ func isSpecialParam(char rune) bool {
 }
 
 func (sw *shellWord) getEnv(name string) string {
-	for _, env := range sw.envs {
-		i := strings.Index(env, "=")
-		if i < 0 {
-			if EqualEnvKeys(name, env) {
-				// Should probably never get here, but just in case treat
-				// it like "var" and "var=" are the same
-				return ""
-			}
-			continue
+	for key, value := range sw.envs {
+		if EqualEnvKeys(name, key) {
+			return value
 		}
-		compareName := env[:i]
-		if !EqualEnvKeys(name, compareName) {
-			continue
-		}
-		return env[i+1:]
 	}
 	return ""
 }
+
+func buildEnvs(env []string) map[string]string {
+	envs := map[string]string{}
+
+	for _, e := range env {
+		i := strings.Index(e, "=")
+
+		if i < 0 {
+			envs[e] = ""
+		} else {
+			k := e[:i]
+			v := e[i+1:]
+
+			// If key already exists, keep previous value.
+			if _, ok := envs[k]; ok {
+				continue
+			}
+			envs[k] = v
+		}
+	}
+
+	return envs
+}

frontend/dockerfile/shell/lex_test.go

@@ -22,6 +22,7 @@ func TestShellParser4EnvVars(t *testing.T) {
 	shlex := NewLex('\\')
 	scanner := bufio.NewScanner(file)
 	envs := []string{"PWD=/home", "SHELL=bash", "KOREAN=한국어"}
+	envsMap := buildEnvs(envs)
 	for scanner.Scan() {
 		line := scanner.Text()
 		lineCount++
@@ -56,6 +57,14 @@ func TestShellParser4EnvVars(t *testing.T) {
 				assert.Check(t, err, "at line %d of %s", lineCount, fn)
 				assert.Check(t, is.Equal(newWord, expected), "at line %d of %s", lineCount, fn)
 			}
+
+			newWord, err = shlex.ProcessWordWithMap(source, envsMap)
+			if expected == "error" {
+				assert.Check(t, is.ErrorContains(err, ""), "input: %q, result: %q", source, newWord)
+			} else {
+				assert.Check(t, err, "at line %d of %s", lineCount, fn)
+				assert.Check(t, is.Equal(newWord, expected), "at line %d of %s", lineCount, fn)
+			}
 		}
 	}
 }
@@ -114,27 +123,27 @@ func TestShellParser4Words(t *testing.T) {
 func TestGetEnv(t *testing.T) {
 	sw := &shellWord{envs: nil}
 
-	sw.envs = []string{}
+	sw.envs = buildEnvs([]string{})
 	if sw.getEnv("foo") != "" {
 		t.Fatal("2 - 'foo' should map to ''")
 	}
 
-	sw.envs = []string{"foo"}
+	sw.envs = buildEnvs([]string{"foo"})
 	if sw.getEnv("foo") != "" {
 		t.Fatal("3 - 'foo' should map to ''")
 	}
 
-	sw.envs = []string{"foo="}
+	sw.envs = buildEnvs([]string{"foo="})
 	if sw.getEnv("foo") != "" {
 		t.Fatal("4 - 'foo' should map to ''")
 	}
 
-	sw.envs = []string{"foo=bar"}
+	sw.envs = buildEnvs([]string{"foo=bar"})
 	if sw.getEnv("foo") != "bar" {
 		t.Fatal("5 - 'foo' should map to 'bar'")
 	}
 
-	sw.envs = []string{"foo=bar", "car=hat"}
+	sw.envs = buildEnvs([]string{"foo=bar", "car=hat"})
 	if sw.getEnv("foo") != "bar" {
 		t.Fatal("6 - 'foo' should map to 'bar'")
 	}
@@ -143,7 +152,7 @@ func TestGetEnv(t *testing.T) {
 	}
 
 	// Make sure we grab the first 'car' in the list
-	sw.envs = []string{"foo=bar", "car=hat", "car=bike"}
+	sw.envs = buildEnvs([]string{"foo=bar", "car=hat", "car=bike"})
 	if sw.getEnv("car") != "hat" {
 		t.Fatal("8 - 'car' should map to 'hat'")
 	}