FEATURE: support Hijacker, CloseNotifier, Flusher; export Size() and Written()

Add support and tests for http's Hijacker, CloseNotifier, and Flusher interfaces in the ResponseWriter.

Also, export some more methods on ResponseWriter: Size() and Written() which return how many bytes were written and whether the header was written yet.

Author: cypriss

TODO

@@ -1,39 +1,76 @@
 package web
 
 import (
+	"bufio"
+	"fmt"
+	"net"
 	"net/http"
 )
 
 // ResponseWriter includes net/http's ResponseWriter and adds a StatusCode() method to obtain the written status code.
 // A ResponseWriter is sent to handlers on each request.
 type ResponseWriter interface {
 	http.ResponseWriter
+	http.Flusher
+	http.Hijacker
+	http.CloseNotifier
 
 	// StatusCode returns the written status code, or 0 if none has been written yet.
 	StatusCode() int
+	// Written returns whether the header has been written yet.
+	Written() bool
+	// Size returns the size in bytes of the body written so far.
+	Size() int
 }
 
 type appResponseWriter struct {
 	http.ResponseWriter
 	statusCode int
-	written    bool
+	size       int
 }
 
 // Don't need this yet because we get it for free:
 func (w *appResponseWriter) Write(data []byte) (n int, err error) {
-	if !w.written {
+	if w.statusCode == 0 {
 		w.statusCode = http.StatusOK
-		w.written = true
 	}
-	return w.ResponseWriter.Write(data)
+	size, err := w.ResponseWriter.Write(data)
+	w.size += size
+	return size, err
 }
 
 func (w *appResponseWriter) WriteHeader(statusCode int) {
 	w.statusCode = statusCode
-	w.written = true
 	w.ResponseWriter.WriteHeader(statusCode)
 }
 
 func (w *appResponseWriter) StatusCode() int {
 	return w.statusCode
 }
+
+func (w *appResponseWriter) Written() bool {
+	return w.statusCode != 0
+}
+
+func (w *appResponseWriter) Size() int {
+	return w.size
+}
+
+func (w *appResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	hijacker, ok := w.ResponseWriter.(http.Hijacker)
+	if !ok {
+		return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
+	}
+	return hijacker.Hijack()
+}
+
+func (w *appResponseWriter) CloseNotify() <-chan bool {
+	return w.ResponseWriter.(http.CloseNotifier).CloseNotify()
+}
+
+func (w *appResponseWriter) Flush() {
+	flusher, ok := w.ResponseWriter.(http.Flusher)
+	if ok {
+		flusher.Flush()
+	}
+}

response_writer.go

@@ -0,0 +1,113 @@
+package web
+
+import (
+	"bufio"
+	"github.com/stretchr/testify/assert"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+)
+
+type hijackableResponse struct {
+	Hijacked bool
+}
+
+func (h *hijackableResponse) Header() http.Header {
+	return nil
+}
+func (h *hijackableResponse) Write(buf []byte) (int, error) {
+	return 0, nil
+}
+func (h *hijackableResponse) WriteHeader(code int) {
+	// no-op
+}
+func (h *hijackableResponse) Flush() {
+	// no-op
+}
+func (h *hijackableResponse) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	h.Hijacked = true
+	return nil, nil, nil
+}
+func (h *hijackableResponse) CloseNotify() <-chan bool {
+	return nil
+}
+
+type closeNotifyingRecorder struct {
+	*httptest.ResponseRecorder
+	closed chan bool
+}
+
+func (c *closeNotifyingRecorder) close() {
+	c.closed <- true
+}
+
+func (c *closeNotifyingRecorder) CloseNotify() <-chan bool {
+	return c.closed
+}
+
+func TestResponseWriterWrite(t *testing.T) {
+	rec := httptest.NewRecorder()
+	rw := ResponseWriter(&appResponseWriter{ResponseWriter: rec})
+
+	assert.Equal(t, rw.Written(), false)
+
+	n, err := rw.Write([]byte("Hello world"))
+	assert.Equal(t, n, 11)
+	assert.NoError(t, err)
+
+	assert.Equal(t, n, 11)
+	assert.Equal(t, rec.Code, rw.StatusCode())
+	assert.Equal(t, rec.Code, http.StatusOK)
+	assert.Equal(t, rec.Body.String(), "Hello world")
+	assert.Equal(t, rw.Size(), 11)
+	assert.Equal(t, rw.Written(), true)
+}
+
+func TestResponseWriterWriteHeader(t *testing.T) {
+	rec := httptest.NewRecorder()
+	rw := ResponseWriter(&appResponseWriter{ResponseWriter: rec})
+
+	rw.WriteHeader(http.StatusNotFound)
+	assert.Equal(t, rec.Code, rw.StatusCode())
+	assert.Equal(t, rec.Code, http.StatusNotFound)
+}
+
+func TestResponseWriterHijack(t *testing.T) {
+	hijackable := &hijackableResponse{}
+	rw := ResponseWriter(&appResponseWriter{ResponseWriter: hijackable})
+	hijacker, ok := rw.(http.Hijacker)
+	assert.True(t, ok)
+	_, _, err := hijacker.Hijack()
+	assert.NoError(t, err)
+	assert.True(t, hijackable.Hijacked)
+}
+
+func TestResponseWriterHijackNotOK(t *testing.T) {
+	rw := ResponseWriter(&appResponseWriter{ResponseWriter: httptest.NewRecorder()})
+	_, _, err := rw.Hijack()
+	assert.Error(t, err)
+}
+
+func TestResponseWriterFlush(t *testing.T) {
+	rw := ResponseWriter(&appResponseWriter{ResponseWriter: httptest.NewRecorder()})
+	rw.Flush()
+}
+
+func TestResponseWriterCloseNotify(t *testing.T) {
+	rec := &closeNotifyingRecorder{
+		httptest.NewRecorder(),
+		make(chan bool, 1),
+	}
+	rw := ResponseWriter(&appResponseWriter{ResponseWriter: rec})
+	closed := false
+	notifier := rw.(http.CloseNotifier).CloseNotify()
+	rec.close()
+	select {
+	case <-notifier:
+		closed = true
+	case <-time.After(time.Second):
+	}
+	assert.True(t, closed)
+}