Move code to pkg

pkg/responsemodifiers/headers.go

@@ -0,0 +1,55 @@
+package responsemodifiers
+
+import (
+	"net/http"
+
+	"github.com/containous/traefik/pkg/config"
+	"github.com/unrolled/secure"
+)
+
+func buildHeaders(headers *config.Headers) func(*http.Response) error {
+	opt := secure.Options{
+		BrowserXssFilter:        headers.BrowserXSSFilter,
+		ContentTypeNosniff:      headers.ContentTypeNosniff,
+		ForceSTSHeader:          headers.ForceSTSHeader,
+		FrameDeny:               headers.FrameDeny,
+		IsDevelopment:           headers.IsDevelopment,
+		SSLRedirect:             headers.SSLRedirect,
+		SSLForceHost:            headers.SSLForceHost,
+		SSLTemporaryRedirect:    headers.SSLTemporaryRedirect,
+		STSIncludeSubdomains:    headers.STSIncludeSubdomains,
+		STSPreload:              headers.STSPreload,
+		ContentSecurityPolicy:   headers.ContentSecurityPolicy,
+		CustomBrowserXssValue:   headers.CustomBrowserXSSValue,
+		CustomFrameOptionsValue: headers.CustomFrameOptionsValue,
+		PublicKey:               headers.PublicKey,
+		ReferrerPolicy:          headers.ReferrerPolicy,
+		SSLHost:                 headers.SSLHost,
+		AllowedHosts:            headers.AllowedHosts,
+		HostsProxyHeaders:       headers.HostsProxyHeaders,
+		SSLProxyHeaders:         headers.SSLProxyHeaders,
+		STSSeconds:              headers.STSSeconds,
+	}
+
+	return func(resp *http.Response) error {
+		if headers.HasCustomHeadersDefined() {
+			// Loop through Custom response headers
+			for header, value := range headers.CustomResponseHeaders {
+				if value == "" {
+					resp.Header.Del(header)
+				} else {
+					resp.Header.Set(header, value)
+				}
+			}
+		}
+
+		if headers.HasSecureHeadersDefined() {
+			err := secure.New(opt).ModifyResponseHeaders(resp)
+			if err != nil {
+				return err
+			}
+		}
+
+		return nil
+	}
+}

pkg/responsemodifiers/log.go

@@ -0,0 +1,13 @@
+package responsemodifiers
+
+import (
+	"context"
+
+	"github.com/containous/traefik/pkg/log"
+	"github.com/sirupsen/logrus"
+)
+
+// getLogger creates a logger configured with the middleware fields.
+func getLogger(ctx context.Context, middleware string, middlewareType string) logrus.FieldLogger {
+	return log.FromContext(ctx).WithField(log.MiddlewareName, middleware).WithField(log.MiddlewareType, middlewareType)
+}

pkg/responsemodifiers/response_modifier.go

@@ -0,0 +1,51 @@
+package responsemodifiers
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/containous/traefik/pkg/config"
+)
+
+// NewBuilder creates a builder.
+func NewBuilder(configs map[string]*config.Middleware) *Builder {
+	return &Builder{configs: configs}
+}
+
+// Builder holds builder configuration.
+type Builder struct {
+	configs map[string]*config.Middleware
+}
+
+// Build Builds the response modifier.
+func (f *Builder) Build(ctx context.Context, names []string) func(*http.Response) error {
+	var modifiers []func(*http.Response) error
+
+	for _, middleName := range names {
+		if conf, ok := f.configs[middleName]; ok {
+			if conf.Headers != nil {
+				getLogger(ctx, middleName, "Headers").Debug("Creating Middleware (ResponseModifier)")
+
+				modifiers = append(modifiers, buildHeaders(conf.Headers))
+			} else if conf.Chain != nil {
+				getLogger(ctx, middleName, "Chain").Debug("Creating Middleware (ResponseModifier)")
+
+				modifiers = append(modifiers, f.Build(ctx, conf.Chain.Middlewares))
+			}
+		}
+	}
+
+	if len(modifiers) > 0 {
+		return func(resp *http.Response) error {
+			for i := len(modifiers); i > 0; i-- {
+				err := modifiers[i-1](resp)
+				if err != nil {
+					return err
+				}
+			}
+			return nil
+		}
+	}
+
+	return func(response *http.Response) error { return nil }
+}

pkg/responsemodifiers/response_modifier_test.go

