Repository: gin-gonic/gin
Pull request #2322 (2020-04-16 09:44:06)

Add CustomRecovery builtin middleware 
- With pull requests:
  - Open your pull request against `master`
  - Your pull request should have no more than two commits, if not you should squash them.
  - It should pass all tests in the available continuous integration systems such as TravisCI.
  - You should add/modify tests to cover your proposed code changes.
  - If your pull request contains a new feature, please document it on the README.

CustomRecovery allows the user to control what happens in the case of a panic without having to re-implement the Recovery middleware

codecov[bot]: # [Codecov](https://codecov.io/gh/gin-gonic/gin/pull/2322?src=pr&el=h1) Report
> Merging [#2322](https://codecov.io/gh/gin-gonic/gin/pull/2322?src=pr&el=desc) into [master](https://codecov.io/gh/gin-gonic/gin/commit/4f208887e1231459672a2a9fc1b2aa40486825d4&el=desc) will **decrease** coverage by `0.08%`.
> The diff coverage is `71.42%`.

[![Impacted file tree graph](https://codecov.io/gh/gin-gonic/gin/pull/2322/graphs/tree.svg?width=650&height=150&src=pr&token=Qpqxhwery8)](https://codecov.io/gh/gin-gonic/gin/pull/2322?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #2322      +/-   ##
==========================================
- Coverage   98.39%   98.31%   -0.09%     
==========================================
  Files          41       41              
  Lines        2304     2310       +6     
==========================================
+ Hits         2267     2271       +4     
- Misses         21       23       +2     
  Partials       16       16              
```


| [Impacted Files](https://codecov.io/gh/gin-gonic/gin/pull/2322?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [recovery.go](https://codecov.io/gh/gin-gonic/gin/pull/2322/diff?src=pr&el=tree#diff-cmVjb3ZlcnkuZ28=) | `94.87% <71.42%> (-2.36%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/gin-gonic/gin/pull/2322?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/gin-gonic/gin/pull/2322?src=pr&el=footer). Last update [4f20888...6427797](https://codecov.io/gh/gin-gonic/gin/pull/2322?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

theonlyjohnny: :wave: Could I get a maintainer's :eyes: on this plz? Currently using my fork and would like to not have to

Patch:
diff --git a/README.md b/README.md
index 771b577..7ff5999 100644
--- a/README.md
+++ b/README.md
@@ -496,6 +496,39 @@ func main() {
 }
 ```
 
+### Custom Recovery behavior
+```go
+func main() {
+	// Creates a router without any middleware by default
+	r := gin.New()
+
+	// Global middleware
+	// Logger middleware will write the logs to gin.DefaultWriter even if you set with GIN_MODE=release.
+	// By default gin.DefaultWriter = os.Stdout
+	r.Use(gin.Logger())
+
+	// Recovery middleware recovers from any panics and writes a 500 if there was one.
+	r.Use(gin.CustomRecovery(func(c *gin.Context, recovered interface{}) {
+		if err, ok := recovered.(string); ok {
+			c.String(http.StatusInternalServerError, fmt.Sprintf("error: %s", err))
+		}
+		c.AbortWithStatus(http.StatusInternalServerError)
+	}))
+
+	r.GET("/panic", func(c *gin.Context) {
+		// panic with a string -- the custom middleware could save this to a database or report it to the user
+		panic("foo")
+	})
+
+	r.GET("/", func(c *gin.Context) {
+		c.String(http.StatusOK, "ohai")
+	})
+
+	// Listen and serve on 0.0.0.0:8080
+	r.Run(":8080")
+}
+```
+
 ### How to write log file
 ```go
 func main() {
diff --git a/recovery.go b/recovery.go
index 8cf0932..d02b829 100644
--- a/recovery.go
+++ b/recovery.go
@@ -26,13 +26,29 @@ 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 RecoveryWithWriter(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 {
+func RecoveryWithWriter(out io.Writer, recovery ...RecoveryFunc) HandlerFunc {
+	if len(recovery) > 0 {
+		return CustomRecoveryWithWriter(out, recovery[0])
+	}
+	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 +86,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 +99,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


Test patch:
diff --git a/recovery_test.go b/recovery_test.go
index 21a0a48..6cc2a47 100644
--- a/recovery_test.go
+++ b/recovery_test.go
@@ -62,7 +62,7 @@ func TestPanicInHandler(t *testing.T) {
 	assert.Equal(t, http.StatusInternalServerError, 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(), "TestPanicInHandler")
+	assert.Contains(t, buffer.String(), t.Name())
 	assert.NotContains(t, buffer.String(), "GET /recovery")
 
 	// Debug mode prints the request
@@ -144,3 +144,107 @@ 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(), t.Name())
+	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)
+}
+
+func TestCustomRecovery(t *testing.T) {
+	errBuffer := new(bytes.Buffer)
+	buffer := new(bytes.Buffer)
+	router := New()
+	DefaultErrorWriter = buffer
+	handleRecovery := func(c *Context, err interface{}) {
+		errBuffer.WriteString(err.(string))
+		c.AbortWithStatus(http.StatusBadRequest)
+	}
+	router.Use(CustomRecovery(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(), t.Name())
+	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)
+}
+
+func TestRecoveryWithWriterWithCustomRecovery(t *testing.T) {
+	errBuffer := new(bytes.Buffer)
+	buffer := new(bytes.Buffer)
+	router := New()
+	DefaultErrorWriter = buffer
+	handleRecovery := func(c *Context, err interface{}) {
+		errBuffer.WriteString(err.(string))
+		c.AbortWithStatus(http.StatusBadRequest)
+	}
+	router.Use(RecoveryWithWriter(DefaultErrorWriter, 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(), t.Name())
+	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)
+}