1 files changed, 75 insertions, 0 deletions
diff --git a/ext.go b/ext.go
new file mode 100644
index 0000000..f6a0af8
--- /dev/null
+++ b/ext.go
@@ -0,0 +1,75 @@
+package main
+
+import (
+	"io"
+	"net/http"
+)
+
+type ResponseWriter interface {
+	http.ResponseWriter
+	Status() int
+}
+
+type responseWriter struct {
+	http.ResponseWriter
+	pendingStatus int
+	status        int
+	size          int
+}
+
+func NewResponseWriter(rw http.ResponseWriter) ResponseWriter {
+	return &responseWriter{
+		ResponseWriter: rw,
+	}
+}
+
+func (rw *responseWriter) WriteHeader(s int) {
+	if rw.Written() {
+		return
+	}
+
+	rw.pendingStatus = s
+
+	if rw.Written() {
+		return
+	}
+
+	rw.status = s
+	rw.ResponseWriter.WriteHeader(s)
+}
+
+func (rw *responseWriter) Write(b []byte) (int, error) {
+	if !rw.Written() {
+		// The status will be StatusOK if WriteHeader has not been called yet
+		rw.WriteHeader(http.StatusOK)
+	}
+	size, err := rw.ResponseWriter.Write(b)
+	rw.size += size
+	return size, err
+}
+
+func (rw *responseWriter) ReadFrom(r io.Reader) (n int64, err error) {
+	if !rw.Written() {
+		// The status will be StatusOK if WriteHeader has not been called yet
+		rw.WriteHeader(http.StatusOK)
+	}
+	n, err = io.Copy(rw.ResponseWriter, r)
+	rw.size += int(n)
+	return
+}
+
+func (rw *responseWriter) Unwrap() http.ResponseWriter {
+	return rw.ResponseWriter
+}
+
+func (rw *responseWriter) Status() int {
+	if rw.Written() {
+		return rw.status
+	}
+
+	return rw.pendingStatus
+}
+
+func (rw *responseWriter) Written() bool {
+	return rw.status != 0
+}