1 files changed, 48 insertions, 0 deletions
diff --git a/buffer.go b/buffer.go
new file mode 100644
index 0000000..e31bebd
--- /dev/null
+++ b/buffer.go
@@ -0,0 +1,48 @@
+package main
+
+// Taken from https://github.com/goods/httpbuf
+
+import (
+	"bytes"
+	"net/http"
+	"sync"
+)
+
+//Buffer is a type that implements http.ResponseWriter but buffers all the data
+//and headers.
+type Buffer struct {
+	bytes.Buffer
+	Status  int
+	headers http.Header
+	once    sync.Once
+}
+
+//Header implements the header method of http.ResponseWriter
+func (b *Buffer) Header() http.Header {
+	b.once.Do(func() {
+		b.headers = make(http.Header)
+	})
+	return b.headers
+}
+
+//WriteHeader implements the WriteHeader method of http.ResponseWriter
+func (b *Buffer) WriteHeader(resp int) {
+	b.Status = resp
+}
+
+//Apply takes an http.ResponseWriter and calls the required methods on it to
+//output the buffered headers, response code, and data. It returns the number
+//of bytes written and any errors flushing.
+func (b *Buffer) Apply(w http.ResponseWriter) (n int, err error) {
+	if len(b.headers) > 0 {
+		h := w.Header()
+		for key, val := range b.headers {
+			h[key] = val
+		}
+	}
+	if b.Status > 0 {
+		w.WriteHeader(b.Status)
+	}
+	n, err = w.Write(b.Bytes())
+	return
+}