narco/response_writter.go

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,
		}
	}

}