forked from ant0ine/go-json-rest
-
Notifications
You must be signed in to change notification settings - Fork 0
/
recorder.go
76 lines (64 loc) · 2.08 KB
/
recorder.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
package rest
import (
"net/http"
)
// recorderMiddleware keeps a record of the HTTP status code of the response.
// The result is available to the wrapping handlers in request.Env["STATUS_CODE"] as an int.
type recorderMiddleware struct{}
func (mw *recorderMiddleware) MiddlewareFunc(h HandlerFunc) HandlerFunc {
return func(w ResponseWriter, r *Request) {
writer := &recorderResponseWriter{w, 0, false}
// call the handler
h(writer, r)
r.Env["STATUS_CODE"] = writer.statusCode
}
}
// Private responseWriter intantiated by the recorder middleware.
// It keeps a record of the HTTP status code of the response.
// It implements the following interfaces:
// ResponseWriter
// http.ResponseWriter
// http.Flusher
// http.CloseNotifier
type recorderResponseWriter struct {
ResponseWriter
statusCode int
wroteHeader bool
}
// Record the status code.
func (w *recorderResponseWriter) WriteHeader(code int) {
w.ResponseWriter.WriteHeader(code)
w.statusCode = code
w.wroteHeader = true
}
// Make sure the local WriteHeader is called, and call the parent WriteJson.
func (w *recorderResponseWriter) WriteJson(v interface{}) error {
if !w.wroteHeader {
w.WriteHeader(http.StatusOK)
}
return w.ResponseWriter.WriteJson(v)
}
// Make sure the local WriteHeader is called, and call the parent Flush.
// Provided in order to implement the http.Flusher interface.
func (w *recorderResponseWriter) Flush() {
if !w.wroteHeader {
w.WriteHeader(http.StatusOK)
}
flusher := w.ResponseWriter.(http.Flusher)
flusher.Flush()
}
// Call the parent CloseNotify.
// Provided in order to implement the http.CloseNotifier interface.
func (w *recorderResponseWriter) CloseNotify() <-chan bool {
notifier := w.ResponseWriter.(http.CloseNotifier)
return notifier.CloseNotify()
}
// Make sure the local WriteHeader is called, and call the parent Write.
// Provided in order to implement the http.ResponseWriter interface.
func (w *recorderResponseWriter) Write(b []byte) (int, error) {
if !w.wroteHeader {
w.WriteHeader(http.StatusOK)
}
writer := w.ResponseWriter.(http.ResponseWriter)
return writer.Write(b)
}