diff --git a/middleware/replacer.go b/middleware/replacer.go index 10f0e35af..8a3d202a2 100644 --- a/middleware/replacer.go +++ b/middleware/replacer.go @@ -98,28 +98,20 @@ func NewReplacer(r *http.Request, rr *responseRecorder, emptyValue string) Repla // the string with the replaced values. func (r replacer) Replace(s string) string { // Header replacements - these are case-insensitive, so we can't just use strings.Replace() - startPos := strings.Index(s, headerReplacer) - for startPos > -1 { - // carefully find end of placeholder - endOffset := strings.Index(s[startPos+1:], "}") - if endOffset == -1 { - startPos = strings.Index(s[startPos+len(headerReplacer):], headerReplacer) - continue + for strings.Contains(s, headerReplacer) { + idxStart := strings.Index(s, headerReplacer) + endOffset := idxStart + len(headerReplacer) + idxEnd := strings.Index(s[endOffset:], "}") + if idxEnd > -1 { + placeholder := strings.ToLower(s[idxStart : endOffset+idxEnd+1]) + replacement := r.replacements[placeholder] + if replacement == "" { + replacement = r.emptyValue + } + s = s[:idxStart] + replacement + s[endOffset+idxEnd+1:] + } else { + break } - endPos := startPos + len(headerReplacer) + endOffset - - // look for replacement, case-insensitive - placeholder := strings.ToLower(s[startPos:endPos]) - replacement := r.replacements[placeholder] - if replacement == "" { - replacement = r.emptyValue - } - - // do the replacement manually - s = s[:startPos] + replacement + s[endPos:] - - // move to next one - startPos = strings.Index(s[endOffset:], headerReplacer) } // Regular replacements - these are easier because they're case-sensitive diff --git a/middleware/replacer_test.go b/middleware/replacer_test.go index 46d7b8f90..d98bd2de1 100644 --- a/middleware/replacer_test.go +++ b/middleware/replacer_test.go @@ -45,7 +45,8 @@ func TestReplace(t *testing.T) { if err != nil { t.Fatal("Request Formation Failed\n") } - request.Header.Set("Custom", "fooBar") + request.Header.Set("Custom", "foobarbaz") + request.Header.Set("ShorterVal", "1") repl := NewReplacer(request, recordRequest, "-") if expected, actual := "This host is localhost.", repl.Replace("This host is {host}."); expected != actual { @@ -57,12 +58,12 @@ func TestReplace(t *testing.T) { if expected, actual := "The response status is 200.", repl.Replace("The response status is {status}."); expected != actual { t.Errorf("{status} replacement: expected '%s', got '%s'", expected, actual) } - if expected, actual := "The Custom header is fooBar.", repl.Replace("The Custom header is {>Custom}."); expected != actual { + if expected, actual := "The Custom header is foobarbaz.", repl.Replace("The Custom header is {>Custom}."); expected != actual { t.Errorf("{>Custom} replacement: expected '%s', got '%s'", expected, actual) } // Test header case-insensitivity - if expected, actual := "The cUsToM header is fooBar...", repl.Replace("The cUsToM header is {>cUsToM}..."); expected != actual { + if expected, actual := "The cUsToM header is foobarbaz...", repl.Replace("The cUsToM header is {>cUsToM}..."); expected != actual { t.Errorf("{>cUsToM} replacement: expected '%s', got '%s'", expected, actual) } @@ -80,6 +81,16 @@ func TestReplace(t *testing.T) { if expected, actual := "Bad {>Custom placeholder", repl.Replace("Bad {>Custom placeholder"); expected != actual { t.Errorf("bad header placeholder: expected '%s', got '%s'", expected, actual) } + + // Test bad header placeholder with valid one later + if expected, actual := "Bad -", repl.Replace("Bad {>Custom placeholder {>ShorterVal}"); expected != actual { + t.Errorf("bad header placeholders: expected '%s', got '%s'", expected, actual) + } + + // Test shorter header value with multiple placeholders + if expected, actual := "Short value 1 then foobarbaz.", repl.Replace("Short value {>ShorterVal} then {>Custom}."); expected != actual { + t.Errorf("short value: expected '%s', got '%s'", expected, actual) + } } func TestSet(t *testing.T) {