/
headers.go
72 lines (63 loc) · 1.96 KB
/
headers.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
// Package headers could be used to rewrite headers of requests and responses.
package headers
import (
"encoding/json"
"fmt"
"net/http"
"golang.org/x/net/context"
"github.com/ironsmile/nedomi/config"
"github.com/ironsmile/nedomi/types"
"github.com/ironsmile/nedomi/utils"
"github.com/ironsmile/nedomi/utils/httputils"
)
// Headers rewrites headers
type Headers struct {
next types.RequestHandler
request headersRewrite
response headersRewrite
}
type headersConfig struct {
Request config.HeadersRewrite `json:"request"`
Response config.HeadersRewrite `json:"response"`
}
// RequestHandle rewrites the headers of the given request
func (h *Headers) RequestHandle(ctx context.Context, w http.ResponseWriter, r *http.Request) {
if !h.request.isEmpty() {
h.request.rewrite(r.Header)
}
if !h.response.isEmpty() {
w = h.wrapResponseWriter(w)
}
h.next.RequestHandle(ctx, w, r)
}
// New creates and returns a ready to used ServerStatusHandler.
func New(cfg *config.Handler, l *types.Location, next types.RequestHandler) (*Headers, error) {
var hr headersConfig
if len(cfg.Settings) != 0 {
if err := json.Unmarshal(cfg.Settings, &hr); err != nil {
return nil, err
}
}
return NewHeaders(next, hr.Request, hr.Response)
}
// NewHeaders is a more convinient constructor
func NewHeaders(next types.RequestHandler, request, response config.HeadersRewrite) (*Headers, error) {
if next == nil {
return nil, fmt.Errorf("headers handler requires next handler")
}
return &Headers{
next: next,
request: headersRewrite(request),
response: headersRewrite(response),
}, nil
}
func (h *Headers) wrapResponseWriter(w http.ResponseWriter) http.ResponseWriter {
var newW = httputils.NewFlexibleResponseWriter(func(frw *httputils.FlexibleResponseWriter) {
httputils.CopyHeaders(frw.Header(), w.Header())
h.response.rewrite(w.Header())
frw.BodyWriter = utils.AddCloser(w)
w.WriteHeader(frw.Code)
})
httputils.CopyHeaders(w.Header(), newW.Header())
return newW
}