diff --git a/README.md b/README.md index b9a7d11..5c72f18 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,8 @@ by replacing a search regex by a replacement string. To configure the `Rewrite Body` plugin you should create a [middleware](https://docs.traefik.io/middlewares/overview/) in your dynamic configuration as explained [here](https://docs.traefik.io/middlewares/overview/). The following example creates -and uses the `rewritebody` middleware plugin to replace all foo occurences by bar in the HTTP response body. +and uses the `rewritebody` middleware plugin to replace all foo occurences by bar and the first occurrence of baz by qux in +the HTTP response body. If you want to apply some limits on the response body, you can chain this middleware plugin with the [Buffering middleware](https://docs.traefik.io/middlewares/buffering/) from Traefik. @@ -42,6 +43,12 @@ If you want to apply some limits on the response body, you can chain this middle regex = "foo" replacement = "bar" + # Rewrites only the first "baz" occurence by "qux" + [[http.middlewares.rewrite-foo.plugin.rewritebody.rewrites]] + regex = "baz" + replacement = "qux" + replaceOnce = true + [http.services] [http.services.my-service] [http.services.my-service.loadBalancer] diff --git a/rewritebody.go b/rewritebody.go index 966ccfb..1308ed5 100644 --- a/rewritebody.go +++ b/rewritebody.go @@ -16,6 +16,7 @@ import ( type Rewrite struct { Regex string `json:"regex,omitempty"` Replacement string `json:"replacement,omitempty"` + ReplaceOnce bool `json:"replaceOnce,omitempty"` } // Config holds the plugin configuration. @@ -32,6 +33,7 @@ func CreateConfig() *Config { type rewrite struct { regex *regexp.Regexp replacement []byte + replaceOnce bool } type rewriteBody struct { @@ -54,6 +56,7 @@ func New(_ context.Context, next http.Handler, config *Config, name string) (htt rewrites[i] = rewrite{ regex: regex, replacement: []byte(rewriteConfig.Replacement), + replaceOnce: rewriteConfig.ReplaceOnce, } } @@ -86,7 +89,14 @@ func (r *rewriteBody) ServeHTTP(rw http.ResponseWriter, req *http.Request) { } for _, rwt := range r.rewrites { - bodyBytes = rwt.regex.ReplaceAll(bodyBytes, rwt.replacement) + if rwt.replaceOnce { + firstOccurence := rwt.regex.Find(bodyBytes) + if firstOccurence != nil { + bodyBytes = bytes.Replace(bodyBytes, firstOccurence, rwt.replacement, 1) + } + } else { + bodyBytes = rwt.regex.ReplaceAll(bodyBytes, rwt.replacement) + } } if _, err := rw.Write(bodyBytes); err != nil { diff --git a/rewritebody_test.go b/rewritebody_test.go index 9b7bad4..c6d9e50 100644 --- a/rewritebody_test.go +++ b/rewritebody_test.go @@ -171,3 +171,69 @@ func TestNew(t *testing.T) { }) } } + +func TestReplaceOnce(t *testing.T) { + tests := []struct { + desc string + rewrites []Rewrite + resBody string + expResBody string + }{ + { + desc: "should replace only once", + rewrites: []Rewrite{ + { + Regex: "foo", + Replacement: "bar", + ReplaceOnce: true, + }, + }, + resBody: "my foo is the best of foos", + expResBody: "my bar is the best of foos", + }, + { + desc: "should replace every occurence", + rewrites: []Rewrite{ + { + Regex: "foo", + Replacement: "bar", + }, + }, + resBody: "my foo is the best of foos", + expResBody: "my bar is the best of bars", + }, + } + for _, test := range tests { + t.Run(test.desc, func(t *testing.T) { + config := &Config{ + Rewrites: test.rewrites, + } + + next := func(rw http.ResponseWriter, req *http.Request) { + rw.Header().Set("Content-Length", strconv.Itoa(len(test.resBody))) + rw.WriteHeader(http.StatusOK) + + _, _ = fmt.Fprintf(rw, test.resBody) + } + + _, err := New(context.Background(), nil, config, "rewriteBody") + if err != nil { + t.Fatal(err) + } + + recorder := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/", nil) + + rewriteBody, err := New(context.Background(), http.HandlerFunc(next), config, "rewriteBody") + if err != nil { + t.Fatal(err) + } + + rewriteBody.ServeHTTP(recorder, req) + + if !bytes.Equal([]byte(test.expResBody), recorder.Body.Bytes()) { + t.Errorf("got body %q, want %q", recorder.Body.Bytes(), test.expResBody) + } + }) + } +}