rewrite: query string enh.; substring replace; add tests (see #2891)

2 files changed, 341 insertions, 50 deletions

diff --git a/modules/caddyhttp/rewrite/rewrite.go b/modules/caddyhttp/rewrite/rewrite.go
index f610658..a142080 100644
--- a/modules/caddyhttp/rewrite/rewrite.go
+++ b/modules/caddyhttp/rewrite/rewrite.go
@@ -35,8 +35,9 @@ type Rewrite struct {
 	Method string `json:"method,omitempty"`
 	URI    string `json:"uri,omitempty"`
 
-	StripPathPrefix string `json:"strip_path_prefix,omitempty"`
-	StripPathSuffix string `json:"strip_path_suffix,omitempty"`
+	StripPathPrefix string     `json:"strip_path_prefix,omitempty"`
+	StripPathSuffix string     `json:"strip_path_suffix,omitempty"`
+	URISubstring    []replacer `json:"uri_substring,omitempty"`
 
 	HTTPRedirect caddyhttp.WeakString `json:"http_redirect,omitempty"`
 	Rehandle     bool                 `json:"rehandle,omitempty"`
@@ -61,49 +62,76 @@ func (rewr *Rewrite) Provision(ctx caddy.Context) error {
 // Validate ensures rewr's configuration is valid.
 func (rewr Rewrite) Validate() error {
 	if rewr.HTTPRedirect != "" && rewr.Rehandle {
-		return fmt.Errorf("cannot be configured to both write a redirect response and rehandle internally")
+		return fmt.Errorf("cannot be configured to both redirect externally and rehandle internally")
 	}
 	return nil
 }
 
 func (rewr Rewrite) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
 	repl := r.Context().Value(caddy.ReplacerCtxKey).(caddy.Replacer)
-	var changed bool
 
 	logger := rewr.logger.With(
 		zap.Object("request", caddyhttp.LoggableHTTPRequest{Request: r}),
 	)
 
-	// rewrite the method
+	changed := rewr.rewrite(r, repl, logger)
+
+	if changed {
+		logger.Debug("rewrote request",
+			zap.String("method", r.Method),
+			zap.String("uri", r.RequestURI),
+		)
+		if rewr.Rehandle {
+			return caddyhttp.ErrRehandle
+		}
+		if rewr.HTTPRedirect != "" {
+			statusCode, err := strconv.Atoi(repl.ReplaceAll(rewr.HTTPRedirect.String(), ""))
+			if err != nil {
+				return caddyhttp.Error(http.StatusInternalServerError, err)
+			}
+			w.Header().Set("Location", r.RequestURI)
+			w.WriteHeader(statusCode)
+			return nil
+		}
+	}
+
+	return next.ServeHTTP(w, r)
+}
+
+// rewrite performs the rewrites on r using repl, which
+// should have been obtained from r, but is passed in for
+// efficiency. It returns true if any changes were made to r.
+func (rewr Rewrite) rewrite(r *http.Request, repl caddy.Replacer, logger *zap.Logger) bool {
+	oldMethod := r.Method
+	oldURI := r.RequestURI
+
+	// method
 	if rewr.Method != "" {
-		method := r.Method
 		r.Method = strings.ToUpper(repl.ReplaceAll(rewr.Method, ""))
-		if r.Method != method {
-			changed = true
-		}
 	}
 
-	// rewrite the URI
+	// uri (which consists of path, query string, and maybe fragment?)
 	if rewr.URI != "" {
-		oldURI := r.RequestURI
 		newURI := repl.ReplaceAll(rewr.URI, "")
 
-		u, err := url.Parse(newURI)
+		newU, err := url.Parse(newURI)
 		if err != nil {
-			return caddyhttp.Error(http.StatusInternalServerError, err)
+			logger.Error("parsing new URI",
+				zap.String("raw_input", rewr.URI),
+				zap.String("input", newURI),
+				zap.Error(err),
+			)
 		}
 
-		r.RequestURI = newURI
-		r.URL.Path = u.Path
-		if u.RawQuery != "" {
-			r.URL.RawQuery = u.RawQuery
+		if newU.Path != "" {
+			r.URL.Path = newU.Path
 		}
-		if u.Fragment != "" {
-			r.URL.Fragment = u.Fragment
+		if newU.RawQuery != "" {
+			newU.RawQuery = strings.TrimPrefix(newU.RawQuery, "&")
+			r.URL.RawQuery = newU.RawQuery
 		}
-
-		if newURI != oldURI {
-			changed = true
+		if newU.Fragment != "" {
+			r.URL.Fragment = newU.Fragment
 		}
 	}
 
@@ -111,42 +139,63 @@ func (rewr Rewrite) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddy
 	if rewr.StripPathPrefix != "" {
 		prefix := repl.ReplaceAll(rewr.StripPathPrefix, "")
 		r.URL.Path = strings.TrimPrefix(r.URL.Path, prefix)
-		newURI := r.URL.String()
-		if newURI != r.RequestURI {
-			changed = true
-		}
-		r.RequestURI = newURI
 	}
 	if rewr.StripPathSuffix != "" {
 		suffix := repl.ReplaceAll(rewr.StripPathSuffix, "")
 		r.URL.Path = strings.TrimSuffix(r.URL.Path, suffix)
-		newURI := r.URL.String()
-		if newURI != r.RequestURI {
-			changed = true
-		}
-		r.RequestURI = newURI
 	}
 
-	if changed {
-		logger.Debug("rewrote request",
-			zap.String("method", r.Method),
-			zap.String("uri", r.RequestURI),
-		)
-		if rewr.Rehandle {
-			return caddyhttp.ErrRehandle
-		}
-		if rewr.HTTPRedirect != "" {
-			statusCode, err := strconv.Atoi(repl.ReplaceAll(rewr.HTTPRedirect.String(), ""))
-			if err != nil {
-				return caddyhttp.Error(http.StatusInternalServerError, err)
-			}
-			w.Header().Set("Location", r.RequestURI)
-			w.WriteHeader(statusCode)
-			return nil
-		}
+	// substring replacements in URI
+	for _, rep := range rewr.URISubstring {
+		rep.do(r, repl)
 	}
 
-	return next.ServeHTTP(w, r)
+	// update the encoded copy of the URI
+	r.RequestURI = r.URL.RequestURI()
+
+	// return true if anything changed
+	return r.Method != oldMethod || r.RequestURI != oldURI
+}
+
+// replacer describes a simple and fast substring replacement.
+type replacer struct {
+	// The substring to find. Supports placeholders.
+	Find string `json:"find,omitempty"`
+
+	// The substring to replace. Supports placeholders.
+	Replace string `json:"replace,omitempty"`
+
+	// Maximum number of replacements per string.
+	// Set to <= 0 for no limit (default).
+	Limit int `json:"limit,omitempty"`
+}
+
+// do performs the replacement on r and returns true if any changes were made.
+func (rep replacer) do(r *http.Request, repl caddy.Replacer) bool {
+	if rep.Find == "" || rep.Replace == "" {
+		return false
+	}
+
+	lim := rep.Limit
+	if lim == 0 {
+		lim = -1
+	}
+
+	find := repl.ReplaceAll(rep.Find, "")
+	replace := repl.ReplaceAll(rep.Replace, "")
+
+	oldPath := r.URL.Path
+	oldQuery := r.URL.RawQuery
+
+	r.URL.Path = strings.Replace(oldPath, find, replace, lim)
+	r.URL.RawQuery = strings.Replace(oldQuery, find, replace, lim)
+
+	// changed := r.URL.Path != oldPath && r.URL.RawQuery != oldQuery
+	// if changed {
+	// 	r.RequestURI = r.URL.RequestURI()
+	// }
+
+	return r.URL.Path != oldPath && r.URL.RawQuery != oldQuery
 }
 
 // Interface guard
diff --git a/modules/caddyhttp/rewrite/rewrite_test.go b/modules/caddyhttp/rewrite/rewrite_test.go
new file mode 100644
index 0000000..f496f0a
--- /dev/null
+++ b/modules/caddyhttp/rewrite/rewrite_test.go
@@ -0,0 +1,242 @@
+// Copyright 2015 Matthew Holt and The Caddy Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package rewrite
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/caddyserver/caddy/v2"
+)
+
+func TestRewrite(t *testing.T) {
+	repl := caddy.NewReplacer()
+
+	for i, tc := range []struct {
+		input, expect *http.Request
+		rule          Rewrite
+	}{
+		{
+			input:  newRequest(t, "GET", "/"),
+			expect: newRequest(t, "GET", "/"),
+		},
+		{
+			rule:   Rewrite{Method: "GET", URI: "/"},
+			input:  newRequest(t, "GET", "/"),
+			expect: newRequest(t, "GET", "/"),
+		},
+		{
+			rule:   Rewrite{Method: "POST"},
+			input:  newRequest(t, "GET", "/"),
+			expect: newRequest(t, "POST", "/"),
+		},
+
+		{
+			rule:   Rewrite{URI: "/foo"},
+			input:  newRequest(t, "GET", "/"),
+			expect: newRequest(t, "GET", "/foo"),
+		},
+		{
+			rule:   Rewrite{URI: "/foo"},
+			input:  newRequest(t, "GET", "/bar"),
+			expect: newRequest(t, "GET", "/foo"),
+		},
+		{
+			rule:   Rewrite{URI: "foo"},
+			input:  newRequest(t, "GET", "/"),
+			expect: newRequest(t, "GET", "foo"),
+		},
+		{
+			rule:   Rewrite{URI: "/foo{http.request.uri.path}"},
+			input:  newRequest(t, "GET", "/bar"),
+			expect: newRequest(t, "GET", "/foo/bar"),
+		},
+		{
+			rule:   Rewrite{URI: "/?c=d"},
+			input:  newRequest(t, "GET", "/"),
+			expect: newRequest(t, "GET", "/?c=d"),
+		},
+		{
+			rule:   Rewrite{URI: "/?c=d"},
+			input:  newRequest(t, "GET", "/?a=b"),
+			expect: newRequest(t, "GET", "/?c=d"),
+		},
+		{
+			rule:   Rewrite{URI: "?c=d"},
+			input:  newRequest(t, "GET", "/foo"),
+			expect: newRequest(t, "GET", "/foo?c=d"),
+		},
+		{
+			rule:   Rewrite{URI: "/?c=d"},
+			input:  newRequest(t, "GET", "/foo"),
+			expect: newRequest(t, "GET", "/?c=d"),
+		},
+		{
+			rule:   Rewrite{URI: "/?{http.request.uri.query}&c=d"},
+			input:  newRequest(t, "GET", "/"),
+			expect: newRequest(t, "GET", "/?c=d"),
+		},
+		{
+			rule:   Rewrite{URI: "/foo?{http.request.uri.query}&c=d"},
+			input:  newRequest(t, "GET", "/"),
+			expect: newRequest(t, "GET", "/foo?c=d"),
+		},
+		{
+			rule:   Rewrite{URI: "?{http.request.uri.query}&c=d"},
+			input:  newRequest(t, "GET", "/foo"),
+			expect: newRequest(t, "GET", "/foo?c=d"),
+		},
+		{
+			rule:   Rewrite{URI: "{http.request.uri.path}?{http.request.uri.query}&c=d"},
+			input:  newRequest(t, "GET", "/foo"),
+			expect: newRequest(t, "GET", "/foo?c=d"),
+		},
+		{
+			rule:   Rewrite{URI: "{http.request.uri.path}{http.request.uri.query_string}&c=d"},
+			input:  newRequest(t, "GET", "/foo"),
+			expect: newRequest(t, "GET", "/foo?c=d"),
+		},
+		{
+			rule:   Rewrite{URI: "/index.php?{http.request.uri.query}&c=d"},
+			input:  newRequest(t, "GET", "/foo"),
+			expect: newRequest(t, "GET", "/index.php?c=d"),
+		},
+		{
+			rule:   Rewrite{URI: "?a=b&c=d"},
+			input:  newRequest(t, "GET", "/foo"),
+			expect: newRequest(t, "GET", "/foo?a=b&c=d"),
+		},
+		{
+			rule:   Rewrite{URI: "/index.php?{http.request.uri.query}&c=d"},
+			input:  newRequest(t, "GET", "/?a=b"),
+			expect: newRequest(t, "GET", "/index.php?a=b&c=d"),
+		},
+		{
+			rule:   Rewrite{URI: "/index.php?c=d&{http.request.uri.query}"},
+			input:  newRequest(t, "GET", "/?a=b"),
+			expect: newRequest(t, "GET", "/index.php?c=d&a=b"),
+		},
+		{
+			rule:   Rewrite{URI: "/index.php?{http.request.uri.query}&p={http.request.uri.path}"},
+			input:  newRequest(t, "GET", "/foo/bar?a=b"),
+			expect: newRequest(t, "GET", "/index.php?a=b&p=/foo/bar"),
+		},
+
+		{
+			rule:   Rewrite{StripPathPrefix: "/prefix"},
+			input:  newRequest(t, "GET", "/foo/bar"),
+			expect: newRequest(t, "GET", "/foo/bar"),
+		},
+		{
+			rule:   Rewrite{StripPathPrefix: "/prefix"},
+			input:  newRequest(t, "GET", "/prefix/foo/bar"),
+			expect: newRequest(t, "GET", "/foo/bar"),
+		},
+		{
+			rule:   Rewrite{StripPathPrefix: "/prefix"},
+			input:  newRequest(t, "GET", "/foo/prefix/bar"),
+			expect: newRequest(t, "GET", "/foo/prefix/bar"),
+		},
+
+		{
+			rule:   Rewrite{StripPathSuffix: "/suffix"},
+			input:  newRequest(t, "GET", "/foo/bar"),
+			expect: newRequest(t, "GET", "/foo/bar"),
+		},
+		{
+			rule:   Rewrite{StripPathSuffix: "suffix"},
+			input:  newRequest(t, "GET", "/foo/bar/suffix"),
+			expect: newRequest(t, "GET", "/foo/bar/"),
+		},
+		{
+			rule:   Rewrite{StripPathSuffix: "/suffix"},
+			input:  newRequest(t, "GET", "/foo/suffix/bar"),
+			expect: newRequest(t, "GET", "/foo/suffix/bar"),
+		},
+
+		{
+			rule:   Rewrite{URISubstring: []replacer{{Find: "findme", Replace: "replaced"}}},
+			input:  newRequest(t, "GET", "/foo/bar"),
+			expect: newRequest(t, "GET", "/foo/bar"),
+		},
+		{
+			rule:   Rewrite{URISubstring: []replacer{{Find: "findme", Replace: "replaced"}}},
+			input:  newRequest(t, "GET", "/foo/findme/bar"),
+			expect: newRequest(t, "GET", "/foo/replaced/bar"),
+		},
+	} {
+		// copy the original input just enough so that we can
+		// compare it after the rewrite to see if it changed
+		originalInput := &http.Request{
+			Method:     tc.input.Method,
+			RequestURI: tc.input.RequestURI,
+			URL:        &*tc.input.URL,
+		}
+
+		// populate the replacer just enough for our tests
+		repl.Set("http.request.uri.path", tc.input.URL.Path)
+		repl.Set("http.request.uri.query", tc.input.URL.RawQuery)
+		repl.Set("http.request.uri.query_string", "?"+tc.input.URL.Query().Encode())
+
+		changed := tc.rule.rewrite(tc.input, repl, nil)
+
+		if expected, actual := !reqEqual(originalInput, tc.input), changed; expected != actual {
+			t.Errorf("Test %d: Expected changed=%t but was %t", i, expected, actual)
+		}
+		if expected, actual := tc.expect.Method, tc.input.Method; expected != actual {
+			t.Errorf("Test %d: Expected Method='%s' but got '%s'", i, expected, actual)
+		}
+		if expected, actual := tc.expect.RequestURI, tc.input.RequestURI; expected != actual {
+			t.Errorf("Test %d: Expected RequestURI='%s' but got '%s'", i, expected, actual)
+		}
+		if expected, actual := tc.expect.URL.String(), tc.input.URL.String(); expected != actual {
+			t.Errorf("Test %d: Expected URL='%s' but got '%s'", i, expected, actual)
+		}
+		if expected, actual := tc.expect.URL.RequestURI(), tc.input.URL.RequestURI(); expected != actual {
+			t.Errorf("Test %d: Expected URL.RequestURI()='%s' but got '%s'", i, expected, actual)
+		}
+	}
+}
+
+func newRequest(t *testing.T, method, uri string) *http.Request {
+	req, err := http.NewRequest(method, uri, nil)
+	if err != nil {
+		t.Fatalf("error creating request: %v", err)
+	}
+	req.RequestURI = req.URL.RequestURI() // simulate incoming request
+	return req
+}
+
+// reqEqual if r1 and r2 are equal enough for our purposes.
+func reqEqual(r1, r2 *http.Request) bool {
+	if r1.Method != r2.Method {
+		return false
+	}
+	if r1.RequestURI != r2.RequestURI {
+		return false
+	}
+	if (r1.URL == nil && r2.URL != nil) || (r1.URL != nil && r2.URL == nil) {
+		return false
+	}
+	if r1.URL == nil && r2.URL == nil {
+		return true
+	}
+	return r1.URL.Scheme == r2.URL.Scheme &&
+		r1.URL.Host == r2.URL.Host &&
+		r1.URL.Path == r2.URL.Path &&
+		r1.URL.RawPath == r2.URL.RawPath &&
+		r1.URL.RawQuery == r2.URL.RawQuery &&
+		r1.URL.Fragment == r2.URL.Fragment
+}