forked from holygeek/traefik
-
Notifications
You must be signed in to change notification settings - Fork 0
/
error_pages.go
174 lines (148 loc) · 5.13 KB
/
error_pages.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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
package middlewares
import (
"bufio"
"bytes"
"net"
"net/http"
"strconv"
"strings"
"github.com/containous/traefik/log"
"github.com/containous/traefik/types"
"github.com/vulcand/oxy/forward"
"github.com/vulcand/oxy/utils"
)
// Compile time validation that the response recorder implements http interfaces correctly.
var _ Stateful = &errorPagesResponseRecorderWithCloseNotify{}
//ErrorPagesHandler is a middleware that provides the custom error pages
type ErrorPagesHandler struct {
HTTPCodeRanges types.HTTPCodeRanges
BackendURL string
errorPageForwarder *forward.Forwarder
}
//NewErrorPagesHandler initializes the utils.ErrorHandler for the custom error pages
func NewErrorPagesHandler(errorPage *types.ErrorPage, backendURL string) (*ErrorPagesHandler, error) {
fwd, err := forward.New()
if err != nil {
return nil, err
}
httpCodeRanges, err := types.NewHTTPCodeRanges(errorPage.Status)
if err != nil {
return nil, err
}
return &ErrorPagesHandler{
HTTPCodeRanges: httpCodeRanges,
BackendURL: backendURL + errorPage.Query,
errorPageForwarder: fwd},
nil
}
func (ep *ErrorPagesHandler) ServeHTTP(w http.ResponseWriter, req *http.Request, next http.HandlerFunc) {
recorder := newErrorPagesResponseRecorder(w)
next.ServeHTTP(recorder, req)
w.WriteHeader(recorder.GetCode())
//check the recorder code against the configured http status code ranges
for _, block := range ep.HTTPCodeRanges {
if recorder.GetCode() >= block[0] && recorder.GetCode() <= block[1] {
log.Errorf("Caught HTTP Status Code %d, returning error page", recorder.GetCode())
finalURL := strings.Replace(ep.BackendURL, "{status}", strconv.Itoa(recorder.GetCode()), -1)
if newReq, err := http.NewRequest(http.MethodGet, finalURL, nil); err != nil {
w.Write([]byte(http.StatusText(recorder.GetCode())))
} else {
ep.errorPageForwarder.ServeHTTP(w, newReq)
}
return
}
}
//did not catch a configured status code so proceed with the request
utils.CopyHeaders(w.Header(), recorder.Header())
w.Write(recorder.GetBody().Bytes())
}
type errorPagesResponseRecorder interface {
http.ResponseWriter
http.Flusher
GetCode() int
GetBody() *bytes.Buffer
IsStreamingResponseStarted() bool
}
// newErrorPagesResponseRecorder returns an initialized responseRecorder.
func newErrorPagesResponseRecorder(rw http.ResponseWriter) errorPagesResponseRecorder {
recorder := &errorPagesResponseRecorderWithoutCloseNotify{
HeaderMap: make(http.Header),
Body: new(bytes.Buffer),
Code: http.StatusOK,
responseWriter: rw,
}
if _, ok := rw.(http.CloseNotifier); ok {
return &errorPagesResponseRecorderWithCloseNotify{recorder}
}
return recorder
}
// errorPagesResponseRecorderWithoutCloseNotify is an implementation of http.ResponseWriter that
// records its mutations for later inspection.
type errorPagesResponseRecorderWithoutCloseNotify struct {
Code int // the HTTP response code from WriteHeader
HeaderMap http.Header // the HTTP response headers
Body *bytes.Buffer // if non-nil, the bytes.Buffer to append written data to
responseWriter http.ResponseWriter
err error
streamingResponseStarted bool
}
type errorPagesResponseRecorderWithCloseNotify struct {
*errorPagesResponseRecorderWithoutCloseNotify
}
// CloseNotify returns a channel that receives at most a
// single value (true) when the client connection has gone
// away.
func (rw *errorPagesResponseRecorderWithCloseNotify) CloseNotify() <-chan bool {
return rw.responseWriter.(http.CloseNotifier).CloseNotify()
}
// Header returns the response headers.
func (rw *errorPagesResponseRecorderWithoutCloseNotify) Header() http.Header {
m := rw.HeaderMap
if m == nil {
m = make(http.Header)
rw.HeaderMap = m
}
return m
}
func (rw *errorPagesResponseRecorderWithoutCloseNotify) GetCode() int {
return rw.Code
}
func (rw *errorPagesResponseRecorderWithoutCloseNotify) GetBody() *bytes.Buffer {
return rw.Body
}
func (rw *errorPagesResponseRecorderWithoutCloseNotify) IsStreamingResponseStarted() bool {
return rw.streamingResponseStarted
}
// Write always succeeds and writes to rw.Body, if not nil.
func (rw *errorPagesResponseRecorderWithoutCloseNotify) Write(buf []byte) (int, error) {
if rw.err != nil {
return 0, rw.err
}
return rw.Body.Write(buf)
}
// WriteHeader sets rw.Code.
func (rw *errorPagesResponseRecorderWithoutCloseNotify) WriteHeader(code int) {
rw.Code = code
}
// Hijack hijacks the connection
func (rw *errorPagesResponseRecorderWithoutCloseNotify) Hijack() (net.Conn, *bufio.ReadWriter, error) {
return rw.responseWriter.(http.Hijacker).Hijack()
}
// Flush sends any buffered data to the client.
func (rw *errorPagesResponseRecorderWithoutCloseNotify) Flush() {
if !rw.streamingResponseStarted {
utils.CopyHeaders(rw.responseWriter.Header(), rw.Header())
rw.responseWriter.WriteHeader(rw.Code)
rw.streamingResponseStarted = true
}
_, err := rw.responseWriter.Write(rw.Body.Bytes())
if err != nil {
log.Errorf("Error writing response in responseRecorder: %s", err)
rw.err = err
}
rw.Body.Reset()
flusher, ok := rw.responseWriter.(http.Flusher)
if ok {
flusher.Flush()
}
}