traefik/middlewares/headers.go
2017-11-09 00:48:03 +01:00

80 lines
2.3 KiB
Go

package middlewares
//Middleware based on https://github.com/unrolled/secure
import (
"net/http"
"github.com/containous/traefik/types"
)
// 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)
}
}