From 969a98fe142231e43a4c371b8148363b979db28b Mon Sep 17 00:00:00 2001
From: Eric Zhang <ericwork@google.com>
Date: Wed, 3 Apr 2024 14:55:14 -0400
Subject: [PATCH] fix:
 stepresult-intepolations-does-not-accept-multiple-matches

---
 pkg/entrypoint/entrypointer.go      | 37 ++++++++++----------
 pkg/entrypoint/entrypointer_test.go | 54 +++++++++++++++++++++++------
 pkg/internal/resultref/resultref.go |  2 +-
 3 files changed, 64 insertions(+), 29 deletions(-)

diff --git a/pkg/entrypoint/entrypointer.go b/pkg/entrypoint/entrypointer.go
index f327160ffd4..8ae4c1f2cfd 100644
--- a/pkg/entrypoint/entrypointer.go
+++ b/pkg/entrypoint/entrypointer.go
@@ -439,41 +439,42 @@ func replaceEnv(stepDir string) error {
 	for _, e := range os.Environ() {
 		pair := strings.SplitN(e, "=", 2)
 		matches := resultref.StepResultRegex.FindAllStringSubmatch(pair[1], -1)
+		v := pair[1]
 		for _, m := range matches {
 			replaceWith, _, err := findReplacement(stepDir, m[0])
 			if err != nil {
 				return err
 			}
-			os.Setenv(pair[0], strings.ReplaceAll(pair[1], m[0], replaceWith))
+			v = strings.ReplaceAll(v, m[0], replaceWith)
 		}
+		os.Setenv(pair[0], v)
 	}
 	return nil
 }
 
 // replaceCommandAndArgs performs replacements for step results in e.Command
 func replaceCommandAndArgs(command []string, stepDir string) ([]string, error) {
-	newCommand := []string{}
+	var newCommand []string
 	for _, c := range command {
 		matches := resultref.StepResultRegex.FindAllStringSubmatch(c, -1)
-		if len(matches) > 0 {
-			for _, m := range matches {
-				replaceWithString, replaceWithArray, err := findReplacement(stepDir, m[0])
-				if err != nil {
-					return nil, err
-				}
-				// if replacing with an array
-				if len(replaceWithArray) > 1 {
-					// append with the array items
-					newCommand = append(newCommand, replaceWithArray...)
-				} else {
-					// append with replaced string
-					c = strings.ReplaceAll(c, m[0], replaceWithString)
-					newCommand = append(newCommand, c)
+		newC := []string{c}
+		for _, m := range matches {
+			replaceWithString, replaceWithArray, err := findReplacement(stepDir, m[0])
+			if err != nil {
+				return []string{}, fmt.Errorf("failed to find replacement for %s to replace %s", m[0], c)
+			}
+			// replaceWithString and replaceWithArray are mutually exclusive
+			if len(replaceWithArray) > 0 {
+				if c != m[0] {
+					// it has to be exact in "$(steps.<step-name>.results.<result-name>[*])" format, without anything else in the original string
+					return nil, errors.New("value must be in \"$(steps.<step-name>.results.<result-name>[*])\" format, when using array results")
 				}
+				newC = replaceWithArray
+			} else {
+				newC[0] = strings.ReplaceAll(newC[0], m[0], replaceWithString)
 			}
-		} else {
-			newCommand = append(newCommand, c)
 		}
+		newCommand = append(newCommand, newC...)
 	}
 	return newCommand, nil
 }
diff --git a/pkg/entrypoint/entrypointer_test.go b/pkg/entrypoint/entrypointer_test.go
index 2809716732e..40aa53158f3 100644
--- a/pkg/entrypoint/entrypointer_test.go
+++ b/pkg/entrypoint/entrypointer_test.go
@@ -786,15 +786,24 @@ func TestApplyStepResultSubstitutions_Env(t *testing.T) {
 		envValue:   "$(steps.foo.results.res.hello)",
 		want:       "World",
 		wantErr:    false,
-	}, {
-		name:       "bad-result-format",
-		stepName:   "foo",
-		resultName: "res",
-		result:     "{\"hello\":\"World\"}",
-		envValue:   "echo $(steps.foo.results.res.hello.bar)",
-		want:       "echo $(steps.foo.results.res.hello.bar)",
-		wantErr:    true,
-	}}
+	},
+		{
+			name:       "interpolation multiple matches",
+			stepName:   "foo",
+			resultName: "res",
+			result:     `{"first":"hello", "second":"world"}`,
+			envValue:   "$(steps.foo.results.res.first)-$(steps.foo.results.res.second)",
+			want:       "hello-world",
+			wantErr:    false,
+		}, {
+			name:       "bad-result-format",
+			stepName:   "foo",
+			resultName: "res",
+			result:     "{\"hello\":\"World\"}",
+			envValue:   "echo $(steps.foo.results.res.hello.bar)",
+			want:       "echo $(steps.foo.results.res.hello.bar)",
+			wantErr:    true,
+		}}
 	stepDir := createTmpDir(t, "env-steps")
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
@@ -875,7 +884,32 @@ func TestApplyStepResultSubstitutions_Command(t *testing.T) {
 		command:    []string{"echo $(steps.foo.results.res.hello.bar)"},
 		want:       []string{"echo $(steps.foo.results.res.hello.bar)"},
 		wantErr:    true,
-	}}
+	}, {
+		name:       "array param no index, with extra string",
+		stepName:   "foo",
+		resultName: "res",
+		result:     "[\"Hello\",\"World\"]",
+		command:    []string{"start", "$(steps.foo.results.res[*])bbb", "stop"},
+		want:       []string{"start", "$(steps.foo.results.res[*])bbb", "stop"},
+		wantErr:    true,
+	}, {
+		name:       "array param, multiple matches",
+		stepName:   "foo",
+		resultName: "res",
+		result:     "[\"Hello\",\"World\"]",
+		command:    []string{"$(steps.foo.results.res[0])-$(steps.foo.results.res[1])"},
+		want:       []string{"Hello-World"},
+		wantErr:    false,
+	}, {
+		name:       "object param, multiple matches",
+		stepName:   "foo",
+		resultName: "res",
+		result:     `{"first":"hello", "second":"world"}`,
+		command:    []string{"$(steps.foo.results.res.first)-$(steps.foo.results.res.second)"},
+		want:       []string{"hello-world"},
+		wantErr:    false,
+	},
+	}
 	stepDir := createTmpDir(t, "command-steps")
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
diff --git a/pkg/internal/resultref/resultref.go b/pkg/internal/resultref/resultref.go
index b42d0055a4d..f2c0152f768 100644
--- a/pkg/internal/resultref/resultref.go
+++ b/pkg/internal/resultref/resultref.go
@@ -42,7 +42,7 @@ const (
 
 	// arrayIndexing will match all `[int]` and `[*]` for parseExpression
 	arrayIndexing          = `\[([0-9])*\*?\]`
-	stepResultUsagePattern = `\$\(steps\..*\.results\..*\)`
+	stepResultUsagePattern = `\$\(steps\..*?\.results\..*?\)`
 )
 
 // arrayIndexingRegex is used to match `[int]` and `[*]`