rewrite: `uri query` replace operation (#6165)

* Implemented query replace oeration

* Modified replace operation to use regexes in caddyfile

* Added more tests to uri query operations

4 files changed, 263 insertions, 1 deletions

diff --git a/caddytest/integration/caddyfile_adapt/uri_query_operations.caddyfiletest b/caddytest/integration/caddyfile_adapt/uri_query_operations.caddyfiletest
new file mode 100644
index 000000000..a53462480
--- /dev/null
+++ b/caddytest/integration/caddyfile_adapt/uri_query_operations.caddyfiletest
@@ -0,0 +1,106 @@
+:9080
+uri query +foo bar
+uri query -baz
+uri query taz test
+uri query key=value example
+uri query changethis>changed
+uri query {
+	findme value replacement
+	+foo1 baz
+}
+
+respond "{query}"
+----------
+{
+	"apps": {
+		"http": {
+			"servers": {
+				"srv0": {
+					"listen": [
+						":9080"
+					],
+					"routes": [
+						{
+							"handle": [
+								{
+									"handler": "rewrite",
+									"query": {
+										"add": [
+											{
+												"key": "foo",
+												"val": "bar"
+											}
+										]
+									}
+								},
+								{
+									"handler": "rewrite",
+									"query": {
+										"delete": [
+											"baz"
+										]
+									}
+								},
+								{
+									"handler": "rewrite",
+									"query": {
+										"set": [
+											{
+												"key": "taz",
+												"val": "test"
+											}
+										]
+									}
+								},
+								{
+									"handler": "rewrite",
+									"query": {
+										"set": [
+											{
+												"key": "key=value",
+												"val": "example"
+											}
+										]
+									}
+								},
+								{
+									"handler": "rewrite",
+									"query": {
+										"rename": [
+											{
+												"key": "changethis",
+												"val": "changed"
+											}
+										]
+									}
+								},
+								{
+									"handler": "rewrite",
+									"query": {
+										"add": [
+											{
+												"key": "foo1",
+												"val": "baz"
+											}
+										],
+										"replace": [
+											{
+												"key": "findme",
+												"replace": "replacement",
+												"search_regexp": "value"
+											}
+										]
+									}
+								},
+								{
+									"body": "{http.request.uri.query}",
+									"handler": "static_response"
+								}
+							]
+						}
+					]
+				}
+			}
+		}
+	}
+}
\ No newline at end of file
diff --git a/caddytest/integration/caddyfile_test.go b/caddytest/integration/caddyfile_test.go
index 5d1fa3f08..628363a52 100644
--- a/caddytest/integration/caddyfile_test.go
+++ b/caddytest/integration/caddyfile_test.go
@@ -569,6 +569,93 @@ func TestRenameAndOtherOps(t *testing.T) {
 	tester.AssertGetResponse("http://localhost:9080/endpoint?foo=bar", 200, "bar=taz&bar=baz")
 }
 
+func TestReplaceOps(t *testing.T) {
+	tester := caddytest.NewTester(t)
+
+	tester.InitServer(`
+	{
+		admin localhost:2999
+		http_port     9080
+	}
+	:9080
+	uri query foo bar baz	
+	respond "{query}"`, "caddyfile")
+
+	tester.AssertGetResponse("http://localhost:9080/endpoint?foo=bar", 200, "foo=baz")
+}
+
+func TestReplaceWithReplacementPlaceholder(t *testing.T) {
+	tester := caddytest.NewTester(t)
+	tester.InitServer(`
+	{
+		admin localhost:2999
+		http_port     9080
+	}
+	:9080
+	uri query foo bar {query.placeholder}	
+	respond "{query}"`, "caddyfile")
+
+	tester.AssertGetResponse("http://localhost:9080/endpoint?placeholder=baz&foo=bar", 200, "foo=baz&placeholder=baz")
+
+}
+
+func TestReplaceWithKeyPlaceholder(t *testing.T) {
+	tester := caddytest.NewTester(t)
+	tester.InitServer(`
+	{
+		admin localhost:2999
+		http_port     9080
+	}
+	:9080
+	uri query {query.placeholder} bar baz	
+	respond "{query}"`, "caddyfile")
+
+	tester.AssertGetResponse("http://localhost:9080/endpoint?placeholder=foo&foo=bar", 200, "foo=baz&placeholder=foo")
+}
+
+func TestPartialReplacement(t *testing.T) {
+	tester := caddytest.NewTester(t)
+	tester.InitServer(`
+	{
+		admin localhost:2999
+		http_port     9080
+	}
+	:9080
+	uri query foo ar az	
+	respond "{query}"`, "caddyfile")
+
+	tester.AssertGetResponse("http://localhost:9080/endpoint?foo=bar", 200, "foo=baz")
+}
+
+func TestNonExistingSearch(t *testing.T) {
+	tester := caddytest.NewTester(t)
+	tester.InitServer(`
+	{
+		admin localhost:2999
+		http_port     9080
+	}
+	:9080
+	uri query foo var baz	
+	respond "{query}"`, "caddyfile")
+
+	tester.AssertGetResponse("http://localhost:9080/endpoint?foo=bar", 200, "foo=bar")
+}
+
+func TestReplaceAllOps(t *testing.T) {
+	tester := caddytest.NewTester(t)
+
+	tester.InitServer(`
+	{
+		admin localhost:2999
+		http_port     9080
+	}
+	:9080
+	uri query * bar baz	
+	respond "{query}"`, "caddyfile")
+
+	tester.AssertGetResponse("http://localhost:9080/endpoint?foo=bar&baz=bar", 200, "baz=baz&foo=baz")
+}
+
 func TestUriOpsBlock(t *testing.T) {
 	tester := caddytest.NewTester(t)
 
diff --git a/modules/caddyhttp/rewrite/caddyfile.go b/modules/caddyhttp/rewrite/caddyfile.go
index 31f7e9b48..0ce5c41d2 100644
--- a/modules/caddyhttp/rewrite/caddyfile.go
+++ b/modules/caddyhttp/rewrite/caddyfile.go
@@ -213,6 +213,9 @@ func applyQueryOps(h httpcaddyfile.Helper, qo *queryOps, args []string) error {
 		renameValKey := strings.Split(key, ">")
 		qo.Rename = append(qo.Rename, queryOpsArguments{Key: renameValKey[0], Val: renameValKey[1]})
 
+	case len(args) == 3:
+		qo.Replace = append(qo.Replace, &queryOpsReplacement{Key: key, SearchRegexp: args[1], Replace: args[2]})
+
 	default:
 		if len(args) != 2 {
 			return h.ArgErr()
diff --git a/modules/caddyhttp/rewrite/rewrite.go b/modules/caddyhttp/rewrite/rewrite.go
index 1859f9df2..3479f0649 100644
--- a/modules/caddyhttp/rewrite/rewrite.go
+++ b/modules/caddyhttp/rewrite/rewrite.go
@@ -118,6 +118,12 @@ func (rewr *Rewrite) Provision(ctx caddy.Context) error {
 		rep.re = re
 	}
 
+	for _, replacementOp := range rewr.Query.Replace {
+		err := replacementOp.Provision(ctx)
+		if err != nil {
+			return fmt.Errorf("compiling regular expression %s in query rewrite replace operation: %v", replacementOp.SearchRegexp, err)
+		}
+	}
 	return nil
 }
 
@@ -490,13 +496,27 @@ type queryOps struct {
 	// and only appends an additional value for that key if any already exist.
 	Add []queryOpsArguments `json:"add,omitempty"`
 
+	// Replaces query parameters.
+	Replace []*queryOpsReplacement `json:"replace,omitempty"`
+
 	// Deletes a given query key by name.
 	Delete []string `json:"delete,omitempty"`
 }
 
+// Provision compiles the query replace operation regex.
+func (replacement *queryOpsReplacement) Provision(_ caddy.Context) error {
+	if replacement.SearchRegexp != "" {
+		re, err := regexp.Compile(replacement.SearchRegexp)
+		if err != nil {
+			return fmt.Errorf("replacement for query field '%s': %v", replacement.Key, err)
+		}
+		replacement.re = re
+	}
+	return nil
+}
+
 func (q *queryOps) do(r *http.Request, repl *caddy.Replacer) {
 	query := r.URL.Query()
-
 	for _, renameParam := range q.Rename {
 		key := repl.ReplaceAll(renameParam.Key, "")
 		val := repl.ReplaceAll(renameParam.Val, "")
@@ -525,6 +545,36 @@ func (q *queryOps) do(r *http.Request, repl *caddy.Replacer) {
 		query[key] = append(query[key], val)
 	}
 
+	for _, replaceParam := range q.Replace {
+		key := repl.ReplaceAll(replaceParam.Key, "")
+		search := repl.ReplaceKnown(replaceParam.Search, "")
+		replace := repl.ReplaceKnown(replaceParam.Replace, "")
+
+		// replace all query keys...
+		if key == "*" {
+			for fieldName, vals := range query {
+				for i := range vals {
+					if replaceParam.re != nil {
+						query[fieldName][i] = replaceParam.re.ReplaceAllString(query[fieldName][i], replace)
+					} else {
+						query[fieldName][i] = strings.ReplaceAll(query[fieldName][i], search, replace)
+					}
+				}
+			}
+			continue
+		}
+
+		for fieldName, vals := range query {
+			for i := range vals {
+				if replaceParam.re != nil {
+					query[fieldName][i] = replaceParam.re.ReplaceAllString(query[fieldName][i], replace)
+				} else {
+					query[fieldName][i] = strings.ReplaceAll(query[fieldName][i], search, replace)
+				}
+			}
+		}
+	}
+
 	for _, deleteParam := range q.Delete {
 		param := repl.ReplaceAll(deleteParam, "")
 		if param == "" {
@@ -546,5 +596,21 @@ type queryOpsArguments struct {
 	Val string `json:"val,omitempty"`
 }
 
+type queryOpsReplacement struct {
+	// The key to replace in the query string.
+	Key string `json:"key,omitempty"`
+
+	// The substring to search for.
+	Search string `json:"search,omitempty"`
+
+	// The regular expression to search with.
+	SearchRegexp string `json:"search_regexp,omitempty"`
+
+	// The string with which to replace matches.
+	Replace string `json:"replace,omitempty"`
+
+	re *regexp.Regexp
+}
+
 // Interface guard
 var _ caddyhttp.MiddlewareHandler = (*Rewrite)(nil)