Add CustomRecovery and CustomRecoveryWithWriter methods

recovery.go

@@ -26,13 +26,26 @@ var (
 	slash     = []byte("/")
 )
 
+// RecoveryFunc defines the function passable to CustomRecovery.
+type RecoveryFunc func(c *Context, err interface{})
+
 // Recovery returns a middleware that recovers from any panics and writes a 500 if there was one.
 func Recovery() HandlerFunc {
 	return RecoveryWithWriter(DefaultErrorWriter)
 }
 
+//CustomRecovery returns a middleware that recovers from any panics and calls the provided handle func to handle it.
+func CustomRecovery(handle RecoveryFunc) HandlerFunc {
+	return CustomRecoveryWithWriter(DefaultErrorWriter, handle)
+}
+
 // 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 {
+	return CustomRecoveryWithWriter(out, defaultHandleRecovery)
+}
+
+// CustomRecoveryWithWriter returns a middleware for a given writer that recovers from any panics and calls the provided handle func to handle it.
+func CustomRecoveryWithWriter(out io.Writer, handle RecoveryFunc) HandlerFunc {
 	var logger *log.Logger
 	if out != nil {
 		logger = log.New(out, "\n\n\x1b[31m", log.LstdFlags)
@@ -70,13 +83,12 @@ func RecoveryWithWriter(out io.Writer) HandlerFunc {
 							timeFormat(time.Now()), err, stack, reset)
 					}
 				}
-
-				// If the connection is dead, we can't write a status to it.
 				if brokenPipe {
+					// If the connection is dead, we can't write a status to it.
 					c.Error(err.(error)) // nolint: errcheck
 					c.Abort()
 				} else {
-					c.AbortWithStatus(http.StatusInternalServerError)
+					handle(c, err)
 				}
 			}
 		}()
@@ -84,6 +96,10 @@ func RecoveryWithWriter(out io.Writer) HandlerFunc {
 	}
 }
 
+func defaultHandleRecovery(c *Context, err interface{}) {
+	c.AbortWithStatus(http.StatusInternalServerError)
+}
+
 // stack returns a nicely formatted stack frame, skipping skip frames.
 func stack(skip int) []byte {
 	buf := new(bytes.Buffer) // the returned data

recovery_test.go

@@ -144,3 +144,37 @@ func TestPanicWithBrokenPipe(t *testing.T) {
 		})
 	}
 }
+
+func TestCustomRecoveryWithWriter(t *testing.T) {
+	errBuffer := new(bytes.Buffer)
+	buffer := new(bytes.Buffer)
+	router := New()
+	handleRecovery := func(c *Context, err interface{}) {
+		errBuffer.WriteString(err.(string))
+		c.AbortWithStatus(http.StatusBadRequest)
+	}
+	router.Use(CustomRecoveryWithWriter(buffer, handleRecovery))
+	router.GET("/recovery", func(_ *Context) {
+		panic("Oupps, Houston, we have a problem")
+	})
+	// RUN
+	w := performRequest(router, "GET", "/recovery")
+	// TEST
+	assert.Equal(t, http.StatusBadRequest, w.Code)
+	assert.Contains(t, buffer.String(), "panic recovered")
+	assert.Contains(t, buffer.String(), "Oupps, Houston, we have a problem")
+	assert.Contains(t, buffer.String(), "TestCustomRecoveryWithWriter")
+	assert.NotContains(t, buffer.String(), "GET /recovery")
+
+	// Debug mode prints the request
+	SetMode(DebugMode)
+	// RUN
+	w = performRequest(router, "GET", "/recovery")
+	// TEST
+	assert.Equal(t, http.StatusBadRequest, w.Code)
+	assert.Contains(t, buffer.String(), "GET /recovery")
+
+	assert.Equal(t, strings.Repeat("Oupps, Houston, we have a problem", 2), errBuffer.String())
+
+	SetMode(TestMode)
+}