homelab/traefik
1
0
Fork 0
Code Activity

Fix custom headers replacement

Browse source
This commit is contained in:
Michael 2017-11-23 17:40:03 +01:00 • committed by Traefiker
parent 24368747ab
commit f6181ef3e2
4 changed files with 115 additions and 44 deletions
Download patch file Download diff file Expand all files Collapse all files

53
middlewares/headers.go
View file

@ -35,46 +35,35 @@ func NewHeaderFromStruct(headers types.Headers) *HeaderStruct {
}
}
// 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)
s.ModifyRequestHeaders(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) {
// ModifyRequestHeaders set or delete request headers
func (s *HeaderStruct) ModifyRequestHeaders(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)
if value == "" {
r.Header.Del(header)
} else {
r.Header.Set(header, value)
}
}
}
// ModifyResponseHeaders set or delete response headers
func (s *HeaderStruct) ModifyResponseHeaders(res *http.Response) error {
// Loop through Custom response headers
for header, value := range s.opt.CustomResponseHeaders {
if value == "" {
res.Header.Del(header)
} else {
res.Header.Set(header, value)
}
}
return nil
}

75
middlewares/headers_test.go
View file

@ -1,6 +1,6 @@
package middlewares
//Middleware tests based on https://github.com/unrolled/secure
// Middleware tests based on https://github.com/unrolled/secure
import (
"net/http"
@ -15,36 +15,66 @@ var myHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte("bar"))
})
// 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,
}
}
func TestNoConfig(t *testing.T) {
s := NewHeader()
header := newHeader()
res := httptest.NewRecorder()
req := testhelpers.MustNewRequest(http.MethodGet, "http://example.com/foo", nil)
s.Handler(myHandler).ServeHTTP(res, req)
header.ServeHTTP(res, req, myHandler)
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{
func TestModifyResponseHeaders(t *testing.T) {
header := newHeader(HeaderOptions{
CustomResponseHeaders: map[string]string{
"X-Custom-Response-Header": "test_response",
},
})
res := httptest.NewRecorder()
req := testhelpers.MustNewRequest(http.MethodGet, "/foo", nil)
res.HeaderMap.Add("X-Custom-Response-Header", "test_response")
s.Handler(myHandler).ServeHTTP(res, req)
header.ModifyResponseHeaders(res.Result())
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")
res = httptest.NewRecorder()
res.HeaderMap.Add("X-Custom-Response-Header", "")
header.ModifyResponseHeaders(res.Result())
assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
assert.Equal(t, "", res.Header().Get("X-Custom-Response-Header"), "Did not get expected header")
res = httptest.NewRecorder()
res.HeaderMap.Add("X-Custom-Response-Header", "test_override")
header.ModifyResponseHeaders(res.Result())
assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
assert.Equal(t, "test_override", res.Header().Get("X-Custom-Response-Header"), "Did not get expected header")
}
func TestCustomRequestHeader(t *testing.T) {
s := NewHeader(HeaderOptions{
header := newHeader(HeaderOptions{
CustomRequestHeaders: map[string]string{
"X-Custom-Request-Header": "test_request",
},
@ -53,8 +83,35 @@ func TestCustomRequestHeader(t *testing.T) {
res := httptest.NewRecorder()
req := testhelpers.MustNewRequest(http.MethodGet, "/foo", nil)
s.Handler(myHandler).ServeHTTP(res, req)
header.ServeHTTP(res, req, nil)
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")
}
func TestCustomRequestHeaderEmptyValue(t *testing.T) {
header := newHeader(HeaderOptions{
CustomRequestHeaders: map[string]string{
"X-Custom-Request-Header": "test_request",
},
})
res := httptest.NewRecorder()
req := testhelpers.MustNewRequest(http.MethodGet, "/foo", nil)
header.ServeHTTP(res, req, nil)
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")
header = newHeader(HeaderOptions{
CustomRequestHeaders: map[string]string{
"X-Custom-Request-Header": "",
},
})
header.ServeHTTP(res, req, nil)
assert.Equal(t, http.StatusOK, res.Code, "Status not OK")
assert.Equal(t, "", req.Header.Get("X-Custom-Request-Header"), "This header is not expected")
}