feat: replace once implemented

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]

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 {

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)
+			}
+		})
+	}
+}