middleware/recoverer/recoverer.go

package recoverer

import (
	"fmt"
	"net/http"
	"runtime/debug"

	"go.pkg.cx/middleware"
)

// DefaultOptions represents default recoverer middleware options
var DefaultOptions = Options(
	SetLogStackFn(defaultLogStackFn),
	SetLogRecoverFn(defaultLogRecoverFn),
	SetResponseHandler(RespondWithInternalServerError),
)

// Options turns a list of option instances into an option
func Options(opts ...Option) Option {
	return func(r *recoverer) {
		for _, opt := range opts {
			opt(r)
		}
	}
}

// SetLogStackFn sets log stack function
func SetLogStackFn(fn func(stack []byte)) Option {
	return func(r *recoverer) {
		r.logStackFn = fn
	}
}

// SetLogRecoverFn sets log recover information function
func SetLogRecoverFn(fn func(rec interface{})) Option {
	return func(r *recoverer) {
		r.logRecoverFn = fn
	}
}

// SetResponseHandler sets response handler
func SetResponseHandler(fn middleware.ResponseHandle) Option {
	return func(r *recoverer) {
		r.responseHandler = fn
	}
}

// Option configures recoverer middleware
type Option func(r *recoverer)

type recoverer struct {
	logStackFn      func(stack []byte)
	logRecoverFn    func(rec interface{})
	responseHandler middleware.ResponseHandle
}

// Middleware is a recoverer middleware that recovers from panic
func Middleware(opts ...Option) func(next http.Handler) http.Handler {
	rvr := &recoverer{}
	opts = append([]Option{DefaultOptions}, opts...)

	for _, opt := range opts {
		opt(rvr)
	}

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			defer func() {
				if rec := recover(); rec != nil {
					rvr.logRecoverFn(rec)
					rvr.logStackFn(debug.Stack())
					rvr.responseHandler(w, r, nil)
				}
			}()

			next.ServeHTTP(w, r)
		})
	}
}

// RespondWithInternalServerError is a default response handler
func RespondWithInternalServerError(w http.ResponseWriter, _ *http.Request, _ error) {
	http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
}

func defaultLogStackFn(stack []byte) {
	fmt.Println(string(stack))
}

func defaultLogRecoverFn(rec interface{}) {
	fmt.Printf("request panic, %v", rec)
}