@@ -0,0 +1,181 @@
+package responsemodifiers
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/containous/traefik/pkg/config"
+	"github.com/containous/traefik/pkg/middlewares/headers"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func stubResponse(_ map[string]*config.Middleware) *http.Response {
+	return &http.Response{Header: make(http.Header)}
+}
+
+func TestBuilderBuild(t *testing.T) {
+	testCases := []struct {
+		desc        string
+		middlewares []string
+		// buildResponse is needed because secure use a private context key
+		buildResponse  func(map[string]*config.Middleware) *http.Response
+		conf           map[string]*config.Middleware
+		assertResponse func(*testing.T, *http.Response)
+	}{
+		{
+			desc:           "no configuration",
+			middlewares:    []string{"foo", "bar"},
+			buildResponse:  stubResponse,
+			conf:           map[string]*config.Middleware{},
+			assertResponse: func(t *testing.T, resp *http.Response) {},
+		},
+		{
+			desc:          "one modifier",
+			middlewares:   []string{"foo", "bar"},
+			buildResponse: stubResponse,
+			conf: map[string]*config.Middleware{
+				"foo": {
+					Headers: &config.Headers{
+						CustomResponseHeaders: map[string]string{"X-Foo": "foo"},
+					},
+				},
+			},
+			assertResponse: func(t *testing.T, resp *http.Response) {
+				t.Helper()
+
+				assert.Equal(t, resp.Header.Get("X-Foo"), "foo")
+			},
+		},
+		{
+			desc:        "secure: one modifier",
+			middlewares: []string{"foo", "bar"},
+			buildResponse: func(middlewares map[string]*config.Middleware) *http.Response {
+				ctx := context.Background()
+
+				var request *http.Request
+				next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+					request = req
+				})
+
+				headerM := *middlewares["foo"].Headers
+				handler, err := headers.New(ctx, next, headerM, "secure")
+				require.NoError(t, err)
+
+				handler.ServeHTTP(httptest.NewRecorder(),
+					httptest.NewRequest(http.MethodGet, "http://foo.com", nil))
+
+				return &http.Response{Header: make(http.Header), Request: request}
+			},
+			conf: map[string]*config.Middleware{
+				"foo": {
+					Headers: &config.Headers{
+						ReferrerPolicy: "no-referrer",
+					},
+				},
+				"bar": {
+					Headers: &config.Headers{
+						CustomResponseHeaders: map[string]string{"X-Bar": "bar"},
+					},
+				},
+			},
+			assertResponse: func(t *testing.T, resp *http.Response) {
+				t.Helper()
+
+				assert.Equal(t, resp.Header.Get("Referrer-Policy"), "no-referrer")
+			},
+		},
+		{
+			desc:          "two modifiers",
+			middlewares:   []string{"foo", "bar"},
+			buildResponse: stubResponse,
+			conf: map[string]*config.Middleware{
+				"foo": {
+					Headers: &config.Headers{
+						CustomResponseHeaders: map[string]string{"X-Foo": "foo"},
+					},
+				},
+				"bar": {
+					Headers: &config.Headers{
+						CustomResponseHeaders: map[string]string{"X-Bar": "bar"},
+					},
+				},
+			},
+			assertResponse: func(t *testing.T, resp *http.Response) {
+				t.Helper()
+
+				assert.Equal(t, resp.Header.Get("X-Foo"), "foo")
+				assert.Equal(t, resp.Header.Get("X-Bar"), "bar")
+			},
+		},
+		{
+			desc:          "modifier order",
+			middlewares:   []string{"foo", "bar"},
+			buildResponse: stubResponse,
+			conf: map[string]*config.Middleware{
+				"foo": {
+					Headers: &config.Headers{
+						CustomResponseHeaders: map[string]string{"X-Foo": "foo"},
+					},
+				},
+				"bar": {
+					Headers: &config.Headers{
+						CustomResponseHeaders: map[string]string{"X-Foo": "bar"},
+					},
+				},
+			},
+			assertResponse: func(t *testing.T, resp *http.Response) {
+				t.Helper()
+
+				assert.Equal(t, resp.Header.Get("X-Foo"), "foo")
+			},
+		},
+		{
+			desc:          "chain",
+			middlewares:   []string{"chain"},
+			buildResponse: stubResponse,
+			conf: map[string]*config.Middleware{
+				"foo": {
+					Headers: &config.Headers{
+						CustomResponseHeaders: map[string]string{"X-Foo": "foo"},
+					},
+				},
+				"bar": {
+					Headers: &config.Headers{
+						CustomResponseHeaders: map[string]string{"X-Foo": "bar"},
+					},
+				},
+				"chain": {
+					Chain: &config.Chain{
+						Middlewares: []string{"foo", "bar"},
+					},
+				},
+			},
+			assertResponse: func(t *testing.T, resp *http.Response) {
+				t.Helper()
+
+				assert.Equal(t, resp.Header.Get("X-Foo"), "foo")
+			},
+		},
+	}
+
+	for _, test := range testCases {
+		test := test
+		t.Run(test.desc, func(t *testing.T) {
+			t.Parallel()
+
+			builder := NewBuilder(test.conf)
+
+			rm := builder.Build(context.Background(), test.middlewares)
+
+			resp := test.buildResponse(test.conf)
+
+			err := rm(resp)
+			require.NoError(t, err)
+
+			test.assertResponse(t, resp)
+		})
+	}
+}