-
Notifications
You must be signed in to change notification settings - Fork 0
/
recovery.go
103 lines (85 loc) · 2.07 KB
/
recovery.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
package recovery
import (
"encoding/json"
"net/http"
"runtime/debug"
"github.com/TanmoySG/wunderDB/internal/server/response"
"github.com/gofiber/fiber/v2"
)
var (
defaultMessage string = "panic on request"
)
type recoveryMessage struct {
Status string `json:"status,omitempty"`
Message *string `json:"message,omitempty"`
Stack *string `json:"stack,omitempty"`
}
func (rm recoveryMessage) Marshal() []byte {
responseJsonBytes, err := json.Marshal(rm)
if err != nil {
return nil
}
return responseJsonBytes
}
// Config defines the config for recovery middleware.
type Config struct {
Next func(c *fiber.Ctx) bool
SendMessage bool
Message *string
EnableStackTrace bool
RecoveryHandler func(c *fiber.Ctx)
}
// ConfigDefault is the default config
var DefaultConfig = Config{
Next: nil,
SendMessage: true,
Message: &defaultMessage,
EnableStackTrace: true,
RecoveryHandler: nil,
}
// Helper function to set default values
func configDefault(config ...Config) Config {
// Return default config if nothing provided
if len(config) < 1 {
return DefaultConfig
}
// Override default config
cfg := config[0]
return cfg
}
// New creates a new middleware handler
func New(config ...Config) fiber.Handler {
// Set default config
cfg := configDefault(config...)
// Return new handler
return func(c *fiber.Ctx) (err error) {
if cfg.Next != nil && cfg.Next(c) {
return c.Next()
}
// Catch panics
defer func() {
if r := recover(); r != nil {
if cfg.SendMessage {
var stackTrace *string
if cfg.EnableStackTrace {
stackTraceString := string(debug.Stack())
stackTrace = &stackTraceString
}
r := recoveryMessage{
Status: response.StatusFailure,
Message: cfg.Message,
Stack: stackTrace,
}
c.Status(http.StatusInternalServerError)
c.Send(r.Marshal())
// extra recovery handler
if cfg.RecoveryHandler != nil {
cfg.RecoveryHandler(c)
}
}
}
}()
// Return err if exist, else move to next handler
return c.Next()
}
}