BufferedResponseWriter

BufferedResponseWriter.go

package webutil

import (
	"bufio"
	"bytes"
	"net"
	"net/http"
)

type BufferedResponseWriter struct {
	header http.Header
	hijacked bool
	hijacker http.Hijacker
	responseCode int
	writeBuffer *bytes.Buffer
}

func NewBufferedResponseWriter(writeBuffer *bytes.Buffer) *BufferedResponseWriter {
	return &BufferedResponseWriter{
		header: make(http.Header),
		responseCode: 0,
		writeBuffer: writeBuffer,
	}
}

func (b *BufferedResponseWriter) Header() http.Header {
	return b.header
}

func (b *BufferedResponseWriter) Write(data []byte) (int, error) {
	return b.writeBuffer.Write(data)
}

func (b *BufferedResponseWriter) WriteHeader(responseCode int) {
	if b.responseCode != 0 {
		// Response code already written, do nothing.
		return
	}

	b.responseCode = responseCode
}

func (b *BufferedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if b.hijacker == nil {
		panic("BufferedResponseWriter cannot support Hijack as it has no inner Hijacker.")
	}
	b.hijacked = true

	return b.hijacker.Hijack()
}

// Apply all the buffered HTTP response data to the
// http.ResponseWriter parameter <out>.
func (b *BufferedResponseWriter) Flush(out http.ResponseWriter) (bytesWritten int64, err error) {
	if b.hijacked {
		return
	}
	outHeader := out.Header()
	for k, v := range b.Header() {
		outHeader[k] = v
	}
	if b.responseCode == 0 {
		// Set response code to http.StatusOK if no other value
		// was passed, matching a standard http.ResponseWriter.
		b.responseCode = http.StatusOK
	}
	out.WriteHeader(b.responseCode)
	return b.writeBuffer.WriteTo(out)
}