本站源代码
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

92 lines
2.1KB

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