1 files changed, 113 insertions, 0 deletions

diff --git a/echo/middleware/recover_middleware.go b/echo/middleware/recover_middleware.go
new file mode 100644
index 0000000..310e146
--- /dev/null
+++ b/echo/middleware/recover_middleware.go
@@ -0,0 +1,113 @@
+package middleware
+
+import (
+        "fmt"
+        "net/http"
+        "runtime"
+
+        "github.com/labstack/echo/v4"
+        "github.com/labstack/echo/v4/middleware"
+        "github.com/labstack/gommon/log"
+)
+
+// This is mostly copied from the upstream and modified to:
+//
+// * Report a stack up to 1MB
+// * Print the stack trace to the output if in debug
+
+type (
+        // RecoverConfig defines the config for Recover middleware.
+        RecoverConfig struct {
+                // Skipper defines a function to skip middleware.
+                Skipper middleware.Skipper
+
+                // Size of the stack to be printed.
+                // Optional. Default value 4KB.
+                StackSize int `yaml:"stack_size"`
+
+                // DisableStackAll disables formatting stack traces of all other goroutines
+                // into buffer after the trace for the current goroutine.
+                // Optional. Default value false.
+                DisableStackAll bool `yaml:"disable_stack_all"`
+
+                // DisablePrintStack disables printing stack trace.
+                // Optional. Default value as false.
+                DisablePrintStack bool `yaml:"disable_print_stack"`
+
+                // LogLevel is log level to printing stack trace.
+                // Optional. Default value 0 (Print).
+                LogLevel log.Lvl
+        }
+)
+
+var (
+        // DefaultRecoverConfig is the default Recover middleware config.
+        DefaultRecoverConfig = RecoverConfig{
+                Skipper:           middleware.DefaultSkipper,
+                StackSize:         1024 << 10, // 1 MB
+                DisableStackAll:   false,
+                DisablePrintStack: false,
+                LogLevel:          0,
+        }
+)
+
+// Recover returns a middleware which recovers from panics anywhere in the chain
+// and handles the control to the centralized HTTPErrorHandler.
+func Recover() echo.MiddlewareFunc {
+        return RecoverWithConfig(DefaultRecoverConfig)
+}
+
+// RecoverWithConfig returns a Recover middleware with config.
+// See: `Recover()`.
+func RecoverWithConfig(config RecoverConfig) echo.MiddlewareFunc {
+        // Defaults
+        if config.Skipper == nil {
+                config.Skipper = DefaultRecoverConfig.Skipper
+        }
+        if config.StackSize == 0 {
+                config.StackSize = DefaultRecoverConfig.StackSize
+        }
+
+        return func(next echo.HandlerFunc) echo.HandlerFunc {
+                return func(c echo.Context) error {
+                        if config.Skipper(c) {
+                                return next(c)
+                        }
+
+                        defer func() {
+                                if r := recover(); r != nil {
+                                        err, ok := r.(error)
+                                        if !ok {
+                                                err = fmt.Errorf("%v", r)
+                                        }
+                                        stack := make([]byte, config.StackSize)
+                                        length := runtime.Stack(stack, !config.DisableStackAll)
+                                        if !config.DisablePrintStack {
+                                                msg := fmt.Sprintf("[PANIC RECOVER] %v %s\n", err, stack[:length])
+                                                switch config.LogLevel {
+                                                case log.DEBUG:
+                                                        c.Logger().Debug(msg)
+                                                case log.INFO:
+                                                        c.Logger().Info(msg)
+                                                case log.WARN:
+                                                        c.Logger().Warn(msg)
+                                                case log.ERROR:
+                                                        c.Logger().Error(msg)
+                                                case log.OFF:
+                                                        // None.
+                                                default:
+                                                        c.Logger().Print(msg)
+                                                }
+                                        }
+                                        if c.Echo().Debug {
+                                                msg := fmt.Sprintf("%v %s\n", err, stack[:length])
+                                                c.String(http.StatusInternalServerError, msg)
+                                        } else {
+                                                c.Error(err)
+                                        }
+                                }
+                        }()
+                        return next(c)
+                }
+        }
+}