github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/handlers/recovery.go (about)

     1  package handlers
     2  
     3  import (
     4  	"log"
     5  	"runtime/debug"
     6  
     7  	http "github.com/hxx258456/ccgo/gmhttp"
     8  )
     9  
    10  // RecoveryHandlerLogger is an interface used by the recovering handler to print logs.
    11  type RecoveryHandlerLogger interface {
    12  	Println(...interface{})
    13  }
    14  
    15  type recoveryHandler struct {
    16  	handler    http.Handler
    17  	logger     RecoveryHandlerLogger
    18  	printStack bool
    19  }
    20  
    21  // RecoveryOption provides a functional approach to define
    22  // configuration for a handler; such as setting the logging
    23  // whether or not to print stack traces on panic.
    24  type RecoveryOption func(http.Handler)
    25  
    26  func parseRecoveryOptions(h http.Handler, opts ...RecoveryOption) http.Handler {
    27  	for _, option := range opts {
    28  		option(h)
    29  	}
    30  
    31  	return h
    32  }
    33  
    34  // RecoveryHandler is HTTP middleware that recovers from a panic,
    35  // logs the panic, writes http.StatusInternalServerError, and
    36  // continues to the next handler.
    37  //
    38  // Example:
    39  //
    40  //  r := mux.NewRouter()
    41  //  r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    42  //  	panic("Unexpected error!")
    43  //  })
    44  //
    45  //  http.ListenAndServe(":1123", handlers.RecoveryHandler()(r))
    46  func RecoveryHandler(opts ...RecoveryOption) func(h http.Handler) http.Handler {
    47  	return func(h http.Handler) http.Handler {
    48  		r := &recoveryHandler{handler: h}
    49  		return parseRecoveryOptions(r, opts...)
    50  	}
    51  }
    52  
    53  // RecoveryLogger is a functional option to override
    54  // the default logger
    55  func RecoveryLogger(logger RecoveryHandlerLogger) RecoveryOption {
    56  	return func(h http.Handler) {
    57  		r := h.(*recoveryHandler)
    58  		r.logger = logger
    59  	}
    60  }
    61  
    62  // PrintRecoveryStack is a functional option to enable
    63  // or disable printing stack traces on panic.
    64  func PrintRecoveryStack(print bool) RecoveryOption {
    65  	return func(h http.Handler) {
    66  		r := h.(*recoveryHandler)
    67  		r.printStack = print
    68  	}
    69  }
    70  
    71  func (h recoveryHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
    72  	defer func() {
    73  		if err := recover(); err != nil {
    74  			w.WriteHeader(http.StatusInternalServerError)
    75  			h.log(err)
    76  		}
    77  	}()
    78  
    79  	h.handler.ServeHTTP(w, req)
    80  }
    81  
    82  func (h recoveryHandler) log(v ...interface{}) {
    83  	if h.logger != nil {
    84  		h.logger.Println(v...)
    85  	} else {
    86  		log.Println(v...)
    87  	}
    88  
    89  	if h.printStack {
    90  		stack := string(debug.Stack())
    91  		if h.logger != nil {
    92  			h.logger.Println(stack)
    93  		} else {
    94  			log.Println(stack)
    95  		}
    96  	}
    97  }