Golang: Demo #2 of panic recovery in HTTP handlers and sending HTTP 500

main.go

// This example demonstrates a pluggable middleware approach to
// recovering from panics in HTTP handlers and sending HTTP 500
// responses to the client when panics happen.
//
// The middleware allows to use in any handlers without them
// being aware of anything special.
//
// Utilises middleware concepts, see:
//  - https://github.com/justinas/alice
//  - https://www.alexedwards.net/blog/making-and-using-middleware
//
//
package main

import (
	"bytes"
	"fmt"
	"log"
	"net/http"
	"unicode"
)

// HttpBuffer is a fully buffered implementation of http.ResponseWriter.
// Main use case is to intercept panics and replaced responses with HTTP 500.
// Since this implements http.ResponseWriter, it can be given to the http
// handlers instead of the default one.
type HttpBuffer struct {
	statusCode int
	headerMap  http.Header
	body       *bytes.Buffer
}

// NewHttpBuffer creates new buffered http.ReponseWriter with 200 status code
// and empty headers.
func NewHttpBuffer() *HttpBuffer {
	return &HttpBuffer{
		headerMap:  make(http.Header),
		body:       new(bytes.Buffer),
		statusCode: 200,
	}
}

// SetStatusCode sets the status code of the response.
func (hb *HttpBuffer) SetStatusCode(code int) {
	hb.statusCode = code
}

// GetStatusCode gets the status code of the response.
func (hb *HttpBuffer) GetStatusCode() int {
	return hb.statusCode
}

// GetBody returns the in-memory buffer containing response body.
func (hb *HttpBuffer) GetBody() *bytes.Buffer {
	return hb.body
}

// Send writes and flushes the full response to the outgoing response writer.
// This normally can be done only once per HTTP response.
// Content-Length header will be added.
// The buffer can be reused for the subsequent responses with or without modifications.
func (hb *HttpBuffer) Send(r http.ResponseWriter) error {
	bodyBytes := hb.body.Bytes()
	bodyLength := fmt.Sprintf("%d", len(bodyBytes))

	// Copy headers
	targetHeaders := r.Header()
	for k, vals := range hb.Header() {
		for _, v := range vals {
			targetHeaders.Add(k, v)
		}
	}

	// Set content-length
	targetHeaders.Set("Content-Length", bodyLength)

	// Send over the headers with status
	r.WriteHeader(hb.statusCode)

	// Write the body
	_, err := r.Write(hb.body.Bytes())

	// Flush response if writer supports it
	if flusher, ok := r.(http.Flusher); ok {
		flusher.Flush()
	}

	return err
}

// http.ReponseWriter implementation

// WriteHeader is from http.ReponseWriter interface.
// Calls SetStatusCode. Can be called multiple times as we
// buffer everything.
func (hb *HttpBuffer) WriteHeader(statusCode int) {
	hb.SetStatusCode(statusCode)
}

// Header is from http.ReponseWriter interface
// Allows to set/delete etc response headers.
func (hb *HttpBuffer) Header() http.Header {
	return hb.headerMap
}

// Write is from http.ReponseWriter interface.
// All writes are buffered in memory until Send method is called.
func (hb *HttpBuffer) Write(b []byte) (int, error) {
	return hb.body.Write(b)
}

/////////////////////////////////////////////////////////////////////////////////
// Middleware
//
/////////////////////////////////////////////////////////////////////////////////

// Middleware is a function prototype for a middleware.
// Optionally does something, calls next, then optionally does something again.
// Same as in https://github.com/justinas/alice/blob/master/chain.go
type Middleware func(next http.Handler) http.Handler

// bufferedResponse middleware buffers the entire response in memory
// before sending to the client.
func bufferedResponse(next http.Handler) http.Handler {
	h := func(w http.ResponseWriter, r *http.Request) {
		// log.Printf("bufferedResponse")

		buffer := NewHttpBuffer()
		buffer.Header().Add("X-bufferedResponse", "yes")
		next.ServeHTTP(buffer, r)
		buffer.Send(w)
	}

	return http.HandlerFunc(h)
}

// noPanicResponse middleware recovers from panic in underlying handlers
// and sends HTTP 500 to the client when panic happens.
func noPanicResponse(next http.Handler) http.Handler {
	h := func(w http.ResponseWriter, r *http.Request) {
		// log.Printf("noPanicResponse")

		defer func() {
			if re := recover(); re != nil {
				buffer := NewHttpBuffer()
				buffer.Header().Add("X-noPanicResponse", "yes")
				buffer.SetStatusCode(http.StatusInternalServerError)
				fmt.Fprintf(buffer, "500 - Something bad happened: %v\n", re)
				buffer.Send(w)
			}
		}()

		buffer := NewHttpBuffer()
		buffer.Header().Add("X-noPanicResponse", "yes")
		next.ServeHTTP(buffer, r) // panic here will be recovered
		buffer.Send(w)
	}

	return http.HandlerFunc(h)
}

