iris/response_writer.go

package iris

import (
	"bufio"
	"net"
	"net/http"
	"sync"

	"github.com/kataras/go-errors"
	"github.com/kataras/go-fs"
	"github.com/klauspost/compress/gzip"
)

type gzipResponseWriter struct {
	http.ResponseWriter
	http.Flusher
	gzipWriter *gzip.Writer
}

var gzpool = sync.Pool{New: func() interface{} { return &gzipResponseWriter{} }}

func acquireGzipResponseWriter(underline http.ResponseWriter) *gzipResponseWriter {
	w := gzpool.Get().(*gzipResponseWriter)
	w.ResponseWriter = underline
	w.gzipWriter = fs.AcquireGzipWriter(w.ResponseWriter)
	return w
}

func releaseGzipResponseWriter(w *gzipResponseWriter) {
	fs.ReleaseGzipWriter(w.gzipWriter)
	gzpool.Put(w)
}

// Write compresses and writes that data to the underline response writer
func (w *gzipResponseWriter) Write(contents []byte) (int, error) {
	return w.gzipWriter.Write(contents)
}

var rpool = sync.Pool{New: func() interface{} { return &ResponseWriter{} }}

func acquireResponseWriter(underline http.ResponseWriter) *ResponseWriter {
	w := rpool.Get().(*ResponseWriter)
	w.ResponseWriter = underline
	w.headers = underline.Header()
	return w
}

func releaseResponseWriter(w *ResponseWriter) {
	w.headers = nil
	w.ResponseWriter = nil
	w.statusCode = 0
	w.beforeFlush = nil
	w.ResetBody()
	rpool.Put(w)
}

// A ResponseWriter interface is used by an HTTP handler to
// construct an HTTP response.
//
// A ResponseWriter may not be used after the Handler.ServeHTTP method
// has returned.
type ResponseWriter struct {
	// yes only one callback, we need simplicity here because on EmitError the beforeFlush events should NOT be cleared
	// but the response is cleared.
	// Sometimes is useful to keep the event,
	// so we keep one func only and let the user decide when he/she wants to override it with an empty func before the EmitError (context's behavior)
	beforeFlush func()
	http.ResponseWriter
	// these three fields are setted on flushBody which runs only once on the end of the handler execution.
	// this helps the performance on multi-write and keep tracks the body, status code and headers in order to run each transaction
	// on its own
	chunks     []byte      // keep track of the body in order to be resetable and useful inside custom transactions
	statusCode int         // the saved status code which will be used from the cache service
	headers    http.Header // the saved headers
}

// Header returns the header map that will be sent by
// WriteHeader. Changing the header after a call to
// WriteHeader (or Write) has no effect unless the modified
// headers were declared as trailers by setting the
// "Trailer" header before the call to WriteHeader (see example).
// To suppress implicit response headers, set their value to nil.
func (w *ResponseWriter) Header() http.Header {
	return w.headers
}

// StatusCode returns the status code header value
func (w *ResponseWriter) StatusCode() int {
	return w.statusCode
}

// Adds the contents to the body reply, it writes the contents temporarily
// to a value in order to be flushed at the end of the request,
// this method give us the opportunity to reset the body if needed.
//
// If WriteHeader has not yet been called, Write calls
// WriteHeader(http.StatusOK) before writing the data. If the Header
// does not contain a Content-Type line, Write adds a Content-Type set
// to the result of passing the initial 512 bytes of written data to
// DetectContentType.
//
// Depending on the HTTP protocol version and the client, calling
// Write or WriteHeader may prevent future reads on the
// Request.Body. For HTTP/1.x requests, handlers should read any
// needed request body data before writing the response. Once the
// headers have been flushed (due to either an explicit Flusher.Flush
// call or writing enough data to trigger a flush), the request body
// may be unavailable. For HTTP/2 requests, the Go HTTP server permits
// handlers to continue to read the request body while concurrently
// writing the response. However, such behavior may not be supported
// by all HTTP/2 clients. Handlers should read before writing if
// possible to maximize compatibility.
func (w *ResponseWriter) Write(contents []byte) (int, error) {
	w.chunks = append(w.chunks, contents...)
	return len(w.chunks), nil
}

// Body returns the body tracked from the writer so far
// do not use this for edit.
func (w *ResponseWriter) Body() []byte {
	return w.chunks
}

// SetBodyString overrides the body and sets it to a string value
func (w *ResponseWriter) SetBodyString(s string) {
	w.chunks = []byte(s)
}

