Wrapping a ResponseWriter to capture the status code

example.go

// Create our own MyResponseWriter to wrap a standard http.ResponseWriter
// so we can store the status code.
type MyResponseWriter struct {
	status int 
	http.ResponseWriter
}

func NewMyResponseWriter(res http.ResponseWriter) *MyResponseWriter {
	// Default the status code to 200
	return &MyResponseWriter{200, res}
}

// Give a way to get the status
func (w MyResponseWriter) Status() int {
	return w.status
}

// Satisfy the http.ResponseWriter interface
func (w MyResponseWriter) Header() http.Header {
	return w.ResponseWriter.Header()
}

func (w MyResponseWriter) Write(data []byte) (int, error) {
	return w.ResponseWriter.Write(data)
}

func (w MyResponseWriter) WriteHeader(statusCode int) {
	// Store the status code
	w.status = statusCode

	// Write the status code onward.
	w.ResponseWriter.WriteHeader(statusCode)
}

Embedding the http.ResponseWriter should mean that MyResponseWriter already satisfies the interface?

I found this on a Google search so I thought I'd point out a few issues with the code in this gist for future visitors:

• If the header is being written by the wrapped ResponseWriter implementation, it won't know to call your WriteHeader method so you have to call it yourself
• The function receivers are not pointers so the status update won't actually be persisted on the original MyResponseWriter instance
• (Minor) There's no need to manually forward the methods of embedded types (you don't need define the Header method)

Here's the fixed up code which should work in all cases:

type MyResponseWriter struct {
    http.ResponseWriter
    status      int
    wroteHeader bool
}

func NewMyResponseWriter(w http.ResponseWriter) *MyResponseWriter {
  return &MyResponseWriter{ResponseWriter: w}
}

func (w *MyResponseWriter) Status() int {
    return w.status
}

func (w *MyResponseWriter) Write(p []byte) (n int, err error) {
    if !w.wroteHeader {
        w.WriteHeader(http.StatusOK)
    }
    return w.ResponseWriter.Write(p)
}

func (w *MyResponseWriter) WriteHeader(code int) {
    w.ResponseWriter.WriteHeader(code)
    // Check after in case there's error handling in the wrapped ResponseWriter.
    if w.wroteHeader {
        return
    }
    w.status = code
    w.wroteHeader = true
}

I've created a package called called httpsnoop which fixes multiple problems with the original gist as well as the improved solution below.

The problems with the existing solutions are described here: https://github.com/felixge/httpsnoop#why-this-package-exists

http.Pusher interface can be implemented, as the documentation states

// Push returns ErrNotSupported if the client has disabled push or if push
// is not supported on the underlying connection.

func (w *MyResponseWriter) Push(target string, opts *http.PushOptions) error {
	if p, ok := w.ResponseWriter.(http.Pusher); ok {
		return p.Push(target, opts)
	}
	return http.ErrNotSupported
}

http.Hijacker interface isn't implemented in http2 connections.