Support ResponseHeaderModifier filter

pkg/middlewares/gatewayapi/headermodifier/request_header_modifier.go

@@ -9,7 +9,7 @@ import (
 	"go.opentelemetry.io/otel/trace"
 )
 
-const typeName = "RequestHeaderModifier"
+const requestHeaderModifierTypeName = "RequestHeaderModifier"
 
 // requestHeaderModifier is a middleware used to modify the headers of an HTTP request.
 type requestHeaderModifier struct {
@@ -22,8 +22,8 @@ type requestHeaderModifier struct {
 }
 
 // NewRequestHeaderModifier creates a new request header modifier middleware.
-func NewRequestHeaderModifier(ctx context.Context, next http.Handler, config dynamic.RequestHeaderModifier, name string) http.Handler {
-	logger := middlewares.GetLogger(ctx, name, typeName)
+func NewRequestHeaderModifier(ctx context.Context, next http.Handler, config dynamic.HeaderModifier, name string) http.Handler {
+	logger := middlewares.GetLogger(ctx, name, requestHeaderModifierTypeName)
 	logger.Debug().Msg("Creating middleware")
 
 	return &requestHeaderModifier{
@@ -36,7 +36,7 @@ func NewRequestHeaderModifier(ctx context.Context, next http.Handler, config dyn
 }
 
 func (r *requestHeaderModifier) GetTracingInformation() (string, string, trace.SpanKind) {
-	return r.name, typeName, trace.SpanKindUnspecified
+	return r.name, requestHeaderModifierTypeName, trace.SpanKindUnspecified
 }
 
 func (r *requestHeaderModifier) ServeHTTP(rw http.ResponseWriter, req *http.Request) {

pkg/middlewares/gatewayapi/headermodifier/request_header_modifier_test.go

@@ -14,25 +14,25 @@ import (
 func TestRequestHeaderModifier(t *testing.T) {
 	testCases := []struct {
 		desc            string
-		config          dynamic.RequestHeaderModifier
+		config          dynamic.HeaderModifier
 		requestHeaders  http.Header
 		expectedHeaders http.Header
 	}{
 		{
 			desc:            "no config",
-			config:          dynamic.RequestHeaderModifier{},
+			config:          dynamic.HeaderModifier{},
 			expectedHeaders: map[string][]string{},
 		},
 		{
 			desc: "set header",
-			config: dynamic.RequestHeaderModifier{
+			config: dynamic.HeaderModifier{
 				Set: map[string]string{"Foo": "Bar"},
 			},
 			expectedHeaders: map[string][]string{"Foo": {"Bar"}},
 		},
 		{
 			desc: "set header with existing headers",
-			config: dynamic.RequestHeaderModifier{
+			config: dynamic.HeaderModifier{
 				Set: map[string]string{"Foo": "Bar"},
 			},
 			requestHeaders:  map[string][]string{"Foo": {"Baz"}, "Bar": {"Foo"}},
@@ -40,7 +40,7 @@ func TestRequestHeaderModifier(t *testing.T) {
 		},
 		{
 			desc: "set multiple headers with existing headers",
-			config: dynamic.RequestHeaderModifier{
+			config: dynamic.HeaderModifier{
 				Set: map[string]string{"Foo": "Bar", "Bar": "Foo"},
 			},
 			requestHeaders:  map[string][]string{"Foo": {"Baz"}, "Bar": {"Foobar"}},
@@ -48,14 +48,14 @@ func TestRequestHeaderModifier(t *testing.T) {
 		},
 		{
 			desc: "add header",
-			config: dynamic.RequestHeaderModifier{
+			config: dynamic.HeaderModifier{
 				Add: map[string]string{"Foo": "Bar"},
 			},
 			expectedHeaders: map[string][]string{"Foo": {"Bar"}},
 		},
 		{
 			desc: "add header with existing headers",
-			config: dynamic.RequestHeaderModifier{
+			config: dynamic.HeaderModifier{
 				Add: map[string]string{"Foo": "Bar"},
 			},
 			requestHeaders:  map[string][]string{"Foo": {"Baz"}, "Bar": {"Foo"}},
@@ -63,7 +63,7 @@ func TestRequestHeaderModifier(t *testing.T) {
 		},
 		{
 			desc: "add multiple headers with existing headers",
-			config: dynamic.RequestHeaderModifier{
+			config: dynamic.HeaderModifier{
 				Add: map[string]string{"Foo": "Bar", "Bar": "Foo"},
 			},
 			requestHeaders:  map[string][]string{"Foo": {"Baz"}, "Bar": {"Foobar"}},
@@ -71,14 +71,14 @@ func TestRequestHeaderModifier(t *testing.T) {
 		},
 		{
 			desc: "remove header",
-			config: dynamic.RequestHeaderModifier{
+			config: dynamic.HeaderModifier{
 				Remove: []string{"Foo"},
 			},
 			expectedHeaders: map[string][]string{},
 		},
 		{
 			desc: "remove header with existing headers",
-			config: dynamic.RequestHeaderModifier{
+			config: dynamic.HeaderModifier{
 				Remove: []string{"Foo"},
 			},
 			requestHeaders:  map[string][]string{"Foo": {"Baz"}, "Bar": {"Foo"}},
@@ -86,7 +86,7 @@ func TestRequestHeaderModifier(t *testing.T) {
 		},
 		{
 			desc: "remove multiple headers with existing headers",
-			config: dynamic.RequestHeaderModifier{
+			config: dynamic.HeaderModifier{
 				Remove: []string{"Foo", "Bar"},
 			},
 			requestHeaders:  map[string][]string{"Foo": {"Bar"}, "Bar": {"Foo"}, "Baz": {"Bar"}},
@@ -106,11 +106,11 @@ func TestRequestHeaderModifier(t *testing.T) {
 			handler := NewRequestHeaderModifier(context.Background(), next, test.config, "foo-request-header-modifier")
 
 			req := testhelpers.MustNewRequest(http.MethodGet, "http://localhost", nil)
-			if test.requestHeaders != nil {
-				req.Header = test.requestHeaders
+			for h, v := range test.requestHeaders {
+				req.Header[h] = v
 			}
-
 			resp := httptest.NewRecorder()
+
 			handler.ServeHTTP(resp, req)
 
 			assert.Equal(t, test.expectedHeaders, gotHeaders)

pkg/middlewares/gatewayapi/headermodifier/response_header_modifier.go

@@ -0,0 +1,60 @@
+package headermodifier
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/traefik/traefik/v3/pkg/config/dynamic"
+	"github.com/traefik/traefik/v3/pkg/middlewares"
+	"go.opentelemetry.io/otel/trace"
+)
+
+const responseHeaderModifierTypeName = "ResponseHeaderModifier"
+
+// requestHeaderModifier is a middleware used to modify the headers of an HTTP response.
+type responseHeaderModifier struct {
+	next http.Handler
+	name string
+
+	set    map[string]string
+	add    map[string]string
+	remove []string
+}
+
+// NewResponseHeaderModifier creates a new response header modifier middleware.
+func NewResponseHeaderModifier(ctx context.Context, next http.Handler, config dynamic.HeaderModifier, name string) http.Handler {
+	logger := middlewares.GetLogger(ctx, name, responseHeaderModifierTypeName)
+	logger.Debug().Msg("Creating middleware")
+
+	return &responseHeaderModifier{
+		next:   next,
+		name:   name,
+		set:    config.Set,
+		add:    config.Add,
+		remove: config.Remove,
+	}
+}
+
+func (r *responseHeaderModifier) GetTracingInformation() (string, string, trace.SpanKind) {
+	return r.name, responseHeaderModifierTypeName, trace.SpanKindUnspecified
+}
+
+func (r *responseHeaderModifier) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
+	r.next.ServeHTTP(middlewares.NewResponseModifier(rw, req, r.modifyResponseHeaders), req)
+}
+
+func (r *responseHeaderModifier) modifyResponseHeaders(res *http.Response) error {
+	for headerName, headerValue := range r.set {
+		res.Header.Set(headerName, headerValue)
+	}
+
+	for headerName, headerValue := range r.add {
+		res.Header.Add(headerName, headerValue)
+	}
+
+	for _, headerName := range r.remove {
+		res.Header.Del(headerName)
+	}
+
+	return nil
+}

pkg/middlewares/gatewayapi/headermodifier/response_header_modifier_test.go

@@ -0,0 +1,121 @@
+package headermodifier
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/traefik/traefik/v3/pkg/config/dynamic"
+	"github.com/traefik/traefik/v3/pkg/testhelpers"
+)
+
+func TestResponseHeaderModifier(t *testing.T) {
+	testCases := []struct {
+		desc            string
+		config          dynamic.HeaderModifier
+		responseHeaders http.Header
+		expectedHeaders http.Header
+	}{
+		{
+			desc:            "no config",
+			config:          dynamic.HeaderModifier{},
+			expectedHeaders: map[string][]string{},
+		},
+		{
+			desc: "set header",
+			config: dynamic.HeaderModifier{
+				Set: map[string]string{"Foo": "Bar"},
+			},
+			expectedHeaders: map[string][]string{"Foo": {"Bar"}},
+		},
+		{
+			desc: "set header with existing headers",
+			config: dynamic.HeaderModifier{
+				Set: map[string]string{"Foo": "Bar"},
+			},
+			responseHeaders: map[string][]string{"Foo": {"Baz"}, "Bar": {"Foo"}},
+			expectedHeaders: map[string][]string{"Foo": {"Bar"}, "Bar": {"Foo"}},
+		},
+		{
+			desc: "set multiple headers with existing headers",
+			config: dynamic.HeaderModifier{
+				Set: map[string]string{"Foo": "Bar", "Bar": "Foo"},
+			},
+			responseHeaders: map[string][]string{"Foo": {"Baz"}, "Bar": {"Foobar"}},
+			expectedHeaders: map[string][]string{"Foo": {"Bar"}, "Bar": {"Foo"}},
+		},
+		{
+			desc: "add header",
+			config: dynamic.HeaderModifier{
+				Add: map[string]string{"Foo": "Bar"},
+			},
+			expectedHeaders: map[string][]string{"Foo": {"Bar"}},
+		},
+		{
+			desc: "add header with existing headers",
+			config: dynamic.HeaderModifier{
+				Add: map[string]string{"Foo": "Bar"},
+			},
+			responseHeaders: map[string][]string{"Foo": {"Baz"}, "Bar": {"Foo"}},
+			expectedHeaders: map[string][]string{"Foo": {"Baz", "Bar"}, "Bar": {"Foo"}},
+		},
+		{
+			desc: "add multiple headers with existing headers",
+			config: dynamic.HeaderModifier{
+				Add: map[string]string{"Foo": "Bar", "Bar": "Foo"},
+			},
+			responseHeaders: map[string][]string{"Foo": {"Baz"}, "Bar": {"Foobar"}},
+			expectedHeaders: map[string][]string{"Foo": {"Baz", "Bar"}, "Bar": {"Foobar", "Foo"}},
+		},
+		{
+			desc: "remove header",
+			config: dynamic.HeaderModifier{
+				Remove: []string{"Foo"},
+			},
+			expectedHeaders: map[string][]string{},
+		},
+		{
+			desc: "remove header with existing headers",
+			config: dynamic.HeaderModifier{
+				Remove: []string{"Foo"},
+			},
+			responseHeaders: map[string][]string{"Foo": {"Baz"}, "Bar": {"Foo"}},
+			expectedHeaders: map[string][]string{"Bar": {"Foo"}},
+		},
+		{
+			desc: "remove multiple headers with existing headers",
+			config: dynamic.HeaderModifier{
+				Remove: []string{"Foo", "Bar"},
+			},
+			responseHeaders: map[string][]string{"Foo": {"Bar"}, "Bar": {"Foo"}, "Baz": {"Bar"}},
+			expectedHeaders: map[string][]string{"Baz": {"Bar"}},
+		},
+	}
+
+	for _, test := range testCases {
+		t.Run(test.desc, func(t *testing.T) {
+			t.Parallel()
+
+			var nextCallCount int
+			next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+				nextCallCount++
+				rw.WriteHeader(http.StatusOK)
+			})
+
+			handler := NewResponseHeaderModifier(context.Background(), next, test.config, "foo-response-header-modifier")
+
+			req := testhelpers.MustNewRequest(http.MethodGet, "http://localhost", nil)
+			resp := httptest.NewRecorder()
+			for k, v := range test.responseHeaders {
+				resp.Header()[k] = v
+			}
+
+			handler.ServeHTTP(resp, req)
+
+			assert.Equal(t, 1, nextCallCount)
+			assert.Equal(t, test.expectedHeaders, resp.Header())
+		})
+	}
+}