caddy/modules/caddyhttp/responsewriter.go

package caddyhttp

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

// ResponseWriterWrapper wraps an underlying ResponseWriter and
// promotes its Pusher/Flusher/Hijacker methods as well. To use
// this type, embed a pointer to it within your own struct type
// that implements the http.ResponseWriter interface, then call
// methods on the embedded value. You can make sure your type
// wraps correctly by asserting that it implements the
// HTTPInterfaces interface.
type ResponseWriterWrapper struct {
	http.ResponseWriter
}

// Hijack implements http.Hijacker. It simply calls the underlying
// ResponseWriter's Hijack method if there is one, or returns
// ErrNotImplemented otherwise.
func (rww *ResponseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hj, ok := rww.ResponseWriter.(http.Hijacker); ok {
		return hj.Hijack()
	}
	return nil, nil, ErrNotImplemented
}

// Flush implements http.Flusher. It simply calls the underlying
// ResponseWriter's Flush method if there is one.
func (rww *ResponseWriterWrapper) Flush() {
	if f, ok := rww.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

// Push implements http.Pusher. It simply calls the underlying
// ResponseWriter's Push method if there is one, or returns
// ErrNotImplemented otherwise.
func (rww *ResponseWriterWrapper) Push(target string, opts *http.PushOptions) error {
	if pusher, ok := rww.ResponseWriter.(http.Pusher); ok {
		return pusher.Push(target, opts)
	}
	return ErrNotImplemented
}

// HTTPInterfaces mix all the interfaces that middleware ResponseWriters need to support.
type HTTPInterfaces interface {
	http.ResponseWriter
	http.Pusher
	http.Flusher
	http.Hijacker
}

// ErrNotImplemented is returned when an underlying
// ResponseWriter does not implement the required method.
var ErrNotImplemented = fmt.Errorf("method not implemented")

type responseRecorder struct {
	*ResponseWriterWrapper
	wroteHeader  bool
	statusCode   int
	buf          *bytes.Buffer
	shouldBuffer func(status int) bool
	stream       bool
}

// NewResponseRecorder returns a new ResponseRecorder that can be
// used instead of a real http.ResponseWriter. The recorder is useful
// for middlewares which need to buffer a responder's response and
// process it in its entirety before actually allowing the response to
// be written. Of course, this has a performance overhead, but
// sometimes there is no way to avoid buffering the whole response.
// Still, if at all practical, middlewares should strive to stream
// responses by wrapping Write and WriteHeader methods instead of
// buffering whole response bodies.
//
// Recorders optionally buffer the response. When the headers are
// to be written, shouldBuffer will be called with the status
// code that is being written. The rest of the headers can be read
// from w.Header(). If shouldBuffer returns true, the response
// will be buffered. You can know the response was buffered if
// the Buffered() method returns true. If the response was not
// buffered, Buffered() will return false and that means the
// response bypassed the recorder and was written directly to the
// underlying writer.
//
// Before calling this function in a middleware handler, make a
// new buffer or obtain one from a pool (use the sync.Pool) type.
// Using a pool is generally recommended for performance gains;
// do profiling to ensure this is the case. If using a pool, be
// sure to reset the buffer before using it.
//
// The returned recorder can be used in place of w when calling
// the next handler in the chain. When that handler returns, you
// can read the status code from the recorder's Status() method.
// The response body fills buf if it was buffered, and the headers
// are available via w.Header().
func NewResponseRecorder(w http.ResponseWriter, buf *bytes.Buffer, shouldBuffer func(status int) bool) ResponseRecorder {
	return &responseRecorder{
		ResponseWriterWrapper: &ResponseWriterWrapper{ResponseWriter: w},
		buf:                   buf,
		shouldBuffer:          shouldBuffer,
	}
}

func (rr *responseRecorder) WriteHeader(statusCode int) {
	if rr.wroteHeader {
		return
	}
	rr.statusCode = statusCode
	rr.wroteHeader = true

	// decide whether we should buffer the response
	if rr.shouldBuffer == nil {
		return
	}
	rr.stream = !rr.shouldBuffer(rr.statusCode)
	if rr.stream {
		rr.ResponseWriterWrapper.WriteHeader(rr.statusCode)
	}
}

func (rr *responseRecorder) Write(data []byte) (int, error) {
	rr.WriteHeader(http.StatusOK)
	if rr.stream {
		return rr.ResponseWriterWrapper.Write(data)
	}
	return rr.buf.Write(data)
}

// Status returns the status code that was written, if any.
func (rr *responseRecorder) Status() int {
	return rr.statusCode
}

// Buffer returns the body buffer that rr was created with.
// You should still have your original pointer, though.
func (rr *responseRecorder) Buffer() *bytes.Buffer {
	return rr.buf
}

// Buffered returns whether rr has decided to buffer the response.
func (rr *responseRecorder) Buffered() bool {
	return !rr.stream
}

// ResponseRecorder is a http.ResponseWriter that records
// responses instead of writing them to the client.
type ResponseRecorder interface {
	HTTPInterfaces
	Status() int
	Buffer() *bytes.Buffer
	Buffered() bool
}

// Interface guards
var (
	_ HTTPInterfaces   = (*ResponseWriterWrapper)(nil)
	_ ResponseRecorder = (*responseRecorder)(nil)
)