// SetBody overrides the body and sets it to a slice of bytes value
func (w *ResponseWriter) SetBody(b []byte) {
	w.chunks = b
}

// ResetBody resets the response body
func (w *ResponseWriter) ResetBody() {
	w.chunks = w.chunks[0:0]
}

// ResetHeaders clears the temp headers
func (w *ResponseWriter) ResetHeaders() {
	// original response writer's headers are empty.
	w.headers = w.ResponseWriter.Header()
}

// Reset resets the response body, headers and the status code header
func (w *ResponseWriter) Reset() {
	w.ResetHeaders()
	w.statusCode = 0
	w.ResetBody()
}

// WriteHeader sends an HTTP response header with status code.
// If WriteHeader is not called explicitly, the first call to Write
// will trigger an implicit WriteHeader(http.StatusOK).
// Thus explicit calls to WriteHeader are mainly used to
// send error codes.
func (w *ResponseWriter) WriteHeader(statusCode int) {
	w.statusCode = statusCode
}

// ContentType returns the content type, if not setted returns empty string
func (w *ResponseWriter) ContentType() string {
	return w.headers.Get(contentType)
}

// SetContentType sets the content type header
func (w *ResponseWriter) SetContentType(cType string) {
	w.headers.Set(contentType, cType)
}

var errHijackNotSupported = errors.New("Hijack is not supported to this response writer!")

// Hijack lets the caller take over the connection.
// After a call to Hijack(), the HTTP server library
// will not do anything else with the connection.
//
// It becomes the caller's responsibility to manage
// and close the connection.
//
// The returned net.Conn may have read or write deadlines
// already set, depending on the configuration of the
// Server. It is the caller's responsibility to set
// or clear those deadlines as needed.
func (w *ResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if h, isHijacker := w.ResponseWriter.(http.Hijacker); isHijacker {
		return h.Hijack()
	}

	return nil, nil, errHijackNotSupported
}

// SetBeforeFlush registers the unique callback which called exactly before the response is flushed to the client
func (w *ResponseWriter) SetBeforeFlush(cb func()) {
	w.beforeFlush = cb
}

// flushResponse the full body, headers and status code to the underline response writer
// called automatically at the end of each request, see ReleaseCtx
func (w *ResponseWriter) flushResponse() {

	if w.beforeFlush != nil {
		w.beforeFlush()
	}

	if w.statusCode > 0 {
		w.ResponseWriter.WriteHeader(w.statusCode)
	}

	if w.headers != nil {
		for k, values := range w.headers {
			for i := range values {
				w.ResponseWriter.Header().Add(k, values[i])
			}
		}
	}

	if len(w.chunks) > 0 {
		w.ResponseWriter.Write(w.chunks)
	}
}

// Flush sends any buffered data to the client.
func (w *ResponseWriter) Flush() {
	w.flushResponse()
	// The Flusher interface is implemented by ResponseWriters that allow
	// an HTTP handler to flush buffered data to the client.
	//
	// The default HTTP/1.x and HTTP/2 ResponseWriter implementations
	// support Flusher, but ResponseWriter wrappers may not. Handlers
	// should always test for this ability at runtime.
	//
	// Note that even for ResponseWriters that support Flush,
	// if the client is connected through an HTTP proxy,
	// the buffered data may not reach the client until the response
	// completes.
	if fl, isFlusher := w.ResponseWriter.(http.Flusher); isFlusher {
		fl.Flush()
	}
}

// clone returns a clone of this response writer
// it copies the header, status code, headers and the beforeFlush finally  returns a new ResponseWriter
func (w *ResponseWriter) clone() *ResponseWriter {
	wc := &ResponseWriter{}
	wc.ResponseWriter = w.ResponseWriter
	wc.statusCode = w.statusCode
	wc.headers = w.headers
	wc.chunks = w.chunks[0:]
	wc.beforeFlush = w.beforeFlush
	return wc
}

// writeTo writes a response writer (temp: status code, headers and body) to another response writer
func (w *ResponseWriter) writeTo(to *ResponseWriter) {
	// set the status code, failure status code are first class
	if w.statusCode > 0 {
		to.statusCode = w.statusCode
	}

	// append the headers
	if w.headers != nil {
		for k, values := range w.headers {
			for _, v := range values {
				if to.headers.Get(v) == "" {
					to.headers.Add(k, v)
				}
			}
		}

	}

	// append the body
	if len(w.chunks) > 0 {
		to.Write(w.chunks)
	}

	if w.beforeFlush != nil {
		to.SetBeforeFlush(w.beforeFlush)
	}
}