fix: rule.SetOnResponse not work

should handle response headers after next plugin and before write response statusCode

Author: isayme

htransformation.go

@@ -1,8 +1,10 @@
 package htransformation
 
 import (
+	"bufio"
 	"context"
 	"fmt"
+	"net"
 	"net/http"
 
 	"github.com/tomMoulard/htransformation/pkg/handler/deleter"
@@ -15,9 +17,10 @@ import (
 
 // HeadersTransformation holds the necessary components of a Traefik plugin.
 type HeadersTransformation struct {
-	name     string
-	next     http.Handler
-	handlers []types.Handler
+	name         string
+	next         http.Handler
+	reqHandlers  []types.Handler
+	respHandlers []types.Handler
 }
 
 // Config holds configuration to be passed to the plugin.
@@ -42,7 +45,8 @@ func New(_ context.Context, next http.Handler, config *Config, name string) (htt
 		types.Set:              set.New,
 	}
 
-	handlers := make([]types.Handler, 0, len(config.Rules))
+	reqHandlers := make([]types.Handler, 0, len(config.Rules))
+	respHandlers := make([]types.Handler, 0, len(config.Rules))
 
 	for _, rule := range config.Rules {
 		newHandler, ok := handlerBuilder[rule.Type]
@@ -59,22 +63,79 @@ func New(_ context.Context, next http.Handler, config *Config, name string) (htt
 			return nil, fmt.Errorf("%w: %s", err, rule.Name)
 		}
 
-		handlers = append(handlers, h)
+		if rule.SetOnResponse {
+			respHandlers = append(respHandlers, h)
+		} else {
+			reqHandlers = append(reqHandlers, h)
+		}
 	}
 
 	return &HeadersTransformation{
-		name:     name,
-		next:     next,
-		handlers: handlers,
+		name:         name,
+		next:         next,
+		reqHandlers:  reqHandlers,
+		respHandlers: respHandlers,
 	}, nil
 }
 
 // Iterate over every header to match the ones specified in the config and
 // return nothing if regexp failed.
 func (u *HeadersTransformation) ServeHTTP(responseWriter http.ResponseWriter, request *http.Request) {
-	for _, handler := range u.handlers {
+	for _, handler := range u.reqHandlers {
 		handler.Handle(responseWriter, request)
 	}
 
-	u.next.ServeHTTP(responseWriter, request)
+	wrappedResponseWriter := newWrappedResponseWriter(responseWriter, func(rw http.ResponseWriter) {
+		for _, handler := range u.respHandlers {
+			handler.Handle(rw, request)
+		}
+	})
+
+	u.next.ServeHTTP(wrappedResponseWriter, request)
+}
+
+type wrappedResponseWriter struct {
+	rw         http.ResponseWriter
+	handler    func(http.ResponseWriter)
+	headerSent bool
+}
+
+func newWrappedResponseWriter(rw http.ResponseWriter, handler func(http.ResponseWriter)) http.ResponseWriter {
+	return &wrappedResponseWriter{
+		rw:         rw,
+		handler:    handler,
+		headerSent: false,
+	}
+}
+
+func (wrw *wrappedResponseWriter) handleResponseHeader() {
+	if wrw.headerSent {
+		return
+	}
+
+	wrw.headerSent = true
+	wrw.handler(wrw.rw)
+}
+
+func (wrw *wrappedResponseWriter) Header() http.Header {
+	return wrw.rw.Header()
+}
+
+func (wrw *wrappedResponseWriter) Write(p []byte) (int, error) {
+	wrw.handleResponseHeader()
+	return wrw.rw.Write(p)
+}
+
+func (wrw *wrappedResponseWriter) WriteHeader(statusCode int) {
+	wrw.handleResponseHeader()
+	wrw.rw.WriteHeader(statusCode)
+}
+
+func (wrw *wrappedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	hijacker, ok := wrw.rw.(http.Hijacker)
+	if !ok {
+		return nil, nil, fmt.Errorf("%T is not an http.Hijacker", wrw.rw)
+	}
+
+	return hijacker.Hijack()
 }

htransformation_test.go

@@ -192,3 +192,71 @@ func TestHeaderRules(t *testing.T) {
 		})
 	}
 }
+
+func TestSetOnResponse(t *testing.T) {
+	testCases := []struct {
+		name             string
+		headerName       string
+		headerValue      string
+		rule             types.Rule
+		expectedNewValue string
+	}{
+		{
+			name:        "set rule",
+			headerName:  "Header-A",
+			headerValue: "valueA",
+			rule: types.Rule{
+				Name:          "set rule",
+				Header:        "Header-A",
+				Value:         "newValue",
+				Type:          types.Set,
+				SetOnResponse: true,
+			},
+			expectedNewValue: "newValue",
+		},
+		{
+			name:        "rewrite rule",
+			headerName:  "Header-A",
+			headerValue: "valueAA",
+			rule: types.Rule{
+				Name:          "rewrite rule",
+				Header:        "Header-A",
+				Value:         `value([\w\W]+)`,
+				ValueReplace:  "newValue-$1",
+				Type:          types.RewriteValueRule,
+				SetOnResponse: true,
+			},
+			expectedNewValue: "newValue-AA",
+		},
+	}
+
+	for _, test := range testCases {
+		t.Run(test.name, func(t *testing.T) {
+			cfg := plug.CreateConfig()
+			cfg.Rules = []types.Rule{test.rule}
+
+			ctx := context.Background()
+			next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+				rw.Header().Add(test.headerName, test.headerValue)
+				rw.WriteHeader(200)
+			})
+
+			handler, err := plug.New(ctx, next, cfg, "demo-plugin")
+			require.NoError(t, err)
+
+			recorder := httptest.NewRecorder()
+
+			req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://localhost", nil)
+			require.NoError(t, err)
+
+			handler.ServeHTTP(recorder, req)
+			resp := recorder.Result()
+			statusCode := resp.StatusCode
+			require.NoError(t, resp.Body.Close())
+
+			assert.Equal(t, http.StatusOK, statusCode)
+
+			assert.Equal(t, test.expectedNewValue, resp.Header.Get(test.rule.Header))
+		})
+	}
+}