Add RecoveryWithWriterAndCallback

This commit is contained in:
letsfire 2019-10-13 12:48:50 +08:00
parent f7becac7bc
commit 66fa76f0f0
1 changed files with 13 additions and 8 deletions

View File

@ -33,6 +33,18 @@ func Recovery() HandlerFunc {
// RecoveryWithWriter returns a middleware for a given writer that recovers from any panics and writes a 500 if there was one. // RecoveryWithWriter returns a middleware for a given writer that recovers from any panics and writes a 500 if there was one.
func RecoveryWithWriter(out io.Writer) HandlerFunc { func RecoveryWithWriter(out io.Writer) HandlerFunc {
return RecoveryWithWriterAndCallback(out, func(c *Context, err interface{}, brokenPipe bool) {
// If the connection is dead, we can't write a status to it.
if brokenPipe {
c.Error(err.(error)) // nolint: errcheck
c.Abort()
} else {
c.AbortWithStatus(http.StatusInternalServerError)
}
})
}
func RecoveryWithWriterAndCallback(out io.Writer, fn func(c *Context, err interface{}, brokenPipe bool)) HandlerFunc {
var logger *log.Logger var logger *log.Logger
if out != nil { if out != nil {
logger = log.New(out, "\n\n\x1b[31m", log.LstdFlags) logger = log.New(out, "\n\n\x1b[31m", log.LstdFlags)
@ -70,14 +82,7 @@ func RecoveryWithWriter(out io.Writer) HandlerFunc {
timeFormat(time.Now()), err, stack, reset) timeFormat(time.Now()), err, stack, reset)
} }
} }
fn(c, err, brokenPipe)
// If the connection is dead, we can't write a status to it.
if brokenPipe {
c.Error(err.(error)) // nolint: errcheck
c.Abort()
} else {
c.AbortWithStatus(http.StatusInternalServerError)
}
} }
}() }()
c.Next() c.Next()