rewrite: Implement `uri query` operations (#6120)

* Implemented basic uri query operations

* Added support for query operations block

* Applied Replacer on all query keys and values

* Implemented rename query key opration

* Rewrite struct: Changed QueryOperations field to Query and comments cleanup

* Cleaned up comments, changed the order of operations and added more tests

* Changed order of fields in queryOps struct to match the operations order

3 files changed, 226 insertions, 1 deletions

diff --git a/caddytest/integration/caddyfile_test.go b/caddytest/integration/caddyfile_test.go
index 959783be7..5d1fa3f08 100644
--- a/caddytest/integration/caddyfile_test.go
+++ b/caddytest/integration/caddyfile_test.go
@@ -497,6 +497,97 @@ func TestUriReplace(t *testing.T) {
 	tester.AssertGetResponse("http://localhost:9080/endpoint?test={%20content%20}", 200, "test=%7B%20content%20%7D")
 }
 
+func TestUriOps(t *testing.T) {
+	tester := caddytest.NewTester(t)
+
+	tester.InitServer(`
+	{
+		admin localhost:2999
+		http_port     9080
+	}
+	:9080
+	uri query +foo bar
+	uri query -baz
+	uri query taz test
+	uri query key=value example
+	uri query changethis>changed
+	
+	respond "{query}"`, "caddyfile")
+
+	tester.AssertGetResponse("http://localhost:9080/endpoint?foo=bar0&baz=buz&taz=nottest&changethis=val", 200, "changed=val&foo=bar0&foo=bar&key%3Dvalue=example&taz=test")
+}
+
+func TestSetThenAddQueryParams(t *testing.T) {
+	tester := caddytest.NewTester(t)
+
+	tester.InitServer(`
+	{
+		admin localhost:2999
+		http_port     9080
+	}
+	:9080
+	uri query foo bar
+	uri query +foo baz
+	
+	respond "{query}"`, "caddyfile")
+
+	tester.AssertGetResponse("http://localhost:9080/endpoint", 200, "foo=bar&foo=baz")
+}
+
+func TestSetThenDeleteParams(t *testing.T) {
+	tester := caddytest.NewTester(t)
+
+	tester.InitServer(`
+	{
+		admin localhost:2999
+		http_port     9080
+	}
+	:9080
+	uri query bar foo{query.foo}
+	uri query -foo
+	
+	respond "{query}"`, "caddyfile")
+
+	tester.AssertGetResponse("http://localhost:9080/endpoint?foo=bar", 200, "bar=foobar")
+}
+
+func TestRenameAndOtherOps(t *testing.T) {
+	tester := caddytest.NewTester(t)
+
+	tester.InitServer(`
+	{
+		admin localhost:2999
+		http_port     9080
+	}
+	:9080
+	uri query foo>bar
+	uri query bar taz
+	uri query +bar baz
+	
+	respond "{query}"`, "caddyfile")
+
+	tester.AssertGetResponse("http://localhost:9080/endpoint?foo=bar", 200, "bar=taz&bar=baz")
+}
+
+func TestUriOpsBlock(t *testing.T) {
+	tester := caddytest.NewTester(t)
+
+	tester.InitServer(`
+	{
+		admin localhost:2999
+		http_port     9080
+	}
+	:9080
+	uri query {
+		+foo bar
+		-baz
+		taz test
+	} 
+	respond "{query}"`, "caddyfile")
+
+	tester.AssertGetResponse("http://localhost:9080/endpoint?foo=bar0&baz=buz&taz=nottest", 200, "foo=bar0&foo=bar&taz=test")
+}
+
 func TestHandleErrorSimpleCodes(t *testing.T) {
 	tester := caddytest.NewTester(t)
 	tester.InitServer(`{
diff --git a/modules/caddyhttp/rewrite/caddyfile.go b/modules/caddyhttp/rewrite/caddyfile.go
index 363dbfdb5..31f7e9b48 100644
--- a/modules/caddyhttp/rewrite/caddyfile.go
+++ b/modules/caddyhttp/rewrite/caddyfile.go
@@ -98,7 +98,7 @@ func parseCaddyfileURI(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, err
 	h.Next() // consume directive name
 
 	args := h.RemainingArgs()
-	if len(args) < 2 {
+	if len(args) < 1 {
 		return nil, h.ArgErr()
 	}
 
@@ -158,12 +158,70 @@ func parseCaddyfileURI(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, err
 			Replace: replace,
 		})
 
+	case "query":
+		if len(args) > 4 {
+			return nil, h.ArgErr()
+		}
+		rewr.Query = &queryOps{}
+		var hasArgs bool
+		if len(args) > 1 {
+			hasArgs = true
+			err := applyQueryOps(h, rewr.Query, args[1:])
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		for h.NextBlock(0) {
+			if hasArgs {
+				return nil, h.Err("Cannot specify uri query rewrites in both argument and block")
+			}
+			queryArgs := []string{h.Val()}
+			queryArgs = append(queryArgs, h.RemainingArgs()...)
+			err := applyQueryOps(h, rewr.Query, queryArgs)
+			if err != nil {
+				return nil, err
+			}
+		}
+
 	default:
 		return nil, h.Errf("unrecognized URI manipulation '%s'", args[0])
 	}
 	return rewr, nil
 }
 
+func applyQueryOps(h httpcaddyfile.Helper, qo *queryOps, args []string) error {
+	key := args[0]
+	switch {
+	case strings.HasPrefix(key, "-"):
+		if len(args) != 1 {
+			return h.ArgErr()
+		}
+		qo.Delete = append(qo.Delete, strings.TrimLeft(key, "-"))
+
+	case strings.HasPrefix(key, "+"):
+		if len(args) != 2 {
+			return h.ArgErr()
+		}
+		param := strings.TrimLeft(key, "+")
+		qo.Add = append(qo.Add, queryOpsArguments{Key: param, Val: args[1]})
+
+	case strings.Contains(key, ">"):
+		if len(args) != 1 {
+			return h.ArgErr()
+		}
+		renameValKey := strings.Split(key, ">")
+		qo.Rename = append(qo.Rename, queryOpsArguments{Key: renameValKey[0], Val: renameValKey[1]})
+
+	default:
+		if len(args) != 2 {
+			return h.ArgErr()
+		}
+		qo.Set = append(qo.Set, queryOpsArguments{Key: key, Val: args[1]})
+	}
+	return nil
+}
+
 // parseCaddyfileHandlePath parses the handle_path directive. Syntax:
 //
 //	handle_path [<matcher>] {
diff --git a/modules/caddyhttp/rewrite/rewrite.go b/modules/caddyhttp/rewrite/rewrite.go
index 77ef668bf..1859f9df2 100644
--- a/modules/caddyhttp/rewrite/rewrite.go
+++ b/modules/caddyhttp/rewrite/rewrite.go
@@ -89,6 +89,9 @@ type Rewrite struct {
 	// Performs regular expression replacements on the URI path.
 	PathRegexp []*regexReplacer `json:"path_regexp,omitempty"`
 
+	// Mutates the query string of the URI.
+	Query *queryOps `json:"query,omitempty"`
+
 	logger *zap.Logger
 }
 
@@ -269,6 +272,11 @@ func (rewr Rewrite) Rewrite(r *http.Request, repl *caddy.Replacer) bool {
 		rep.do(r, repl)
 	}
 
+	// apply query operations
+	if rewr.Query != nil {
+		rewr.Query.do(r, repl)
+	}
+
 	// update the encoded copy of the URI
 	r.RequestURI = r.URL.RequestURI()
 
@@ -470,5 +478,73 @@ func changePath(req *http.Request, newVal func(pathOrRawPath string) string) {
 	}
 }
 
+// queryOps describes the operations to perform on query keys: add, set, rename and delete.
+type queryOps struct {
+	// Renames a query key from Key to Val, without affecting the value.
+	Rename []queryOpsArguments `json:"rename,omitempty"`
+
+	// Sets query parameters; overwrites a query key with the given value.
+	Set []queryOpsArguments `json:"set,omitempty"`
+
+	// Adds query parameters; does not overwrite an existing query field,
+	// and only appends an additional value for that key if any already exist.
+	Add []queryOpsArguments `json:"add,omitempty"`
+
+	// Deletes a given query key by name.
+	Delete []string `json:"delete,omitempty"`
+}
+
+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, "")
+		if key == "" || val == "" {
+			continue
+		}
+		query[val] = query[key]
+		delete(query, key)
+	}
+
+	for _, setParam := range q.Set {
+		key := repl.ReplaceAll(setParam.Key, "")
+		if key == "" {
+			continue
+		}
+		val := repl.ReplaceAll(setParam.Val, "")
+		query[key] = []string{val}
+	}
+
+	for _, addParam := range q.Add {
+		key := repl.ReplaceAll(addParam.Key, "")
+		if key == "" {
+			continue
+		}
+		val := repl.ReplaceAll(addParam.Val, "")
+		query[key] = append(query[key], val)
+	}
+
+	for _, deleteParam := range q.Delete {
+		param := repl.ReplaceAll(deleteParam, "")
+		if param == "" {
+			continue
+		}
+		delete(query, param)
+	}
+
+	r.URL.RawQuery = query.Encode()
+}
+
+type queryOpsArguments struct {
+	// A key in the query string. Note that query string keys may appear multiple times.
+	Key string `json:"key,omitempty"`
+
+	// The value for the given operation; for add and set, this is
+	// simply the value of the query, and for rename this is the
+	// query key to rename to.
+	Val string `json:"val,omitempty"`
+}
+
 // Interface guard
 var _ caddyhttp.MiddlewareHandler = (*Rewrite)(nil)