1package main
 2
 3import (
 4	"io"
 5	"net/http"
 6)
 7
 8type ResponseWriter interface {
 9	http.ResponseWriter
10	Status() int
11}
12
13type responseWriter struct {
14	http.ResponseWriter
15	pendingStatus int
16	status        int
17	size          int
18}
19
20func NewResponseWriter(rw http.ResponseWriter) ResponseWriter {
21	return &responseWriter{
22		ResponseWriter: rw,
23	}
24}
25
26func (rw *responseWriter) WriteHeader(s int) {
27	if rw.Written() {
28		return
29	}
30
31	rw.pendingStatus = s
32
33	if rw.Written() {
34		return
35	}
36
37	rw.status = s
38	rw.ResponseWriter.WriteHeader(s)
39}
40
41func (rw *responseWriter) Write(b []byte) (int, error) {
42	if !rw.Written() {
43		// The status will be StatusOK if WriteHeader has not been called yet
44		rw.WriteHeader(http.StatusOK)
45	}
46	size, err := rw.ResponseWriter.Write(b)
47	rw.size += size
48	return size, err
49}
50
51func (rw *responseWriter) ReadFrom(r io.Reader) (n int64, err error) {
52	if !rw.Written() {
53		// The status will be StatusOK if WriteHeader has not been called yet
54		rw.WriteHeader(http.StatusOK)
55	}
56	n, err = io.Copy(rw.ResponseWriter, r)
57	rw.size += int(n)
58	return
59}
60
61func (rw *responseWriter) Unwrap() http.ResponseWriter {
62	return rw.ResponseWriter
63}
64
65func (rw *responseWriter) Status() int {
66	if rw.Written() {
67		return rw.status
68	}
69
70	return rw.pendingStatus
71}
72
73func (rw *responseWriter) Written() bool {
74	return rw.status != 0
75}