From 3bb514f923c4b29c658f77b1d163c57a8f419def Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Sun, 6 Feb 2022 23:23:06 -0800
Subject: [PATCH] shell: add function to detect what env were used

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 frontend/dockerfile/shell/lex.go      | 18 +++++++++++++++++-
 frontend/dockerfile/shell/lex_test.go | 25 ++++++++++++++++++++++++-
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/frontend/dockerfile/shell/lex.go b/frontend/dockerfile/shell/lex.go
index e534f52e..23ab81f2 100644
--- a/frontend/dockerfile/shell/lex.go
+++ b/frontend/dockerfile/shell/lex.go
@@ -56,12 +56,20 @@ func (s *Lex) ProcessWordWithMap(word string, env map[string]string) (string, er
 	return word, err
 }
 
+// ProcessWordWithMatches will use the 'env' list of environment variables,
+// replace any env var references in 'word' and return the env that were used.
+func (s *Lex) ProcessWordWithMatches(word string, env map[string]string) (string, map[string]struct{}, error) {
+	sw := s.init(word, env)
+	word, _, err := sw.process(word)
+	return word, sw.matches, err
+}
+
 func (s *Lex) ProcessWordsWithMap(word string, env map[string]string) ([]string, error) {
 	_, words, err := s.process(word, env)
 	return words, err
 }
 
-func (s *Lex) process(word string, env map[string]string) (string, []string, error) {
+func (s *Lex) init(word string, env map[string]string) *shellWord {
 	sw := &shellWord{
 		envs:              env,
 		escapeToken:       s.escapeToken,
@@ -69,8 +77,14 @@ func (s *Lex) process(word string, env map[string]string) (string, []string, err
 		skipProcessQuotes: s.SkipProcessQuotes,
 		rawQuotes:         s.RawQuotes,
 		rawEscapes:        s.RawEscapes,
+		matches:           make(map[string]struct{}),
 	}
 	sw.scanner.Init(strings.NewReader(word))
+	return sw
+}
+
+func (s *Lex) process(word string, env map[string]string) (string, []string, error) {
+	sw := s.init(word, env)
 	return sw.process(word)
 }
 
@@ -82,6 +96,7 @@ type shellWord struct {
 	rawEscapes        bool
 	skipUnsetEnv      bool
 	skipProcessQuotes bool
+	matches           map[string]struct{}
 }
 
 func (sw *shellWord) process(source string) (string, []string, error) {
@@ -456,6 +471,7 @@ func isSpecialParam(char rune) bool {
 func (sw *shellWord) getEnv(name string) (string, bool) {
 	for key, value := range sw.envs {
 		if EqualEnvKeys(name, key) {
+			sw.matches[name] = struct{}{}
 			return value, true
 		}
 	}
diff --git a/frontend/dockerfile/shell/lex_test.go b/frontend/dockerfile/shell/lex_test.go
index 0a225a4b..b1430715 100644
--- a/frontend/dockerfile/shell/lex_test.go
+++ b/frontend/dockerfile/shell/lex_test.go
@@ -185,7 +185,7 @@ func TestShellParser4Words(t *testing.T) {
 }
 
 func TestGetEnv(t *testing.T) {
-	sw := &shellWord{envs: nil}
+	sw := &shellWord{envs: nil, matches: make(map[string]struct{})}
 
 	getEnv := func(name string) string {
 		value, _ := sw.getEnv(name)
@@ -225,3 +225,26 @@ func TestGetEnv(t *testing.T) {
 		t.Fatal("8 - 'car' should map to 'bike'")
 	}
 }
+
+func TestProcessWithMatches(t *testing.T) {
+	shlex := NewLex('\\')
+
+	w, matches, err := shlex.ProcessWordWithMatches("foo ${BAR} ${UNUSED}", map[string]string{
+		"ANOTHER": "bar",
+		"BAR":     "baz",
+	})
+	require.NoError(t, err)
+	require.Equal(t, "foo baz ", w)
+
+	require.Equal(t, 1, len(matches))
+	_, ok := matches["BAR"]
+	require.True(t, ok)
+
+	w, matches, err = shlex.ProcessWordWithMatches("foo ${BAR:-abc} ${UNUSED}", map[string]string{
+		"ANOTHER": "bar",
+	})
+	require.NoError(t, err)
+	require.Equal(t, "foo abc ", w)
+
+	require.Equal(t, 0, len(matches))
+}