traefik/vendor/github.com/thoas/stats/recorder.go

96 lines
2.3 KiB
Go
Raw Normal View History

2017-02-07 21:33:23 +00:00
package stats
import (
"bufio"
"fmt"
"net"
"net/http"
)
type ResponseWriter interface {
http.ResponseWriter
http.Flusher
// Status returns the status code of the response or 0 if the response has not been written.
Status() int
// Written returns whether or not the ResponseWriter has been written.
Written() bool
// Size returns the size of the response body.
Size() int
// Before allows for a function to be called before the ResponseWriter has been written to. This is
// useful for setting headers or any other operations that must happen before a response has been written.
Before(func(ResponseWriter))
}
type beforeFunc func(ResponseWriter)
type recorderResponseWriter struct {
http.ResponseWriter
status int
size int
beforeFuncs []beforeFunc
2019-03-28 12:46:04 +00:00
written bool
2017-02-07 21:33:23 +00:00
}
func NewRecorderResponseWriter(w http.ResponseWriter, statusCode int) ResponseWriter {
2019-03-28 12:46:04 +00:00
return &recorderResponseWriter{ResponseWriter: w, status: statusCode}
2017-02-07 21:33:23 +00:00
}
func (r *recorderResponseWriter) WriteHeader(code int) {
2019-03-28 12:46:04 +00:00
r.written = true
2017-02-07 21:33:23 +00:00
r.ResponseWriter.WriteHeader(code)
r.status = code
}
func (r *recorderResponseWriter) Flush() {
flusher, ok := r.ResponseWriter.(http.Flusher)
if ok {
flusher.Flush()
}
}
func (r *recorderResponseWriter) Status() int {
return r.status
}
func (r *recorderResponseWriter) Write(b []byte) (int, error) {
if !r.Written() {
// The status will be StatusOK if WriteHeader has not been called yet
r.WriteHeader(http.StatusOK)
}
size, err := r.ResponseWriter.Write(b)
r.size += size
return size, err
}
// Proxy method to Status to add support for gocraft
func (r *recorderResponseWriter) StatusCode() int {
return r.Status()
}
func (r *recorderResponseWriter) Size() int {
return r.size
}
func (r *recorderResponseWriter) Written() bool {
return r.StatusCode() != 0
}
func (r *recorderResponseWriter) CloseNotify() <-chan bool {
return r.ResponseWriter.(http.CloseNotifier).CloseNotify()
}
func (r *recorderResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
2019-03-28 12:46:04 +00:00
if !r.written {
r.status = 0
}
2017-02-07 21:33:23 +00:00
hijacker, ok := r.ResponseWriter.(http.Hijacker)
if !ok {
return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
}
return hijacker.Hijack()
}
func (r *recorderResponseWriter) Before(before func(ResponseWriter)) {
r.beforeFuncs = append(r.beforeFuncs, before)
}