feat: implement secrets in multihttp

Author: d0ugal

internal/prober/interpolation/interpolation.go

@@ -14,7 +14,7 @@ import (
 var VariableRegex = regexp.MustCompile(`\$\{([a-zA-Z_][a-zA-Z0-9_-]*)\}`)
 
 // SecretRegex matches ${secrets.secret_name} patterns
-var SecretRegex = regexp.MustCompile(`\$\{secrets\.([^}]*)\}`)
+var SecretRegex = regexp.MustCompile(`\$\{secrets\.([a-zA-Z0-9_][a-zA-Z0-9_\.\-]*)\}`)
 
 // VariableProvider defines the interface for resolving variables
 type VariableProvider interface {
@@ -219,6 +219,100 @@ func ToJavaScript(value string) string {
 	return s.String()
 }
 
+// ToJavaScriptWithSecrets converts a string with both variable and secret interpolation to JavaScript code
+// This is used by multihttp to generate JavaScript that references both variables and secrets
+func ToJavaScriptWithSecrets(value string) string {
+	if len(value) == 0 {
+		return `''`
+	}
+
+	var s strings.Builder
+	buf := []byte(value)
+
+	// First handle secret variables
+	locs := SecretRegex.FindAllSubmatchIndex(buf, -1)
+	p := 0
+
+	for _, loc := range locs {
+		if len(loc) < 4 {
+			panic("unexpected result while building JavaScript")
+		}
+
+		if s.Len() > 0 {
+			s.WriteRune('+')
+		}
+
+		if pre := buf[p:loc[0]]; len(pre) > 0 {
+			s.WriteRune('\'')
+			escapeJavaScript(&s, pre)
+			s.WriteRune('\'')
+			s.WriteRune('+')
+		}
+
+		// Generate async secret lookup
+		s.WriteString(`await secrets.get('`)
+		s.Write(buf[loc[2]:loc[3]])
+		s.WriteString(`')`)
+
+		p = loc[1]
+	}
+
+	// Then handle regular variables in the remaining text
+	remainingText := buf[p:]
+	if len(remainingText) > 0 {
+		regularLocs := VariableRegex.FindAllSubmatchIndex(remainingText, -1)
+
+		if len(regularLocs) > 0 {
+			if s.Len() > 0 {
+				s.WriteRune('+')
+			}
+
+			p2 := 0
+			for _, loc := range regularLocs {
+				if len(loc) < 4 {
+					panic("unexpected result while building JavaScript")
+				}
+
+				if s.Len() > 0 {
+					s.WriteRune('+')
+				}
+
+				if pre := remainingText[p2:loc[0]]; len(pre) > 0 {
+					s.WriteRune('\'')
+					escapeJavaScript(&s, pre)
+					s.WriteRune('\'')
+					s.WriteRune('+')
+				}
+
+				s.WriteString(`vars['`)
+				s.Write(remainingText[loc[2]:loc[3]])
+				s.WriteString(`']`)
+
+				p2 = loc[1]
+			}
+
+			if len(remainingText[p2:]) > 0 {
+				if s.Len() > 0 {
+					s.WriteRune('+')
+				}
+				s.WriteRune('\'')
+				escapeJavaScript(&s, remainingText[p2:])
+				s.WriteRune('\'')
+			}
+		} else {
+			// No regular variables, just append the remaining text
+			if s.Len() > 0 {
+				s.WriteRune('+')
+			}
+			s.WriteRune('\'')
+			escapeJavaScript(&s, remainingText)
+			s.WriteRune('\'')
+		}
+	}
+
+	return s.String()
+}
+
 // escapeJavaScript escapes a byte slice for use in JavaScript strings
 func escapeJavaScript(s *strings.Builder, buf []byte) {
 	for _, b := range buf {
@@ -235,10 +329,18 @@ func escapeJavaScript(s *strings.Builder, buf []byte) {
 			s.WriteString(`\r`)
 		case '\t':
 			s.WriteString(`\t`)
+		case '=':
+			s.WriteString(`\u003D`)
+		case '>':
+			s.WriteString(`\u003E`)
+		case '<':
+			s.WriteString(`\u003C`)
+		case '&':
+			s.WriteString(`\u0026`)
 		default:
-			if b < 32 || b > 126 {
-				// Escape non-printable characters
-				fmt.Fprintf(s, `\x%02x`, b)
+			if b < 32 {
+				// Escape control characters using Unicode escape
+				fmt.Fprintf(s, `\u%04X`, b)
 			} else {
 				s.WriteByte(b)
 			}

internal/prober/interpolation/interpolation_test.go

@@ -116,8 +116,8 @@ func TestResolver_Resolve(t *testing.T) {
 		"empty secret name": {
 			input:          "${secrets.}",
 			secretEnabled:  true,
-			expectedOutput: "",
-			expectError:    true,
+			expectedOutput: "${secrets.}",
+			expectError:    false,
 		},
 		"invalid secret name": {
 			input:          "${secrets.invalid-name}",

internal/prober/multihttp/script.go

@@ -9,6 +9,7 @@ import (
 	"strings"
 	"text/template"
 
+	"github.com/grafana/synthetic-monitoring-agent/internal/prober/interpolation"
 	sm "github.com/grafana/synthetic-monitoring-agent/pkg/pb/synthetic_monitoring"
 )
 
@@ -20,52 +21,7 @@ var templateFS embed.FS
 var userVariables = regexp.MustCompile(`\$\{([a-zA-Z_][a-zA-Z0-9_]*)\}`)
 
 func performVariableExpansion(in string) string {
-	if len(in) == 0 {
-		return `''`
-	}
-
-	var s strings.Builder
-	buf := []byte(in)
-	locs := userVariables.FindAllSubmatchIndex(buf, -1)
-
-	p := 0
-	for _, loc := range locs {
-		if len(loc) < 4 { // put the bounds checker at ease
-			panic("unexpected result while building URL")
-		}
-
-		if s.Len() > 0 {
-			s.WriteRune('+')
-		}
-
-		if pre := buf[p:loc[0]]; len(pre) > 0 {
-			s.WriteRune('\'')
-			template.JSEscape(&s, pre)
-			s.WriteRune('\'')
-			s.WriteRune('+')
-		}
-
-		s.WriteString(`vars['`)
-		// Because of the capture in the regular expression, the result
-		// has two indices that represent the matched substring, and
-		// two more indices that represent the capture group.
-		s.Write(buf[loc[2]:loc[3]])
-		s.WriteString(`']`)
-
-		p = loc[1]
-	}
-
-	if len(buf[p:]) > 0 {
-		if s.Len() > 0 {
-			s.WriteRune('+')
-		}
-
-		s.WriteRune('\'')
-		template.JSEscape(&s, buf[p:])
-		s.WriteRune('\'')
-	}
-
-	return s.String()
+	return interpolation.ToJavaScriptWithSecrets(in)
 }
 
 // Query params must be appended to a URL that has already been created.
@@ -113,13 +69,19 @@ func interpolateBodyVariables(bodyVarName string, body *sm.HttpRequestBody) []st
 	default:
 		var buf strings.Builder
 
-		matches := userVariables.FindAllString(string(body.Payload), -1)
+		// Find both regular and secret variables using submatch indices
+		regularMatches := userVariables.FindAllSubmatchIndex(body.Payload, -1)
+		secretMatches := interpolation.SecretRegex.FindAllSubmatchIndex(body.Payload, -1)
+
 		parsedMatches := make(map[string]struct{})
-		out := make([]string, 0, len(matches))
+		out := make([]string, 0, len(regularMatches)+len(secretMatches))
 
-		// For every instance of ${variable} in the body,
-		// this block returns {bodyVarName} = {bodyVarName}.replaceAll('${variable}', vars['variable'])
-		for _, m := range matches {
+		// Handle regular variables
+		for _, match := range regularMatches {
+			if len(match) < 4 {
+				continue
+			}
+			m := string(body.Payload[match[0]:match[1]])
 			if _, found := parsedMatches[m]; found {
 				continue
 			}
@@ -132,15 +94,38 @@ func interpolateBodyVariables(bodyVarName string, body *sm.HttpRequestBody) []st
 			buf.WriteString(m)
 			buf.WriteString("', vars['")
 			// writing the variable name from between ${ and }
-			for i := 2; i < len(m)-1; i++ {
-				buf.WriteByte(m[i])
-			}
+			buf.Write(body.Payload[match[2]:match[3]])
 			buf.WriteString("'])")
 			out = append(out, buf.String())
 
 			parsedMatches[m] = struct{}{}
 		}
 
+		// Handle secret variables
+		for _, match := range secretMatches {
+			if len(match) < 4 {
+				continue
+			}
+			m := string(body.Payload[match[0]:match[1]])
+			if _, found := parsedMatches[m]; found {
+				continue
+			}
+
+			buf.Reset()
+			buf.WriteString(bodyVarName)
+			buf.WriteString("=")
+			buf.WriteString(bodyVarName)
+			buf.WriteString(".replaceAll('")
+			buf.WriteString(m)
+			buf.WriteString("', await secrets.get('")
+			// writing the secret name from the capture group
+			buf.Write(body.Payload[match[2]:match[3]])
+			buf.WriteString("'))")
+			out = append(out, buf.String())
+
+			parsedMatches[m] = struct{}{}
+		}
+
 		return out
 	}
 }
@@ -422,6 +407,35 @@ func buildVars(variable *sm.MultiHttpEntryVariable) string {
 	return b.String()
 }
 
+func hasSecretVariables(settings *sm.MultiHttpSettings) bool {
+	for _, entry := range settings.Entries {
+		// Check URL
+		if interpolation.SecretRegex.MatchString(entry.Request.Url) {
+			return true
+		}
+
+		// Check headers
+		for _, header := range entry.Request.Headers {
+			if interpolation.SecretRegex.MatchString(header.Value) {
+				return true
+			}
+		}
+
+		// Check query fields
+		for _, field := range entry.Request.QueryFields {
+			if interpolation.SecretRegex.MatchString(field.Name) || interpolation.SecretRegex.MatchString(field.Value) {
+				return true
+			}
+		}
+
+		// Check body
+		if entry.Request.Body != nil && interpolation.SecretRegex.MatchString(string(entry.Request.Body.Payload)) {
+			return true
+		}
+	}
+	return false
+}
+
 func settingsToScript(settings *sm.MultiHttpSettings) ([]byte, error) {
 	// Convert settings to script using a Go template
 	tmpl, err := template.
@@ -442,9 +456,18 @@ func settingsToScript(settings *sm.MultiHttpSettings) ([]byte, error) {
 
 	var buf bytes.Buffer
 
+	// Create template data with secret variable detection
+	templateData := struct {
+		*sm.MultiHttpSettings
+		HasSecretVariables bool
+	}{
+		MultiHttpSettings:  settings,
+		HasSecretVariables: hasSecretVariables(settings),
+	}
+
 	// TODO(mem): figure out if we need to transform the data in some way
 	// before executing the template
-	if err := tmpl.ExecuteTemplate(&buf, "script.tmpl", settings); err != nil {
+	if err := tmpl.ExecuteTemplate(&buf, "script.tmpl", templateData); err != nil {
 		return nil, fmt.Errorf("executing script template: %w", err)
 	}
 

internal/prober/multihttp/script.tmpl

@@ -7,6 +7,7 @@ import { test } from 'k6/execution';
 import encoding from 'k6/encoding';
 import jsonpath from 'https://jslib.k6.io/jsonpath/1.0.2/index.js';
 import { URL } from 'https://jslib.k6.io/url/1.0.0/index.js';
+{{ if .HasSecretVariables }}import secrets from 'k6/secrets';{{ end }}
 
 export const options = {
 	scenarios: {
@@ -65,7 +66,7 @@ function assertHeader(headers, name, matcher) {
 	return false;
 }
 
-export default function() {
+export default {{ if .HasSecretVariables }}async {{ end }}function() {
 	let response;
 	let body;
 	let url;