recorder.go 2.87 KB
Newer Older
1
package httpserver
Matthew Holt's avatar
Matthew Holt committed
2

3
import (
4 5 6
	"bufio"
	"errors"
	"net"
7 8 9
	"net/http"
	"time"
)
Matthew Holt's avatar
Matthew Holt committed
10

11
// ResponseRecorder is a type of http.ResponseWriter that captures
Matthew Holt's avatar
Matthew Holt committed
12 13 14 15 16
// the status code written to it and also the size of the body
// written in the response. A status code does not have
// to be written, however, in which case 200 must be assumed.
// It is best to have the constructor initialize this type
// with that default status code.
17 18 19 20 21 22
//
// Setting the Replacer field allows middlewares to type-assert
// the http.ResponseWriter to ResponseRecorder and set their own
// placeholder values for logging utilities to use.
//
// Beware when accessing the Replacer value; it may be nil!
23
type ResponseRecorder struct {
Matthew Holt's avatar
Matthew Holt committed
24
	http.ResponseWriter
25 26 27 28
	Replacer Replacer
	status   int
	size     int
	start    time.Time
Matthew Holt's avatar
Matthew Holt committed
29 30
}

31
// NewResponseRecorder makes and returns a new responseRecorder,
Matthew Holt's avatar
Matthew Holt committed
32 33 34 35 36
// which captures the HTTP Status code from the ResponseWriter
// and also the length of the response body written through it.
// Because a status is not set unless WriteHeader is called
// explicitly, this constructor initializes with a status code
// of 200 to cover the default case.
37 38
func NewResponseRecorder(w http.ResponseWriter) *ResponseRecorder {
	return &ResponseRecorder{
Matthew Holt's avatar
Matthew Holt committed
39 40
		ResponseWriter: w,
		status:         http.StatusOK,
41
		start:          time.Now(),
Matthew Holt's avatar
Matthew Holt committed
42 43 44 45 46
	}
}

// WriteHeader records the status code and calls the
// underlying ResponseWriter's WriteHeader method.
47
func (r *ResponseRecorder) WriteHeader(status int) {
Matthew Holt's avatar
Matthew Holt committed
48 49 50 51 52 53
	r.status = status
	r.ResponseWriter.WriteHeader(status)
}

// Write is a wrapper that records the size of the body
// that gets written.
54
func (r *ResponseRecorder) Write(buf []byte) (int, error) {
Matthew Holt's avatar
Matthew Holt committed
55 56 57 58 59 60
	n, err := r.ResponseWriter.Write(buf)
	if err == nil {
		r.size += n
	}
	return n, err
}
61

62 63 64 65 66 67 68 69 70 71
// Size is a Getter to size property
func (r *ResponseRecorder) Size() int {
	return r.size
}

// Status is a Getter to status property
func (r *ResponseRecorder) Status() int {
	return r.status
}

72 73
// Hijack implements http.Hijacker. It simply wraps the underlying
// ResponseWriter's Hijack method if there is one, or returns an error.
74
func (r *ResponseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
75 76 77
	if hj, ok := r.ResponseWriter.(http.Hijacker); ok {
		return hj.Hijack()
	}
78
	return nil, nil, errors.New("not a Hijacker")
79
}
80 81 82 83 84 85

// Flush implements http.Flusher. It simply wraps the underlying
// ResponseWriter's Flush method if there is one, or does nothing.
func (r *ResponseRecorder) Flush() {
	if f, ok := r.ResponseWriter.(http.Flusher); ok {
		f.Flush()
Matt Holt's avatar
Matt Holt committed
86 87
	} else {
		panic("not a Flusher") // should be recovered at the beginning of middleware stack
88 89
	}
}
90 91 92 93 94 95 96 97 98

// CloseNotify implements http.CloseNotifier.
// It just inherits the underlying ResponseWriter's CloseNotify method.
func (r *ResponseRecorder) CloseNotify() <-chan bool {
	if cn, ok := r.ResponseWriter.(http.CloseNotifier); ok {
		return cn.CloseNotify()
	}
	panic("not a CloseNotifier")
}