1 files changed, 131 insertions, 0 deletions
diff --git a/response_writer.go b/response_writer.go
new file mode 100644
index 0000000..65cae31
--- /dev/null
+++ b/response_writer.go
@@ -0,0 +1,131 @@
+// Copyright 2014 Manu Martinez-Almeida. All rights reserved.
+// Use of this source code is governed by a MIT style
+// license that can be found in the LICENSE file.
+
+package ginpb
+
+import (
+	"bufio"
+	"io"
+	"net"
+	"net/http"
+)
+
+const (
+	noWritten     = -1
+	defaultStatus = http.StatusOK
+)
+
+// ResponseWriter ...
+type ResponseWriter interface {
+	http.ResponseWriter
+	http.Hijacker
+	http.Flusher
+	http.CloseNotifier
+
+	// Status returns the HTTP response status code of the current request.
+	Status() int
+
+	// Size returns the number of bytes already written into the response http body.
+	// See Written()
+	Size() int
+
+	// WriteString writes the string into the response body.
+	WriteString(string) (int, error)
+
+	// Written returns true if the response body was already written.
+	Written() bool
+
+	// WriteHeaderNow forces to write the http header (status code + headers).
+	WriteHeaderNow()
+
+	// Pusher get the http.Pusher for server push
+	Pusher() http.Pusher
+}
+
+type responseWriter struct {
+	http.ResponseWriter
+	size   int
+	status int
+}
+
+var _ ResponseWriter = (*responseWriter)(nil)
+
+func (w *responseWriter) Unwrap() http.ResponseWriter {
+	return w.ResponseWriter
+}
+
+func (w *responseWriter) reset(writer http.ResponseWriter) {
+	w.ResponseWriter = writer
+	w.size = noWritten
+	w.status = defaultStatus
+}
+
+func (w *responseWriter) WriteHeader(code int) {
+	if code > 0 && w.status != code {
+		if w.Written() {
+			// debugPrint("[WARNING] Headers were already written. Wanted to override status code %d with %d", w.status, code)
+			return
+		}
+		w.status = code
+	}
+}
+
+func (w *responseWriter) WriteHeaderNow() {
+	if !w.Written() {
+		w.size = 0
+		w.ResponseWriter.WriteHeader(w.status)
+	}
+}
+
+func (w *responseWriter) Write(data []byte) (n int, err error) {
+	w.WriteHeaderNow()
+	n, err = w.ResponseWriter.Write(data)
+	w.size += n
+	return
+}
+
+func (w *responseWriter) WriteString(s string) (n int, err error) {
+	w.WriteHeaderNow()
+	n, err = io.WriteString(w.ResponseWriter, s)
+	w.size += n
+	return
+}
+
+func (w *responseWriter) Status() int {
+	return w.status
+}
+
+func (w *responseWriter) Size() int {
+	return w.size
+}
+
+func (w *responseWriter) Written() bool {
+	return w.size != noWritten
+}
+
+// Hijack implements the http.Hijacker interface.
+func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if w.size < 0 {
+		w.size = 0
+	}
+	return w.ResponseWriter.(http.Hijacker).Hijack()
+}
+
+// CloseNotify implements the http.CloseNotifier interface.
+func (w *responseWriter) CloseNotify() <-chan bool {
+	return w.ResponseWriter.(http.CloseNotifier).CloseNotify()
+}
+
+// Flush implements the http.Flusher interface.
+func (w *responseWriter) Flush() {
+	w.WriteHeaderNow()
+	w.ResponseWriter.(http.Flusher).Flush()
+}
+
+func (w *responseWriter) Pusher() (pusher http.Pusher) {
+	if pusher, ok := w.ResponseWriter.(http.Pusher); ok {
+		return pusher
+	}
+	return nil
+}