From 273f2410fa3a99ce2617d35bc6ae6292e454c341 Mon Sep 17 00:00:00 2001
From: Alexandre Tuleu <alexandre.tuleu.2005@polytechnique.org>
Date: Wed, 12 Aug 2015 14:01:24 +0200
Subject: [PATCH] Adds a custom, type safe ResponseWriter

---
 response_writter.go | 205 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 205 insertions(+)
 create mode 100644 response_writter.go

diff --git a/response_writter.go b/response_writter.go
new file mode 100644
index 0000000..3dc1394
--- /dev/null
+++ b/response_writter.go
@@ -0,0 +1,205 @@
+package narco
+
+import (
+	"fmt"
+	"net/http"
+)
+
+// those are intern utility to get the data
+
+// yep we have to write o
+
+type ResponseWriter interface {
+	http.ResponseWriter
+
+	Status() int
+	Size() int
+
+	HeaderWritten() bool
+	//Could set other stuff if needed
+}
+
+type responseWriterImpl struct {
+	http.ResponseWriter
+	size, status int
+}
+
+func (rw *responseWriterImpl) HeaderWritten() bool {
+	return rw.status > 0
+}
+
+func (rw *responseWriterImpl) Write(data []byte) (int, error) {
+	// do as told in http.Interface
+	if rw.HeaderWritten() == false {
+		rw.WriteHeader(http.StatusOK)
+	}
+
+	size, err := rw.ResponseWriter.Write(data)
+	rw.size += size
+	return size, err
+}
+
+func (rw *responseWriterImpl) WriteHeader(status int) {
+	if status <= 0 {
+		panic(fmt.Sprintf("Invalid providen http status: %s", status))
+	}
+
+	rw.status = status
+	rw.ResponseWriter.WriteHeader(status)
+}
+
+func (rw *responseWriterImpl) Header() http.Header {
+	return rw.ResponseWriter.Header()
+}
+
+func (rw *responseWriterImpl) Size() int {
+	return rw.size
+}
+
+func (rw *responseWriterImpl) Status() int {
+	return rw.status
+}
+
+func newResponseWriterImpl(rw http.ResponseWriter) *responseWriterImpl {
+	return &responseWriterImpl{
+		ResponseWriter: rw,
+		status:         -1,
+		size:           0,
+	}
+}
+
+type responseWriterC struct {
+	ResponseWriter
+	http.CloseNotifier
+}
+
+type responseWriterF struct {
+	ResponseWriter
+	http.Flusher
+}
+
+type responseWriterFC struct {
+	ResponseWriter
+	http.Flusher
+	http.CloseNotifier
+}
+
+type responseWriterH struct {
+	ResponseWriter
+	http.Hijacker
+}
+
+type responseWriterHC struct {
+	ResponseWriter
+	http.Hijacker
+	http.CloseNotifier
+}
+
+type responseWriterHF struct {
+	ResponseWriter
+	http.Hijacker
+	http.Flusher
+}
+
+type responseWriterHFC struct {
+	ResponseWriter
+	http.Hijacker
+	http.Flusher
+	http.CloseNotifier
+}
+
+type Creator func(ResponseWriter, http.Hijacker, http.Flusher, http.CloseNotifier) ResponseWriter
+
+func makeIdent(hasHijacker, hasFlusher, hasCloseNotifier bool) int {
+	res := 0
+	if hasCloseNotifier {
+		res += 1
+	}
+	if hasFlusher {
+		res += 1 << 1
+	}
+	if hasCloseNotifier {
+		res += 1 << 2
+	}
+	return res
+}
+
+var rwFactory = make(map[int]Creator, 8)
+
+func WrapResponseWritter(rw http.ResponseWriter) ResponseWriter {
+	if res, ok := rw.(ResponseWriter); ok == true {
+		return res
+	}
+
+	h, hOk := rw.(http.Hijacker)
+	f, fOk := rw.(http.Flusher)
+	c, cOk := rw.(http.CloseNotifier)
+
+	creator, ok := rwFactory[makeIdent(hOk, fOk, cOk)]
+	if ok == false {
+		panic("Case is not covered")
+	}
+
+	return creator(newResponseWriterImpl(rw), h, f, c)
+}
+
+func init() {
+	rwFactory[makeIdent(false, false, false)] = func(rw ResponseWriter, h http.Hijacker, f http.Flusher, c http.CloseNotifier) ResponseWriter {
+		return rw
+	}
+
+	rwFactory[makeIdent(false, false, true)] = func(rw ResponseWriter, h http.Hijacker, f http.Flusher, c http.CloseNotifier) ResponseWriter {
+		return responseWriterC{
+			ResponseWriter: rw,
+			CloseNotifier:  c,
+		}
+	}
+
+	rwFactory[makeIdent(false, true, false)] = func(rw ResponseWriter, h http.Hijacker, f http.Flusher, c http.CloseNotifier) ResponseWriter {
+		return responseWriterF{
+			ResponseWriter: rw,
+			Flusher:        f,
+		}
+	}
+
+	rwFactory[makeIdent(false, true, true)] = func(rw ResponseWriter, h http.Hijacker, f http.Flusher, c http.CloseNotifier) ResponseWriter {
+		return responseWriterFC{
+			ResponseWriter: rw,
+			Flusher:        f,
+			CloseNotifier:  c,
+		}
+	}
+
+	rwFactory[makeIdent(true, false, false)] = func(rw ResponseWriter, h http.Hijacker, f http.Flusher, c http.CloseNotifier) ResponseWriter {
+		return responseWriterH{
+			ResponseWriter: rw,
+			Hijacker:       h,
+		}
+	}
+
+	rwFactory[makeIdent(true, false, true)] = func(rw ResponseWriter, h http.Hijacker, f http.Flusher, c http.CloseNotifier) ResponseWriter {
+		return responseWriterHC{
+			ResponseWriter: rw,
+			Hijacker:       h,
+			CloseNotifier:  c,
+		}
+	}
+
+	rwFactory[makeIdent(true, true, false)] = func(rw ResponseWriter, h http.Hijacker, f http.Flusher, c http.CloseNotifier) ResponseWriter {
+		return responseWriterHF{
+			ResponseWriter: rw,
+			Hijacker:       h,
+			Flusher:        f,
+		}
+	}
+
+	rwFactory[makeIdent(true, true, true)] = func(rw ResponseWriter, h http.Hijacker, f http.Flusher, c http.CloseNotifier) ResponseWriter {
+		return responseWriterHFC{
+			ResponseWriter: rw,
+			Hijacker:       h,
+			Flusher:        f,
+			CloseNotifier:  c,
+		}
+	}
+
+}