Use responseModifier to override secure headers

server/server.go

@@ -43,6 +43,7 @@ import (
 	"github.com/eapache/channels"
 	"github.com/sirupsen/logrus"
 	thoas_stats "github.com/thoas/stats"
+	"github.com/unrolled/secure"
 	"github.com/urfave/negroni"
 	"github.com/vulcand/oxy/buffer"
 	"github.com/vulcand/oxy/connlimit"
@@ -936,11 +937,9 @@ func (s *Server) loadConfig(configurations types.Configurations, globalConfigura
 					}
 
 					headerMiddleware := middlewares.NewHeaderFromStruct(frontend.Headers)
-					var responseModifier func(res *http.Response) error
-					if headerMiddleware != nil {
-						responseModifier = headerMiddleware.ModifyResponseHeaders
-					}
+					secureMiddleware := middlewares.NewSecure(frontend.Headers)
 
+					var responseModifier = buildModifyResponse(secureMiddleware, headerMiddleware)
 					var fwd http.Handler
 
 					fwd, err = forward.New(
@@ -1136,10 +1135,9 @@ func (s *Server) loadConfig(configurations types.Configurations, globalConfigura
 						n.Use(s.tracingMiddleware.NewNegroniHandlerWrapper("Header", headerMiddleware, false))
 					}
 
-					secureMiddleware := middlewares.NewSecure(frontend.Headers)
 					if secureMiddleware != nil {
 						log.Debugf("Adding secure middleware for frontend %s", frontendName)
-						n.UseFunc(secureMiddleware.HandlerFuncWithNext)
+						n.UseFunc(secureMiddleware.HandlerFuncWithNextForRequestOnly)
 					}
 
 					if config.Backends[frontend.Backend].Buffering != nil {
@@ -1493,3 +1491,21 @@ func (s *Server) buildBufferingMiddleware(handler http.Handler, config *types.Bu
 		buffer.CondSetter(len(config.RetryExpression) > 0, buffer.Retry(config.RetryExpression)),
 	)
 }
+
+func buildModifyResponse(secure *secure.Secure, header *middlewares.HeaderStruct) func(res *http.Response) error {
+	return func(res *http.Response) error {
+		if secure != nil {
+			err := secure.ModifyResponseHeaders(res)
+			if err != nil {
+				return err
+			}
+		}
+		if header != nil {
+			err := header.ModifyResponseHeaders(res)
+			if err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+}

server/server_test.go

@@ -9,6 +9,8 @@ import (
 	"testing"
 	"time"
 
+	"context"
+
 	"github.com/containous/flaeg"
 	"github.com/containous/mux"
 	"github.com/containous/traefik/configuration"
@@ -22,6 +24,7 @@ import (
 	"github.com/davecgh/go-spew/spew"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/unrolled/secure"
 	"github.com/urfave/negroni"
 	"github.com/vulcand/oxy/roundrobin"
 )
@@ -1017,6 +1020,123 @@ func TestBuildRedirectHandler(t *testing.T) {
 	}
 }
 
+type mockContext struct {
+	headers http.Header
+}
+
+func (c mockContext) Deadline() (deadline time.Time, ok bool) {
+	return deadline, ok
+}
+
+func (c mockContext) Done() <-chan struct{} {
+	ch := make(chan struct{})
+	close(ch)
+	return ch
+}
+
+func (c mockContext) Err() error {
+	return context.DeadlineExceeded
+}
+
+func (c mockContext) Value(key interface{}) interface{} {
+	return c.headers
+}
+
+func TestNewServerWithResponseModifiers(t *testing.T) {
+	cases := []struct {
+		desc             string
+		headerMiddleware *middlewares.HeaderStruct
+		secureMiddleware *secure.Secure
+		ctx              context.Context
+		expected         map[string]string
+	}{
+		{
+			desc:             "header and secure nil",
+			headerMiddleware: nil,
+			secureMiddleware: nil,
+			ctx:              mockContext{},
+			expected: map[string]string{
+				"X-Default":       "powpow",
+				"Referrer-Policy": "same-origin",
+			},
+		},
+		{
+			desc: "header middleware not nil",
+			headerMiddleware: middlewares.NewHeaderFromStruct(&types.Headers{
+				CustomResponseHeaders: map[string]string{
+					"X-Default": "powpow",
+				},
+			}),
+			secureMiddleware: nil,
+			ctx:              mockContext{},
+			expected: map[string]string{
+				"X-Default":       "powpow",
+				"Referrer-Policy": "same-origin",
+			},
+		},
+		{
+			desc:             "secure middleware not nil",
+			headerMiddleware: nil,
+			secureMiddleware: middlewares.NewSecure(&types.Headers{
+				ReferrerPolicy: "no-referrer",
+			}),
+			ctx: mockContext{
+				headers: http.Header{"Referrer-Policy": []string{"no-referrer"}},
+			},
+			expected: map[string]string{
+				"X-Default":       "powpow",
+				"Referrer-Policy": "no-referrer",
+			},
+		},
+		{
+			desc: "header and secure middleware not nil",
+			headerMiddleware: middlewares.NewHeaderFromStruct(&types.Headers{
+				CustomResponseHeaders: map[string]string{
+					"Referrer-Policy": "powpow",
+				},
+			}),
+			secureMiddleware: middlewares.NewSecure(&types.Headers{
+				ReferrerPolicy: "no-referrer",
+			}),
+			ctx: mockContext{
+				headers: http.Header{"Referrer-Policy": []string{"no-referrer"}},
+			},
+			expected: map[string]string{
+				"X-Default":       "powpow",
+				"Referrer-Policy": "powpow",
+			},
+		},
+	}
+
+	for _, test := range cases {
+		test := test
+		t.Run(test.desc, func(t *testing.T) {
+			t.Parallel()
+
+			headers := make(http.Header)
+			headers.Add("X-Default", "powpow")
+			headers.Add("Referrer-Policy", "same-origin")
+
+			req := httptest.NewRequest(http.MethodGet, "http://127.0.0.1", nil)
+
+			res := &http.Response{
+				Request: req.WithContext(test.ctx),
+				Header:  headers,
+			}
+
+			responseModifier := buildModifyResponse(test.secureMiddleware, test.headerMiddleware)
+			err := responseModifier(res)
+
+			assert.NoError(t, err)
+			assert.Equal(t, len(test.expected), len(res.Header))
+
+			for k, v := range test.expected {
+				assert.Equal(t, v, res.Header.Get(k))
+			}
+		})
+	}
+}
+
 func buildDynamicConfig(dynamicConfigBuilders ...func(*types.Configuration)) *types.Configuration {
 	config := &types.Configuration{
 		Frontends: make(map[string]*types.Frontend),