// requestLogger middleware logs all incoming requests to stderr.
func requestLogger(next http.Handler) http.Handler {
	h := func(w http.ResponseWriter, r *http.Request) {
		// log.Printf("requestLogger")
		log.Printf("%s - %s - %s\n", r.Host, r.Method, r.RequestURI)
		w.Header().Add("X-requestLogger", "yes")
		next.ServeHTTP(w, r)
	}

	return http.HandlerFunc(h)
}

// stringReverser middleware reverses the words of the body.
// demonstrates how the output of the downstream handlers
// can be modified before sending it out, e.g. compression,
// url rewriting etc.
func stringReverser(next http.Handler) http.Handler {
	h := func(w http.ResponseWriter, r *http.Request) {
		// log.Printf("stringReverser")
		buffer := NewHttpBuffer()
		buffer.Header().Add("X-stringReverser", "yes")
		next.ServeHTTP(buffer, r)

		// for unicode to work, we need to use runes not bytes
		body := buffer.GetBody()
		bodyRunes := bytes.Runes(body.Bytes())

		// output
		outRunes := make([]rune, 0, len(bodyRunes))

		// accumulate words here
		word := make([]rune, 0)

		for _, r := range bodyRunes {
			// letters get accumulated into word buffer for later reversal
			if unicode.IsLetter(r) {
				word = append(word, r)
				continue
			}

			// non-letters trigger writing of accumulated word
			// into the output in reverse order
			for j := len(word) - 1; j >= 0; j-- {
				outRunes = append(outRunes, word[j])
			}
			word = make([]rune, 0)

			// also write non-letters as is
			outRunes = append(outRunes, r)
		}

		// flush any words into the output also in reverse
		for j := len(word) - 1; j >= 0; j-- {
			outRunes = append(outRunes, word[j])
		}
		word = make([]rune, 0)

		// rewrite the output with new one
		body.Reset()
		body.Write([]byte(string(outRunes)))

		// send to upstream
		buffer.Send(w)
	}

	return http.HandlerFunc(h)
}

/////////////////////////////////////////////////////////////////////////////////
// BAD HANDLER
//
/////////////////////////////////////////////////////////////////////////////////

var reqNumber = 0

// badHandler panics on every second request after
// writing most of the content out to the response writer.
// It has no awareness of any panic recovery of buffering.
// This is pretty from-examples http handler code we'd
// normally write.
func badHandler(w http.ResponseWriter, r *http.Request) {
	// log.Printf("badHandler")
	w.Header().Add("X-badHandler", "yes")
	fmt.Fprintf(w, "Hi there, I love %s! And some Russian here: архипелаг гулаг\n", r.URL.Path[1:])

	// panic every second request
	reqNumber++
	if reqNumber%2 == 0 {
		panic("hello panic!")
	}
}

/////////////////////////////////////////////////////////////////////////////////

// makeChain chains together supplied middleware. The final
// handler can be added by invoking the returned value like so:
//
// ```
// 	var middlewareChain = makeChain(bufferedResponse, noPanicResponse, requestLogger)
// 	var regularHandler http.Handler
// 	var finalHandler = middlewareChain(regularHandler)
// 	http.Handle("/", finalHandler)
// ```
//
func makeChain(constructors ...Middleware) Middleware {
	return func(finalHandler http.Handler) http.Handler {
		currentHandler := finalHandler
		for i := len(constructors) - 1; i >= 0; i-- {
			c := constructors[i]
			currentHandler = c(currentHandler)
		}
		return currentHandler
	}
}

// makeHandler attaches the handler to the middleware
func makeHandler(handlerFn func(http.ResponseWriter, *http.Request), middleware ...Middleware) http.Handler {
	chain := makeChain(middleware...)
	handler := http.HandlerFunc(handlerFn)
	return chain(handler)
}

/////////////////////////////////////////////////////////////////////////////////
// MAIN
//
// Noddy HTTP server which panics in the handler every second request.
// This gets intercepted and HTTP 500 is issued to the client.
//
// The interception is done using "middleware" approach so that
// the panicking handler doesn't even need to be aware we are doing
// something about it.
//
/////////////////////////////////////////////////////////////////////////////////

func main() {
	// our middleware chain
	chain := makeChain(bufferedResponse, noPanicResponse, requestLogger, stringReverser)

	// the handler attached to the middleware
	handler := http.HandlerFunc(badHandler)

	// handler attached to the middleware
	attachedHandler := makeHandler(handler, chain)

	http.Handle("/", attachedHandler)
	log.Fatal(http.ListenAndServe(":8080", nil))
}