Create Header Middleware

middlewares/headers.go

@@ -0,0 +1,79 @@
+package middlewares
+
+//Middleware based on https://github.com/unrolled/secure
+
+import (
+	"github.com/containous/traefik/types"
+	"net/http"
+)
+
+// HeaderOptions is a struct for specifying configuration options for the headers middleware.
+type HeaderOptions struct {
+	// If Custom request headers are set, these will be added to the request
+	CustomRequestHeaders map[string]string
+	// If Custom response headers are set, these will be added to the ResponseWriter
+	CustomResponseHeaders map[string]string
+}
+
+// HeaderStruct is a middleware that helps setup a few basic security features. A single headerOptions struct can be
+// provided to configure which features should be enabled, and the ability to override a few of the default values.
+type HeaderStruct struct {
+	// Customize headers with a headerOptions struct.
+	opt HeaderOptions
+}
+
+// NewHeaderFromStruct constructs a new header instance from supplied frontend header struct.
+func NewHeaderFromStruct(headers types.Headers) *HeaderStruct {
+	o := HeaderOptions{
+		CustomRequestHeaders:  headers.CustomRequestHeaders,
+		CustomResponseHeaders: headers.CustomResponseHeaders,
+	}
+
+	return &HeaderStruct{
+		opt: o,
+	}
+}
+
+// NewHeader constructs a new header instance with supplied options.
+func NewHeader(options ...HeaderOptions) *HeaderStruct {
+	var o HeaderOptions
+	if len(options) == 0 {
+		o = HeaderOptions{}
+	} else {
+		o = options[0]
+	}
+
+	return &HeaderStruct{
+		opt: o,
+	}
+}
+
+// Handler implements the http.HandlerFunc for integration with the standard net/http lib.
+func (s *HeaderStruct) Handler(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Let headers process the request.
+		s.Process(w, r)
+		h.ServeHTTP(w, r)
+	})
+}
+
+func (s *HeaderStruct) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
+	s.Process(w, r)
+	// If there is a next, call it.
+	if next != nil {
+		next(w, r)
+	}
+}
+
+// Process runs the actual checks and returns an error if the middleware chain should stop.
+func (s *HeaderStruct) Process(w http.ResponseWriter, r *http.Request) {
+	// Loop through Custom request headers
+	for header, value := range s.opt.CustomRequestHeaders {
+		r.Header.Set(header, value)
+	}
+
+	// Loop through Custom response headers
+	for header, value := range s.opt.CustomResponseHeaders {
+		w.Header().Add(header, value)
+	}
+}

middlewares/headers_test.go

@@ -0,0 +1,59 @@
+package middlewares
+
+//Middleware tests based on https://github.com/unrolled/secure
+
+import (
+	"github.com/containous/traefik/testhelpers"
+	"github.com/stretchr/testify/assert"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+var myHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	w.Write([]byte("bar"))
+})
+
+func TestNoConfig(t *testing.T) {
+	s := NewHeader()
+
+	res := httptest.NewRecorder()
+	req := testhelpers.MustNewRequest(http.MethodGet, "http://example.com/foo", nil)
+
+	s.Handler(myHandler).ServeHTTP(res, req)
+
+	assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
+	assert.Equal(t, "bar", res.Body.String(), "Body not the expected")
+}
+
+func TestCustomResponseHeader(t *testing.T) {
+	s := NewHeader(HeaderOptions{
+		CustomResponseHeaders: map[string]string{
+			"X-Custom-Response-Header": "test_response",
+		},
+	})
+
+	res := httptest.NewRecorder()
+	req := testhelpers.MustNewRequest(http.MethodGet, "/foo", nil)
+
+	s.Handler(myHandler).ServeHTTP(res, req)
+
+	assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
+	assert.Equal(t, "test_response", res.Header().Get("X-Custom-Response-Header"), "Did not get expected header")
+}
+
+func TestCustomRequestHeader(t *testing.T) {
+	s := NewHeader(HeaderOptions{
+		CustomRequestHeaders: map[string]string{
+			"X-Custom-Request-Header": "test_request",
+		},
+	})
+
+	res := httptest.NewRecorder()
+	req := testhelpers.MustNewRequest(http.MethodGet, "/foo", nil)
+
+	s.Handler(myHandler).ServeHTTP(res, req)
+
+	assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
+	assert.Equal(t, "test_request", req.Header.Get("X-Custom-Request-Header"), "Did not get expected header")
+}

middlewares/secure.go

@@ -0,0 +1,31 @@
+package middlewares
+
+import (
+	"github.com/containous/traefik/types"
+	"github.com/unrolled/secure"
+)
+
+// NewSecure constructs a new Secure instance with supplied options.
+func NewSecure(headers types.Headers) *secure.Secure {
+	opt := secure.Options{
+		AllowedHosts:            headers.AllowedHosts,
+		HostsProxyHeaders:       headers.HostsProxyHeaders,
+		SSLRedirect:             headers.SSLRedirect,
+		SSLTemporaryRedirect:    headers.SSLTemporaryRedirect,
+		SSLHost:                 headers.SSLHost,
+		SSLProxyHeaders:         headers.SSLProxyHeaders,
+		STSSeconds:              headers.STSSeconds,
+		STSIncludeSubdomains:    headers.STSIncludeSubdomains,
+		STSPreload:              headers.STSPreload,
+		ForceSTSHeader:          headers.ForceSTSHeader,
+		FrameDeny:               headers.FrameDeny,
+		CustomFrameOptionsValue: headers.CustomFrameOptionsValue,
+		ContentTypeNosniff:      headers.ContentTypeNosniff,
+		BrowserXssFilter:        headers.BrowserXSSFilter,
+		ContentSecurityPolicy:   headers.ContentSecurityPolicy,
+		PublicKey:               headers.PublicKey,
+		ReferrerPolicy:          headers.ReferrerPolicy,
+		IsDevelopment:           headers.IsDevelopment,
+	}
+	return secure.New(opt)
+}